remove gripper from initial observation, and add wrist cam (#57)

* remove gripper from initial observation, and add wrist cam

* small change for now

* update to the right version

* update to reflect the dataset change

* bump to include the new background removal

* add some small changes

* remove symmetry default

* add in wrist cam and remove gripper

configs/commands/rlbench/_train.yaml

@@ -12,7 +12,7 @@ pretraining:
   mode: none
 
 training:
-  max_epochs: 500
+  max_epochs: 2000
   batch_size: 8
   check_val_every_n_epoch: 10
   log_every_n_steps: 100

configs/dataset/rlbench/_default.yaml

@@ -15,11 +15,13 @@ train_dset:
     cached: true
     phase: ${task.phase.name}
     teleport_initial_to_final: true
-    with_symmetry: True
+    with_symmetry: False
     occlusion_cfg: ???
     num_points: 512
     action_mode: "gripper_and_object"
     anchor_mode: "background_robot_removed"
+    include_wrist_cam: True
+    gripper_in_first_phase: False
 
 val_dset:
   demo_dset:
@@ -32,8 +34,10 @@ val_dset:
     cached: true
     phase: ${task.phase.name}
     teleport_initial_to_final: true
-    with_symmetry: True
+    with_symmetry: False
     occlusion_cfg: ???
     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}
+    gripper_in_first_phase: ${...train_dset.demo_dset.gripper_in_first_phase}

configs/eval_rlbench.yaml

@@ -23,6 +23,8 @@ policy_spec:
   model: ${model}
   include_rgb_features: False
   add_random_jitter: True
+  include_wrist_cam: True
+  gripper_in_first_phase: False
 
 
 # Usually only a single checkpoint, but we could have multiple for each phase.

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@11722e4d803581cc79097f333c1f2c2eb536e15f",
+  "rpad-rlbench-utils @ git+https://github.com/r-pad/rlbench_utils.git@5d6167cbb07c25242e3a87a5d5a896c05a7370a7",
 ]
 
 

scripts/README.md

@@ -167,25 +167,25 @@ RLBENCH_10_TASKS = [
 ## Training
 
 
-./launch.sh autobot 5 python scripts/train_residual_flow.py --config-name commands/rlbench/pick_and_lift/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 0 python scripts/train_residual_flow.py --config-name commands/rlbench/pick_and_lift/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 1 python scripts/train_residual_flow.py --config-name commands/rlbench/pick_up_cup/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 1 python scripts/train_residual_flow.py --config-name commands/rlbench/pick_up_cup/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 2 python scripts/train_residual_flow.py --config-name commands/rlbench/put_knife_on_chopping_board/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 2 python scripts/train_residual_flow.py --config-name commands/rlbench/put_knife_on_chopping_board/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_money_in_safe/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 3 python scripts/train_residual_flow.py --config-name commands/rlbench/put_money_in_safe/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 4 python scripts/train_residual_flow.py --config-name commands/rlbench/push_button/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 4 python scripts/train_residual_flow.py --config-name commands/rlbench/push_button/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 0 python scripts/train_residual_flow.py --config-name commands/rlbench/reach_target/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 5 python scripts/train_residual_flow.py --config-name commands/rlbench/reach_target/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 6 python scripts/train_residual_flow.py --config-name commands/rlbench/slide_block_to_target/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 6 python scripts/train_residual_flow.py --config-name commands/rlbench/slide_block_to_target/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 7 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 7 python scripts/train_residual_flow.py --config-name commands/rlbench/stack_wine/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 8 python scripts/train_residual_flow.py --config-name commands/rlbench/take_money_out_safe/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 8 python scripts/train_residual_flow.py --config-name commands/rlbench/take_money_out_safe/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
-./launch.sh autobot 9 python scripts/train_residual_flow.py --config-name commands/rlbench/take_umbrella_out_of_umbrella_stand/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3
+./launch.sh autobot 9 python scripts/train_residual_flow.py --config-name commands/rlbench/take_umbrella_out_of_umbrella_stand/train_taxpose_tc.yaml dm.train_dset.include_rgb_features=True model.feature_channels=3 benchmark.dataset_root=/data/rlbench10_collisions
 
 ## Checkpoints from this training run
 
@@ -293,3 +293,67 @@ take_umbrella_out_of_umbrella_stand:    r-pad/taxpose/model-txvpna0v:v0
 ./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/take_umbrella_out_of_umbrella_stand/taxpose_tc/eval_rlbench.yaml num_trials=100 policy_spec.include_rgb_features=True model.feature_channels=3 checkpoints.ckpt_file=r-pad/taxpose/model-txvpna0v:v0 wandb.group=rlbench_take_umbrella_out_of_umbrella_stand
 
 ```
+
+# Wrist Cam + Gripper Out
+
+## Training
+
+See the branch.
+
+## Precision Eval
+
+Trained models:
+pick_and_lift: hbdcvydg
+pick_up_cup: gwfjqpfk
+push_button: mmc5fhzu
+put_knife_on_chopping_board: rjscih24
+put_money_in_safe: gj3h3o3c
+reach_target: v8vv53tx
+slide_block_to_target: sejd7pz0
+stack_wine: 3hyo3r7q
+take_money_out_safe: 69gka1ew
+take_umbrella_out_of_umbrella_stand: b48mz8e1
+
+```
+
+# pick_and_lift
+
+./configs/commands/rlbench/pick_and_lift/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-hbdcvydg:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# pick_up_cup
+
+./configs/commands/rlbench/pick_up_cup/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-gwfjqpfk:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# push_button
+
+./configs/commands/rlbench/push_button/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-mmc5fhzu:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# put_knife_on_chopping_board
+
+./configs/commands/rlbench/put_knife_on_chopping_board/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-rjscih24:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# put_money_in_safe
+
+./configs/commands/rlbench/put_money_in_safe/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-gj3h3o3c:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# reach_target
+
+./configs/commands/rlbench/reach_target/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-v8vv53tx:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# slide_block_to_target
+
+./configs/commands/rlbench/slide_block_to_target/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-sejd7pz0:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# stack_wine
+
+./configs/commands/rlbench/stack_wine/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-3hyo3r7q:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# take_money_out_safe
+
+./configs/commands/rlbench/take_money_out_safe/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-69gka1ew:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+# take_umbrella_out_of_umbrella_stand
+
+./configs/commands/rlbench/take_umbrella_out_of_umbrella_stand/taxpose_tc/precision_eval/precision_eval.sh dm.train_dset.include_rgb_features=True model.feature_channels=3 checkpoint=r-pad/taxpose/model-b48mz8e1:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+```

scripts/eval_rlbench.py

@@ -347,6 +347,8 @@ def predict(self, obs, phase: str) -> Tuple[np.ndarray, Dict[str, Any]]:
             self.action_mode,
             self.anchor_mode,
             n_pts=self.policy_spec.num_points,
+            include_wrist_cam=self.policy_spec.include_wrist_cam,
+            gripper_in_first_phase=self.policy_spec.gripper_in_first_phase,
         )
 
         model = self.models[phase]
@@ -551,8 +553,16 @@ def obs_to_input(
     action_mode: ActionMode,
     anchor_mode: AnchorMode,
     n_pts: Optional[int] = 1024,
+    include_wrist_cam=False,
+    gripper_in_first_phase=False,
 ):
-    rgb, pc, mask = obs_to_rgb_point_cloud(obs)
+    rgb, pc, mask = obs_to_rgb_point_cloud(obs, include_wrist_cam)
+
+    # Filter out any points with the mask == 16777215
+    mask_ixs = (mask != 16777215).squeeze()
+    mask = mask[mask_ixs]
+    pc = pc[mask_ixs]
+    rgb = rgb[mask_ixs]
 
     ########################################
     # Separate the action and anchor points.
@@ -576,6 +586,7 @@ def obs_to_input(
         task_name,
         phase,
         use_from_simulator=True,
+        gripper_in_first_phase=gripper_in_first_phase,
     )
 
     ##############################