diff --git a/src/envs/pybullet_robots.py b/src/pybullet_helpers/motion_planning.py similarity index 71% rename from src/envs/pybullet_robots.py rename to src/pybullet_helpers/motion_planning.py index 5f70af1e0a..7d2aa33660 100644 --- a/src/envs/pybullet_robots.py +++ b/src/pybullet_helpers/motion_planning.py @@ -1,38 +1,16 @@ -"""Interfaces to PyBullet robots.""" +"""Motion Planning in Pybullet.""" from __future__ import annotations -from typing import TYPE_CHECKING, Collection, Iterator, List, Optional, \ - Sequence +from typing import Collection, Iterator, Optional, Sequence import numpy as np import pybullet as p from predicators.src import utils +from predicators.src.pybullet_helpers.robots import SingleArmPyBulletRobot from predicators.src.settings import CFG from predicators.src.structs import JointsState -if TYPE_CHECKING: - from predicators.src.pybullet_helpers.robots import SingleArmPyBulletRobot - -########################### Other utility functions ########################### - - -def get_kinematic_chain(robot: int, end_effector: int, - physics_client_id: int) -> List[int]: - """Get all of the free joints from robot base to end effector. - - Includes the end effector. - """ - kinematic_chain = [] - while end_effector > -1: - joint_info = p.getJointInfo(robot, - end_effector, - physicsClientId=physics_client_id) - if joint_info[3] > -1: - kinematic_chain.append(end_effector) - end_effector = joint_info[-1] - return kinematic_chain - def run_motion_planning( robot: SingleArmPyBulletRobot, initial_state: JointsState, diff --git a/src/pybullet_helpers/robots/single_arm.py b/src/pybullet_helpers/robots/single_arm.py index d76caf947c..e756d00861 100644 --- a/src/pybullet_helpers/robots/single_arm.py +++ b/src/pybullet_helpers/robots/single_arm.py @@ -8,9 +8,9 @@ from gym.spaces import Box from predicators.src import utils -from predicators.src.envs.pybullet_robots import get_kinematic_chain from predicators.src.pybullet_helpers.inverse_kinematics import \ pybullet_inverse_kinematics +from predicators.src.pybullet_helpers.utils import get_kinematic_chain from predicators.src.settings import CFG from predicators.src.structs import Array, JointsState, Pose3D diff --git a/src/pybullet_helpers/utils.py b/src/pybullet_helpers/utils.py new file mode 100644 index 0000000000..f81f784508 --- /dev/null +++ b/src/pybullet_helpers/utils.py @@ -0,0 +1,23 @@ +"""Other Pybullet utility functions.""" +from __future__ import annotations + +from typing import List + +import pybullet as p + + +def get_kinematic_chain(robot: int, end_effector: int, + physics_client_id: int) -> List[int]: + """Get all of the free joints from robot base to end effector. + + Includes the end effector. + """ + kinematic_chain = [] + while end_effector > -1: + joint_info = p.getJointInfo(robot, + end_effector, + physicsClientId=physics_client_id) + if joint_info[3] > -1: + kinematic_chain.append(end_effector) + end_effector = joint_info[-1] + return kinematic_chain diff --git a/tests/envs/test_pybullet_robots.py b/tests/envs/test_pybullet_robots.py index 69e8fdd46e..1e9b1e17f6 100644 --- a/tests/envs/test_pybullet_robots.py +++ b/tests/envs/test_pybullet_robots.py @@ -6,12 +6,13 @@ from predicators.src import utils from predicators.src.envs.pybullet_env import create_pybullet_block -from predicators.src.envs.pybullet_robots import get_kinematic_chain, \ - run_motion_planning from predicators.src.pybullet_helpers.inverse_kinematics import \ pybullet_inverse_kinematics +from predicators.src.pybullet_helpers.motion_planning import \ + run_motion_planning from predicators.src.pybullet_helpers.robots import FetchPyBulletRobot, \ create_single_arm_pybullet_robot +from predicators.src.pybullet_helpers.utils import get_kinematic_chain from predicators.src.settings import CFG