Tighten _robot_matches_state atol so set_state hint forces reset

The fast-path joint-match check used atol=1e-2, which let a caller's
initial_joint_positions hint be silently treated as "already there"
when live joints were within 1e-2 of initial — leaving the EE pose up
to ~3e-3 off the requested state. State.allclose compares features at
1e-3, so the test then failed reconstruction. Match the State.allclose
tolerance.

Also pick up trailing yapf reformatting in two approach files.

Author: yichao-liang

predicators/approaches/agent_planner_approach.py

@@ -65,9 +65,8 @@ def __init__(self,
         # __init__.
         if CFG.wait_option_terminate_on_atom_change:
             cast(  # pylint: disable=protected-access
-                Any, self._option_model
-            )._abstract_function = (
-                lambda s: utils.abstract(s, self._get_all_predicates()))
+                Any, self._option_model)._abstract_function = (
+                    lambda s: utils.abstract(s, self._get_all_predicates()))
         self._online_learning_cycle = 0
         self._requests_train_task_idxs: Optional[List[int]] = None
         self._run_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

predicators/approaches/bilevel_planning_approach.py

@@ -58,9 +58,9 @@ def __init__(self,
         # and so predicates invented later are reflected at call time.
         if CFG.wait_option_terminate_on_atom_change:
             cast(  # pylint: disable=protected-access
-                Any, self._option_model
-            )._abstract_function = (
-                lambda s: utils.abstract(s, self._get_current_predicates()))
+                Any, self._option_model)._abstract_function = (
+                    lambda s: utils.abstract(s, self._get_current_predicates())
+                )
         self._num_calls = 0
         self._last_plan: List[_Option] = []  # used if plan WITH sim
         self._last_nsrt_plan: List[_GroundNSRT] = []  # plan WITHOUT sim

predicators/envs/pybullet_env.py

@@ -453,10 +453,10 @@ def _set_state(self, state: State) -> None:
         #     constraint, so a kept constraint would leave the held
         #     body behind).
         new_held_id = self._held_obj_id_in_state(state)
-        held_obj_moved = (self._held_obj_id is not None and any(
-            o.id == self._held_obj_id for o in objects_to_reset))
-        rebuild_constraint = (full_reset
-                              or new_held_id != self._held_obj_id
+        held_obj_moved = (self._held_obj_id is not None
+                          and any(o.id == self._held_obj_id
+                                  for o in objects_to_reset))
+        rebuild_constraint = (full_reset or new_held_id != self._held_obj_id
                               or (self._held_obj_id is not None and
                                   (robot_changed or held_obj_moved)))
 
@@ -506,9 +506,7 @@ def _set_state(self, state: State) -> None:
                 logging.warning(
                     "Could not reconstruct state exactly in reset.")
 
-    def _robot_matches_state(self,
-                             state: State,
-                             atol: float = 1e-2) -> bool:
+    def _robot_matches_state(self, state: State, atol: float = 1e-3) -> bool:
         """True if PyBullet's live robot pose already equals state's.
 
         Compares at the joint level. The EE-quaternion path that
@@ -517,6 +515,13 @@ def _robot_matches_state(self,
         fail an EE-pose comparison and trigger a full robot reset on
         every simulate() call (visible jitter).
 
+        ``atol`` matches ``State.allclose``'s feature tolerance: a looser
+        check would let the fast-path skip a reset even when the live EE
+        pose differs from the requested state by more than allclose
+        accepts (e.g. when a caller hands us
+        ``initial_joint_positions`` as a hint and the live joints are
+        only 1e-2 close).
+
         Returns False when ``state`` has no joint_positions — the only
         live caller in that situation is
         ``_add_pybullet_state_to_tasks``, where forcing a reset is
@@ -572,8 +577,8 @@ def _object_pose_matches_state(self,
     def _held_obj_id_in_state(self, state: State) -> Optional[int]:
         """Which PyBullet body id is marked is_held > 0.5 in ``state``.
 
-        Returns None if no object is held in ``state``. Mirrors the
-        per-object logic in _reset_single_object before constraint
+        Returns None if no object is held in ``state``. Mirrors the per-
+        object logic in _reset_single_object before constraint
         management was hoisted out into _set_state.
         """
         for obj in state.data: