From 87e3749b55e6f3e6fda98e0c53969aacfad1e19b Mon Sep 17 00:00:00 2001 From: yichao-liang Date: Fri, 27 Dec 2024 21:12:59 +0000 Subject: [PATCH] refactor _get_expected_finger_normals to base pybullet env --- predicators/envs/pybullet_balance.py | 15 --------------- predicators/envs/pybullet_blocks.py | 15 --------------- predicators/envs/pybullet_coffee.py | 13 ------------- predicators/envs/pybullet_cover.py | 7 ------- predicators/envs/pybullet_env.py | 22 ++++++++++++++-------- predicators/envs/pybullet_grow.py | 13 ------------- 6 files changed, 14 insertions(+), 71 deletions(-) diff --git a/predicators/envs/pybullet_balance.py b/predicators/envs/pybullet_balance.py index aba049a59..1a39dab7a 100644 --- a/predicators/envs/pybullet_balance.py +++ b/predicators/envs/pybullet_balance.py @@ -527,21 +527,6 @@ def _load_task_from_json(self, json_file: Path) -> EnvironmentTask: def _get_object_ids_for_held_check(self) -> List[int]: return sorted(self._block_id_to_block) - def _get_expected_finger_normals(self) -> Dict[int, Array]: - if CFG.pybullet_robot == "panda": - # gripper rotated 90deg so parallel to x-axis - normal = np.array([1., 0., 0.], dtype=np.float32) - elif CFG.pybullet_robot == "fetch": - # gripper parallel to y-axis - normal = np.array([0., 1., 0.], dtype=np.float32) - else: # pragma: no cover - # Shouldn't happen unless we introduce a new robot. - raise ValueError(f"Unknown robot {CFG.pybullet_robot}") - - return { - self._pybullet_robot.left_finger_id: normal, - self._pybullet_robot.right_finger_id: -1 * normal, - } def _force_grasp_object(self, block: Object) -> None: block_to_block_id = {b: i for i, b in self._block_id_to_block.items()} diff --git a/predicators/envs/pybullet_blocks.py b/predicators/envs/pybullet_blocks.py index dd37cd440..b852273b2 100644 --- a/predicators/envs/pybullet_blocks.py +++ b/predicators/envs/pybullet_blocks.py @@ -230,21 +230,6 @@ def _load_task_from_json(self, json_file: Path) -> EnvironmentTask: def _get_object_ids_for_held_check(self) -> List[int]: return sorted(self._block_id_to_block) - def _get_expected_finger_normals(self) -> Dict[int, Array]: - if CFG.pybullet_robot == "panda": - # gripper rotated 90deg so parallel to x-axis - normal = np.array([1., 0., 0.], dtype=np.float32) - elif CFG.pybullet_robot == "fetch": - # gripper parallel to y-axis - normal = np.array([0., 1., 0.], dtype=np.float32) - else: # pragma: no cover - # Shouldn't happen unless we introduce a new robot. - raise ValueError(f"Unknown robot {CFG.pybullet_robot}") - - return { - self._pybullet_robot.left_finger_id: normal, - self._pybullet_robot.right_finger_id: -1 * normal, - } def _force_grasp_object(self, block: Object) -> None: block_to_block_id = {b: i for i, b in self._block_id_to_block.items()} diff --git a/predicators/envs/pybullet_coffee.py b/predicators/envs/pybullet_coffee.py index bcf52b513..4d2be57d1 100644 --- a/predicators/envs/pybullet_coffee.py +++ b/predicators/envs/pybullet_coffee.py @@ -699,19 +699,6 @@ def _get_object_ids_for_held_check(self) -> List[int]: return [self._jug_id, self._plug_id] return [self._jug_id] - def _get_expected_finger_normals(self) -> Dict[int, Array]: - if CFG.pybullet_robot == "fetch": - # gripper parallel to y-axis - normal = np.array([0., 1., 0.], dtype=np.float32) - else: # pragma: no cover - # Shouldn't happen unless we introduce a new robot. - raise ValueError(f"Unknown robot {CFG.pybullet_robot}") - - return { - self._pybullet_robot.left_finger_id: normal, - self._pybullet_robot.right_finger_id: -1 * normal, - } - def _state_to_gripper_orn(self, state: State) -> Quaternion: wrist = state.get(self._robot, "wrist") tilt = state.get(self._robot, "tilt") diff --git a/predicators/envs/pybullet_cover.py b/predicators/envs/pybullet_cover.py index 90da2e07a..091a79d3c 100644 --- a/predicators/envs/pybullet_cover.py +++ b/predicators/envs/pybullet_cover.py @@ -279,13 +279,6 @@ def _get_state(self) -> State: def _get_object_ids_for_held_check(self) -> List[int]: return sorted(self._block_id_to_block) - def _get_expected_finger_normals(self) -> Dict[int, Array]: - # Both fetch and panda have grippers parallel to x-axis - return { - self._pybullet_robot.left_finger_id: np.array([1., 0., 0.]), - self._pybullet_robot.right_finger_id: np.array([-1., 0., 0.]), - } - @classmethod def get_name(cls) -> str: return "pybullet_cover" diff --git a/predicators/envs/pybullet_env.py b/predicators/envs/pybullet_env.py index acc7422fa..830ae163c 100644 --- a/predicators/envs/pybullet_env.py +++ b/predicators/envs/pybullet_env.py @@ -164,15 +164,21 @@ def _get_object_ids_for_held_check(self) -> List[int]: held.""" raise NotImplementedError("Override me!") - @abc.abstractmethod def _get_expected_finger_normals(self) -> Dict[int, Array]: - """Get the expected finger normals, used in detect_held_object(), as a - mapping from finger link index to a unit-length normal vector. - - This is environment-specific because it depends on the end - effector's orientation when grasping. - """ - raise NotImplementedError("Override me!") + if CFG.pybullet_robot == "panda": + # gripper rotated 90deg so parallel to x-axis + normal = np.array([1., 0., 0.], dtype=np.float32) + elif CFG.pybullet_robot == "fetch": + # gripper parallel to y-axis + normal = np.array([0., 1., 0.], dtype=np.float32) + else: # pragma: no cover + # Shouldn't happen unless we introduce a new robot. + raise ValueError(f"Unknown robot {CFG.pybullet_robot}") + + return { + self._pybullet_robot.left_finger_id: normal, + self._pybullet_robot.right_finger_id: -1 * normal, + } @classmethod def _fingers_state_to_joint(cls, pybullet_robot: SingleArmPyBulletRobot, diff --git a/predicators/envs/pybullet_grow.py b/predicators/envs/pybullet_grow.py index 02ae5de10..539973128 100644 --- a/predicators/envs/pybullet_grow.py +++ b/predicators/envs/pybullet_grow.py @@ -269,19 +269,6 @@ def _get_object_ids_for_held_check(self) -> List[int]: assert self._red_jug_id is not None and self._blue_jug_id is not None return [self._red_jug_id, self._blue_jug_id] - def _get_expected_finger_normals(self) -> Dict[int, Array]: - """For the default fetch robot in predicators. - - We assume a certain orientation where the left_finger_id is in - +y direction, the right_finger_id is in -y. - """ - normal_left = np.array([0., 1., 0.], dtype=np.float32) - normal_right = np.array([0., -1., 0.], dtype=np.float32) - return { - self._pybullet_robot.left_finger_id: normal_left, - self._pybullet_robot.right_finger_id: normal_right, - } - # ------------------------------------------------------------------------- # Setting or updating the environment’s state.