Add PyBullet robot

.gitignore

@@ -0,0 +1 @@
+__pycache__

arm_1dof/bullet_arm_1dof.py

@@ -0,0 +1,49 @@
+from tokenize import String
+import pybullet 
+import robot_arm_1dof
+
+import os
+
+class BulletArm1Dof:
+    """Class to initialize pybullet server and load robot"""
+    SDF_MODELS_SUBDIR="arm_1dof"
+    SDF_SEARCH_PATH="/".join([os.environ.get("SDF_MODELS_DIR", "."), SDF_MODELS_SUBDIR])
+    SDF_MODEL_NAME="demo_human_col"
+    SDF_MODEL_FILENAME="/".join([SDF_MODEL_NAME, "skeleton.sdf"])
+    URDF_MODEL_FILENAME="/".join([SDF_MODEL_NAME, "skeleton.urdf"])
+
+    def __init__(self) -> None:
+        self._server_id = -1
+        self._skel_arm = None
+        self._skel_arm_id = -1
+
+    @property
+    def server_id(self):
+        """Pybullet server ID"""
+        return self._server_id
+
+    def InitPybullet(self, g=[0., 0., 0.], bullet_connect=pybullet.DIRECT, sdf_search_path=SDF_SEARCH_PATH) -> None:
+        """Start bullet server and set sdf search path"""
+        # Connect to pybullet
+        self._server_id = pybullet.connect(bullet_connect)
+
+        pybullet.setGravity(g[0], g[1], g[2], physicsClientId=self._server_id)
+        pybullet.setAdditionalSearchPath(path=sdf_search_path, physicsClientId=self._server_id)
+
+        self._timestep = pybullet.getPhysicsEngineParameters(physicsClientId=self._server_id)['fixedTimeStep']
+
+
+    def LoadRobot(self, robot_sdf_filename=URDF_MODEL_FILENAME) -> robot_arm_1dof.RobotArm1Dof:
+        """Load robot from sdf file"""
+        self._skel_arm_id = pybullet.loadURDF(fileName=robot_sdf_filename, useFixedBase=True, physicsClientId=self._server_id)
+
+        self._skel_arm = robot_arm_1dof.RobotArm1Dof(bullet_ctrl=self, bullet_body_id=self._skel_arm_id)
+
+        return self._skel_arm
+
+    def Simulate(self, sim_time) -> None:
+        t = 0
+        while t < sim_time:
+            pybullet.stepSimulation(physicsClientId=self._server_id)
+            t += self._timestep
+            

arm_1dof/robot_arm_1dof.py

@@ -0,0 +1,70 @@
+from ntpath import join
+import numpy as np
+from pandas import array
+import pybullet
+
+class RobotArm1Dof:
+    """Sekeleton arm bullet body"""
+    UPPER_ARM_LINK_ID=1
+    FOREARM_LINK_ID=2
+    HAND_LINK_ID=3
+
+    SHOULDER_A_JOINT_ID=0
+    ELBOW_JOINT_ID=1
+
+    def __init__(self, bullet_ctrl, bullet_body_id) -> None:
+
+        self._bullet_ctrl = bullet_ctrl
+        self._body_id = bullet_body_id
+        pybullet.resetBasePositionAndOrientation(
+            bodyUniqueId=self._body_id, 
+            posObj=[0., 0., 0.], ornObj=[0., 0., 0., 1.], 
+            physicsClientId=self._bullet_ctrl.server_id)
+        
+        # Set mass
+        dyn_info = pybullet.getDynamicsInfo(self._body_id, RobotArm1Dof.UPPER_ARM_LINK_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+        self._upper_arm_mass = dyn_info[0]
+
+        dyn_info = pybullet.getDynamicsInfo(self._body_id, RobotArm1Dof.FOREARM_LINK_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+        self._lower_arm_mass = dyn_info[1]
+
+        self._joint_stats = {}
+        self._ee_link_state = None
+
+        self.UpdateStats()
+
+    def UpdateStats(self):
+        """Query state info from bullet"""
+        # Retrieve joint info
+        self._joint_stats[RobotArm1Dof.SHOULDER_A_JOINT_ID] = pybullet.getJointState(self._body_id, RobotArm1Dof.SHOULDER_A_JOINT_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        self._joint_stats[RobotArm1Dof.ELBOW_JOINT_ID] = pybullet.getJointState(self._body_id, RobotArm1Dof.ELBOW_JOINT_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        self._ee_link_state = pybullet.getLinkState(self._body_id, RobotArm1Dof.HAND_LINK_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+    def JointPos(self, joint_id) -> float:
+        """Get joint info"""
+        return self._joint_stats[joint_id][0]
+
+    def JointVel(self, joint_id) -> float:
+        """Get joint info"""
+        return self._joint_stats[joint_id][1]
+
+    def EEPose(self) -> tuple:
+        """Get end effector pose, linear pos and rotation (as quaternion)"""
+        return (self._ee_link_state[0], self._ee_link_state[1])
+
+    def EEVel(self) -> tuple:
+        """Get end effector vel, linear and angular"""
+        return (self._ee_link_state[6], self._ee_link_state[7])
+
+    def SetJointTorque(self, joint_ids, torques) -> None:
+        """Sets torque applied to specified joint"""
+        pybullet.setJointMotorControl2(self._body_id, joint_ids, \
+            controlMode=pybullet.TORQUE_CONTROL, forces=torques, \
+            physicsClientId=self._bullet_ctrl.server_id)

arm_1dof/test_arm_1dof.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+
+from bullet_arm_1dof import BulletArm1Dof
+from robot_arm_1dof import RobotArm1Dof
+
+import pybullet
+
+if __name__ == "__main__":
+    r_bullet = BulletArm1Dof()
+    r_bullet.InitPybullet(pybullet.GUI)
+
+    r_arm = r_bullet.LoadRobot()
+
+    while True:
+        pass
+
+else:
+    r_bullet = BulletArm1Dof()
+    r_bullet.InitPybullet(bullet_connect=pybullet.GUI)
+
+    r_arm = r_bullet.LoadRobot()