Move remaining Pybullet code to dedicated modules (#1171)

- Using IDE functionality

Author: williamshen-nz

src/envs/pybullet_robots.py src/pybullet_helpers/motion_planning.pysrc/envs/pybullet_robots.py renamed to 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,

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
 

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

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