Fix near_poses behavior

StanfordVL · Sep 29, 2023 · 4504012 · 4504012
1 parent b524661
commit 4504012
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/omnigibson/action_primitives/starter_semantic_action_primitives.py b/omnigibson/action_primitives/starter_semantic_action_primitives.py
@@ -1714,15 +1714,17 @@ def _sample_pose_in_room(self, room: str):
             {"room": room}
         )
 
-    def _sample_pose_with_object_and_predicate(self, predicate, held_obj, target_obj):
+    def _sample_pose_with_object_and_predicate(self, predicate, held_obj, target_obj, near_poses=None, near_poses_threshold=None):
         """
         Returns a pose for the held object relative to the target object that satisfies the predicate
 
         Args:
             predicate (object_states.OnTop or object_states.Inside): Relation between held object and the target object
             held_obj (StatefulObject): Object held by the robot
             target_obj (StatefulObject): Object to sample a pose relative to
-
+            near_poses (Iterable of arrays): Poses in the world frame to sample near
+            near_poses_threshold (float): The distance threshold to check if the sampled pose is near the poses in near_poses
+            
         Returns:
             2-tuple:
                 - 3-array: (x,y,z) Position in the world frame
@@ -1739,6 +1741,12 @@ def _sample_pose_with_object_and_predicate(self, predicate, held_obj, target_obj
             # Get the object pose by subtracting the offset
             sampled_obj_pose = T.pose2mat((sampled_bb_center, sampled_bb_orn)) @ T.pose_inv(T.pose2mat((bb_center_in_base, [0, 0, 0, 1])))
 
+            # Check that the pose is near one of the poses in the near_poses list if provided.
+            if near_poses:
+                sampled_pos = np.array([sampled_obj_pose[0]])
+                if not np.any(np.linalg.norm(near_poses - sampled_pos, axis=1) < near_poses_threshold):
+                    continue
+
             # Return the pose
             return T.mat2pose(sampled_obj_pose)
 

diff --git a/omnigibson/action_primitives/symbolic_semantic_action_primitives.py b/omnigibson/action_primitives/symbolic_semantic_action_primitives.py
@@ -231,7 +231,7 @@ def _place_with_predicate(self, obj, predicate, near_poses=None, near_poses_thre
             )
 
         # Find a spot to put it
-        obj_pose = self._sample_pose_with_object_and_predicate(predicate, obj_in_hand, obj, near_poses=[], near_poses_threshold=None)
+        obj_pose = self._sample_pose_with_object_and_predicate(predicate, obj_in_hand, obj, near_poses=near_poses, near_poses_threshold=near_poses_threshold)
 
         # Get close, release the object.
         # yield from self._navigate_if_needed(obj, pose_on_obj=obj_pose)