From 523656bfc904d6a717b7c97d9c2c4c7174936287 Mon Sep 17 00:00:00 2001
From: Ashay Athalye <aathalye@theaiinstitute.com>
Date: Wed, 18 Sep 2024 20:55:00 -0400
Subject: [PATCH] Add options.

---
 predicators/envs/spot_env.py                  | 74 +++++++++----------
 .../ground_truth_models/spot_env/options.py   | 12 ++-
 2 files changed, 45 insertions(+), 41 deletions(-)

diff --git a/predicators/envs/spot_env.py b/predicators/envs/spot_env.py
index e469031e6..cb3132d43 100644
--- a/predicators/envs/spot_env.py
+++ b/predicators/envs/spot_env.py
@@ -2496,41 +2496,41 @@ def _detection_id_to_obj(self) -> Dict[ObjectDetectionID, Object]:
         return detection_id_to_obj
 
     def _create_operators(self) -> Iterator[STRIPSOperator]:
-        # Pick object
-        robot = Variable("?robot", _robot_type)
-        obj = Variable("?object", _movable_object_type)
-        table = Variable("?table", _movable_object_type)
-        parameters = [robot, obj, table]
-        preconds: Set[LiftedAtom] = {
-            LiftedAtom(_HandEmpty, [robot]),
-            LiftedAtom(_NotHolding, [robot, obj]),
-            LiftedAtom(_VLMOn, [obj, table])
-        }
-        add_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
-        del_effs: Set[LiftedAtom] = {
-            LiftedAtom(_HandEmpty, [robot]),
-            LiftedAtom(_NotHolding, [robot, obj]),
-            LiftedAtom(_VLMOn, [obj, table])
-        }
-        ignore_effs: Set[LiftedAtom] = set()
-        yield STRIPSOperator("Pick", parameters, preconds, add_effs, del_effs,
-                             ignore_effs)
-
-        # Place object
-        robot = Variable("?robot", _robot_type)
-        obj = Variable("?object", _movable_object_type)
-        pan = Variable("?pan", _container_type)
-        parameters = [robot, obj, pan]
-        preconds: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
-        add_effs: Set[LiftedAtom] = {
-            LiftedAtom(_HandEmpty, [robot]),
-            LiftedAtom(_NotHolding, [robot, obj]),
-            LiftedAtom(_VLMOn, [obj, pan])
-        }
-        del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
-        ignore_effs: Set[LiftedAtom] = set()
-        yield STRIPSOperator("Place", parameters, preconds, add_effs, del_effs,
-                             ignore_effs)
+        # # Pick object
+        # robot = Variable("?robot", _robot_type)
+        # obj = Variable("?object", _movable_object_type)
+        # table = Variable("?table", _movable_object_type)
+        # parameters = [robot, obj, table]
+        # preconds: Set[LiftedAtom] = {
+        #     LiftedAtom(_HandEmpty, [robot]),
+        #     LiftedAtom(_NotHolding, [robot, obj]),
+        #     LiftedAtom(_VLMOn, [obj, table])
+        # }
+        # add_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
+        # del_effs: Set[LiftedAtom] = {
+        #     LiftedAtom(_HandEmpty, [robot]),
+        #     LiftedAtom(_NotHolding, [robot, obj]),
+        #     LiftedAtom(_VLMOn, [obj, table])
+        # }
+        # ignore_effs: Set[LiftedAtom] = set()
+        # yield STRIPSOperator("Pick", parameters, preconds, add_effs, del_effs,
+        #                      ignore_effs)
+
+        # # Place object
+        # robot = Variable("?robot", _robot_type)
+        # obj = Variable("?object", _movable_object_type)
+        # pan = Variable("?pan", _container_type)
+        # parameters = [robot, obj, pan]
+        # preconds: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
+        # add_effs: Set[LiftedAtom] = {
+        #     LiftedAtom(_HandEmpty, [robot]),
+        #     LiftedAtom(_NotHolding, [robot, obj]),
+        #     LiftedAtom(_VLMOn, [obj, pan])
+        # }
+        # del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, obj])}
+        # ignore_effs: Set[LiftedAtom] = set()
+        # yield STRIPSOperator("Place", parameters, preconds, add_effs, del_effs,
+        #                      ignore_effs)
         
 
         ##########################################3
@@ -2548,7 +2548,7 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
             LiftedAtom(_NotHolding, [robot, dustpan]),
         }
         ignore_effs: Set[LiftedAtom] = set()
-        yield STRIPSOperator("Pick", parameters, preconds, add_effs, del_effs,
+        yield STRIPSOperator("Pick1", parameters, preconds, add_effs, del_effs,
                              ignore_effs)
 
         # Place(robot, dustpan, mess)
@@ -2581,7 +2581,7 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
             LiftedAtom(_NotHolding, [robot, broom]),
         }
         ignore_effs: Set[LiftedAtom] = set()
-        yield STRIPSOperator("Pick", parameters, preconds, add_effs, del_effs,
+        yield STRIPSOperator("Pick2", parameters, preconds, add_effs, del_effs,
                              ignore_effs)
         
         # Sweep(robot, broom, mess, dustpan)
diff --git a/predicators/ground_truth_models/spot_env/options.py b/predicators/ground_truth_models/spot_env/options.py
index 6e26a7ccf..daf0c096b 100644
--- a/predicators/ground_truth_models/spot_env/options.py
+++ b/predicators/ground_truth_models/spot_env/options.py
@@ -957,8 +957,10 @@ def _teleop(robot: Robot, lease_client: LeaseClient):
     "PrepareContainerForSweeping": Box(-np.inf, np.inf, (3, )),  # dx, dy, dyaw
     "DropNotPlaceableObject": Box(0, 1, (0, )),  # empty
     "MoveToReadySweep": Box(0, 1, (0, )),  # empty
-    "Pick": Box(0, 1, (0, )),  # empty
-    "Place": Box(0, 1, (0, ))  # empty
+    "Pick1": Box(0, 1, (0, )),  # empty
+    "Place": Box(0, 1, (0, )),  # empty
+    "Pick2": Box(0, 1, (0, )),  # empty
+    "Sweep": Box(0, 1, (0, ))  # empty
 }
 
 # NOTE: the policies MUST be unique because they output actions with extra info
@@ -982,8 +984,10 @@ def _teleop(robot: Robot, lease_client: LeaseClient):
     "PrepareContainerForSweeping": _prepare_container_for_sweeping_policy,
     "DropNotPlaceableObject": _drop_not_placeable_object_policy,
     "MoveToReadySweep": _move_to_ready_sweep_policy,
-    "Pick": _teleop_policy,
-    "Place": _teleop_policy
+    "Pick1": _teleop_policy,
+    "Place": _teleop_policy,
+    "Pick2": _teleop_policy,
+    "Sweep": _teleop_policy
 }