Refactor Pybullet Robots (#1169)

* Move Single Arm Robots to their own module

* Fix the linting issue

Author: williamshen-nz

src/envs/pybullet_blocks.py

@@ -11,8 +11,9 @@
 from predicators.src.envs.blocks import BlocksEnv
 from predicators.src.envs.pybullet_env import PyBulletEnv, \
     create_pybullet_block
-from predicators.src.envs.pybullet_robots import _SingleArmPyBulletRobot, \
-    create_change_fingers_option, create_move_end_effector_to_pose_option, \
+from predicators.src.envs.pybullet_robots import \
+    create_change_fingers_option, create_move_end_effector_to_pose_option
+from predicators.src.pybullet_helpers.robots import SingleArmPyBulletRobot, \
     create_single_arm_pybullet_robot
 from predicators.src.settings import CFG
 from predicators.src.structs import Array, Object, ParameterizedOption, \
@@ -214,7 +215,7 @@ def _initialize_pybullet(self) -> None:
                                       self._physics_client_id))
 
     def _create_pybullet_robot(
-            self, physics_client_id: int) -> _SingleArmPyBulletRobot:
+            self, physics_client_id: int) -> SingleArmPyBulletRobot:
         ee_home = (self.robot_init_x, self.robot_init_y, self.robot_init_z)
         return create_single_arm_pybullet_robot(CFG.pybullet_robot, ee_home,
                                                 self._ee_orn,

src/envs/pybullet_cover.py

@@ -10,8 +10,9 @@
 from predicators.src.envs.cover import CoverEnv
 from predicators.src.envs.pybullet_env import PyBulletEnv, \
     create_pybullet_block
-from predicators.src.envs.pybullet_robots import _SingleArmPyBulletRobot, \
-    create_change_fingers_option, create_move_end_effector_to_pose_option, \
+from predicators.src.envs.pybullet_robots import \
+    create_change_fingers_option, create_move_end_effector_to_pose_option
+from predicators.src.pybullet_helpers.robots import SingleArmPyBulletRobot, \
     create_single_arm_pybullet_robot
 from predicators.src.settings import CFG
 from predicators.src.structs import Action, Array, Object, \
@@ -134,7 +135,7 @@ def _initialize_pybullet(self) -> None:
                                       self._physics_client_id))
 
     def _create_pybullet_robot(
-            self, physics_client_id: int) -> _SingleArmPyBulletRobot:
+            self, physics_client_id: int) -> SingleArmPyBulletRobot:
         ee_home = (self._workspace_x, self._robot_init_y, self._workspace_z)
         return create_single_arm_pybullet_robot(CFG.pybullet_robot, ee_home,
                                                 self._ee_orn,

src/envs/pybullet_env.py

@@ -13,7 +13,7 @@
 
 from predicators.src import utils
 from predicators.src.envs import BaseEnv
-from predicators.src.envs.pybullet_robots import _SingleArmPyBulletRobot
+from predicators.src.pybullet_helpers.robots import SingleArmPyBulletRobot
 from predicators.src.settings import CFG
 from predicators.src.structs import Action, Array, Pose3D, State, Task, Video
 
@@ -116,7 +116,7 @@ def _initialize_pybullet(self) -> None:
 
     @abc.abstractmethod
     def _create_pybullet_robot(
-            self, physics_client_id: int) -> _SingleArmPyBulletRobot:
+            self, physics_client_id: int) -> SingleArmPyBulletRobot:
         """Make and return a PyBullet robot object in the given
         physics_client_id.