From a60692ac0a46057862871c852511dd1ea26f4c71 Mon Sep 17 00:00:00 2001
From: Ashay Athalye <aathalye@theaiinstitute.com>
Date: Wed, 18 Sep 2024 21:02:12 -0400
Subject: [PATCH] Add place broom skill.

---
 predicators/envs/spot_env.py                     | 16 +++++++++++++++-
 .../ground_truth_models/spot_env/options.py      | 10 ++++++----
 predicators/perception/spot_perceiver.py         | 12 ++++++++++--
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/predicators/envs/spot_env.py b/predicators/envs/spot_env.py
index cb3132d43..c16068e0b 100644
--- a/predicators/envs/spot_env.py
+++ b/predicators/envs/spot_env.py
@@ -2564,7 +2564,7 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
         }
         del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, dustpan])}
         ignore_effs: Set[LiftedAtom] = set()
-        yield STRIPSOperator("Place", parameters, preconds, add_effs, del_effs,
+        yield STRIPSOperator("Place1", parameters, preconds, add_effs, del_effs,
                              ignore_effs)
         
         # Pick(robot, broom)
@@ -2599,6 +2599,20 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
         ignore_effs: Set[LiftedAtom] = set()
         yield STRIPSOperator("Sweep", parameters, preconds, add_effs, del_effs,
                              ignore_effs)
+        
+        # Place(robot, broom)
+        robot = Variable("?robot", _robot_type)
+        broom = Variable("?broom", _movable_object_type)
+        parameters = [robot, dustpan]
+        preconds: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, broom])}
+        add_effs: Set[LiftedAtom] = {
+            LiftedAtom(_HandEmpty, [robot]),
+            LiftedAtom(_NotHolding, [robot, broom]),
+        }
+        del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, broom])}
+        ignore_effs: Set[LiftedAtom] = set()
+        yield STRIPSOperator("Place2", parameters, preconds, add_effs, del_effs,
+                             ignore_effs)
 
     # def _generate_train_tasks(self) -> List[EnvironmentTask]:
     #     goal = self._generate_goal_description()  # currently just one goal
diff --git a/predicators/ground_truth_models/spot_env/options.py b/predicators/ground_truth_models/spot_env/options.py
index daf0c096b..cbce83f29 100644
--- a/predicators/ground_truth_models/spot_env/options.py
+++ b/predicators/ground_truth_models/spot_env/options.py
@@ -958,9 +958,10 @@ def _teleop(robot: Robot, lease_client: LeaseClient):
     "DropNotPlaceableObject": Box(0, 1, (0, )),  # empty
     "MoveToReadySweep": Box(0, 1, (0, )),  # empty
     "Pick1": Box(0, 1, (0, )),  # empty
-    "Place": Box(0, 1, (0, )),  # empty
+    "Place1": Box(0, 1, (0, )),  # empty
     "Pick2": Box(0, 1, (0, )),  # empty
-    "Sweep": Box(0, 1, (0, ))  # empty
+    "Sweep": Box(0, 1, (0, )),  # empty
+    "Place2": Box(0, 1, (0, ))  # empty
 }
 
 # NOTE: the policies MUST be unique because they output actions with extra info
@@ -985,9 +986,10 @@ def _teleop(robot: Robot, lease_client: LeaseClient):
     "DropNotPlaceableObject": _drop_not_placeable_object_policy,
     "MoveToReadySweep": _move_to_ready_sweep_policy,
     "Pick1": _teleop_policy,
-    "Place": _teleop_policy,
+    "Place1": _teleop_policy,
     "Pick2": _teleop_policy,
-    "Sweep": _teleop_policy
+    "Sweep": _teleop_policy,
+    "Place2": _teleop_policy
 }
 
 
diff --git a/predicators/perception/spot_perceiver.py b/predicators/perception/spot_perceiver.py
index d60b007a5..c6178f26f 100644
--- a/predicators/perception/spot_perceiver.py
+++ b/predicators/perception/spot_perceiver.py
@@ -643,8 +643,10 @@ def _create_goal(self, state: State,
         assert self._curr_env is not None
         pred_name_to_pred = {p.name: p for p in self._curr_env.predicates}
         VLMOn = pred_name_to_pred["VLMOn"]
+        Inside = pred_name_to_pred["Inside"]
+        Holding = pred_name_to_pred["Holding"]
         HandEmpty = pred_name_to_pred["HandEmpty"]
-        
+
 
         if goal_description == "put the cup in the pan":
             robot = Object("robot", _robot_type)
@@ -657,7 +659,13 @@ def _create_goal(self, state: State,
             return goal
 
         if goal_description == "put the mess in the dustpan":
-
+            robot = Object("robot", _robot_type)
+            dustpan = Object("dustpan", _movable_object_type)
+            wrappers = Object("wrappers", _movable_object_type)
+            goal = {
+                GroundAtom(Inside, [wrappers, dustpan]),
+                GroundAtom(Holding, [robot, dustpan])
+            }
 
         raise NotImplementedError("Unrecognized goal description")