diff --git a/arm_1dof/bullet_arm_1dof.py b/arm_1dof/bullet_arm_1dof.py new file mode 100644 index 0000000..5291446 --- /dev/null +++ b/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 + \ No newline at end of file diff --git a/arm_1dof/robot_arm_1dof.py b/arm_1dof/robot_arm_1dof.py new file mode 100644 index 0000000..9081cf9 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/arm_1dof/test_arm_1dof.py b/arm_1dof/test_arm_1dof.py new file mode 100755 index 0000000..7407b4c --- /dev/null +++ b/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() \ No newline at end of file