little efficiency tweaks

scripts/README.md

@@ -149,7 +149,7 @@ mkdir -p configs/checkpoints/rlbench/take_umbrella_out_of_umbrella_stand/pretrai
 
 
 
-# RGB
+# RGB training + evals.
 
 RLBENCH_10_TASKS = [
     "pick_and_lift",
@@ -164,6 +164,8 @@ RLBENCH_10_TASKS = [
     "take_umbrella_out_of_umbrella_stand",
 ]
 
+## 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
 
@@ -184,3 +186,110 @@ RLBENCH_10_TASKS = [
 ./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 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
+
+## Checkpoints from this training run
+
+pick_and_lift:                          r-pad/taxpose/model-9tx1uje9:v0
+pick_up_cup:                            r-pad/taxpose/model-9pfjeq0j:v0
+put_knife_on_chopping_board:            r-pad/taxpose/model-c9u0u4np:v0
+put_money_in_safe:                      r-pad/taxpose/model-lq4b953m:v0
+push_button:                            r-pad/taxpose/model-0kxpww1x:v0
+reach_target:                           r-pad/taxpose/model-w5kjqoph:v0
+slide_block_to_target:                  r-pad/taxpose/model-fct8vbrq:v0
+stack_wine:                             r-pad/taxpose/model-cbe1hgx4:v0
+take_money_out_safe:                    r-pad/taxpose/model-j3swo5k7:v0
+take_umbrella_out_of_umbrella_stand:    r-pad/taxpose/model-txvpna0v:v0
+
+
+## Precision Eval.
+
+
+```
+# 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-9tx1uje9: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-9pfjeq0j: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-c9u0u4np: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-lq4b953m: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-0kxpww1x: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-w5kjqoph: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-fct8vbrq: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-cbe1hgx4: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-j3swo5k7: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-txvpna0v:v0 benchmark.dataset_root=/data/rlbench10_collisions
+
+```
+
+
+
+## RLBench Eval.
+
+```
+# pick_and_lift
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/pick_and_lift/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-9tx1uje9:v0 wandb.group=rlbench_pick_and_lift
+
+# pick_up_cup
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/pick_up_cup/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-9pfjeq0j:v0 wandb.group=rlbench_pick_up_cup
+
+# put_knife_on_chopping_board
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/put_knife_on_chopping_board/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-c9u0u4np:v0 wandb.group=rlbench_put_knife_on_chopping_board
+
+# put_money_in_safe
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/put_money_in_safe/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-lq4b953m:v0 wandb.group=rlbench_put_money_in_safe
+
+# push_button
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/push_button/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-0kxpww1x:v0 wandb.group=rlbench_push_button
+
+# reach_target
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/reach_target/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-w5kjqoph:v0 wandb.group=rlbench_reach_target
+
+# slide_block_to_target
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/slide_block_to_target/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-fct8vbrq:v0 wandb.group=rlbench_slide_block_to_target
+
+# stack_wine
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/stack_wine/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-cbe1hgx4:v0 wandb.group=rlbench_stack_wine
+
+# take_money_out_safe
+
+./launch.sh local-docker 0 python scripts/eval_rlbench.py --config-name commands/rlbench/take_money_out_safe/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-j3swo5k7:v0 wandb.group=rlbench_take_money_out_safe
+
+# take_umbrella_out_of_umbrella_stand
+
+./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
+
+```

scripts/eval_metrics.py

@@ -99,8 +99,6 @@ def main(cfg):
         cfg=cfg.dm,
     )
 
-    dm.setup(stage=cfg.split)
-
     network = create_network(cfg.model)
 
     model = EquivarianceTrainingModule(
@@ -132,13 +130,6 @@ def main(cfg):
 
     # model = torch.compile(model)
 
-    if cfg.split == "train":
-        loader = dm.train_dataloader()
-    elif cfg.split == "val":
-        loader = dm.val_dataloader()
-    elif cfg.split == "test":
-        loader = dm.test_dataloader()
-
     print("setting up train")
     dm.setup(stage="train")
     train_dataloader = dm.train_dataloader()

scripts/eval_rlbench.py

@@ -101,7 +101,7 @@
         "place": False,
     },
     "take_umbrella_out_of_umbrella_stand": {
-        "pregrasp": False,
+        "pregasp": False,
         "grasp": False,
         "lift": True,
     },
@@ -141,8 +141,10 @@ def write_video(self, video_path: str):
         out.release()
 
         # ffmpeg -i "$file" -c:v libx264 -c:a aac "${file%.mp4}.mov"
+
+        # Redirect the output to /dev/null to suppress the output.
         os.system(
-            f"/usr/bin/ffmpeg -i {video_path}.mp4 -c:v libx264 -c:a aac {video_path}.mov"
+            f"/usr/bin/ffmpeg -i {video_path}.mp4 -c:v libx264 -c:a aac {video_path}.mov > /dev/null 2>&1"
         )
 
 
@@ -739,8 +741,8 @@ def pr():
                 T_gripper_world, extras = policy.predict(obs, phase)
                 if "plot" in extras:
                     phase_plots.append((phase, extras["plot"]))
-            except Exception as e:
-                print(e)
+            except Exception as ex:
+
                 phase_results[phase] = FailureReason.PREDICTION_FAILURE
                 success = False
                 return TrialResult(success, pr())
@@ -764,11 +766,14 @@ def pr():
                     ignore_collisions=TASK_TO_IGNORE_COLLISIONS[task_spec.name][phase],
                 )  # Eventually add collision checking.
 
-            except Exception as e:
-                if "workspace" in str(e):
+            except Exception as ex:
+                if "workspace" in str(ex):
                     phase_results[phase] = FailureReason.PREDICTION_OUTSIDE_WORKSPACE
-                else:
+                elif "A path could not be found." in str(ex):
                     phase_results[phase] = FailureReason.MOTION_PLANNING_FAILURE
+                else:
+                    logging.error(f"Unknown error: {ex}")
+                    phase_results[phase] = FailureReason.UNKNOWN_FAILURE
                 success = False
                 return TrialResult(success, pr())
 
@@ -914,6 +919,9 @@ def main(cfg):
         config=OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True),
     )
 
+    # Make sure the model weights can be downloaded.
+    get_weights_path(cfg.checkpoints.ckpt_file, cfg.wandb, run=run)
+
     results = run_trials(
         cfg.policy_spec,
         cfg.task,