Add place broom skill.

bdaiinstitute · Sep 19, 2024 · a60692a · a60692a
1 parent 523656b
commit a60692a
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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")