Add muscle-controlled robot

.gitignore

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

arm_1dof/bullet_arm_1dof.py

@@ -0,0 +1,76 @@
+import arm_1dof.robot_arm_1dof
+import arm_1dof.robot_arm_muscle_1dof
+import pybullet
+from tokenize import String
+from MuscleModelPython import MUSCLE_INDEX, JOINT_INDEX
+
+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
+        self._skel_muscles = None
+
+    @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) -> arm_1dof.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 = arm_1dof.robot_arm_1dof.RobotArm1Dof(bullet_ctrl=self, bullet_body_id=self._skel_arm_id)
+
+        return self._skel_arm
+
+    def LoadMuscleRobot(self, robot_sdf_filename=URDF_MODEL_FILENAME) -> arm_1dof.robot_arm_muscle_1dof.RobotArmMuscle1Dof:
+        """Load robot from sdf file, add muscles"""
+        self._skel_arm_id = pybullet.loadURDF(fileName=robot_sdf_filename, useFixedBase=True, physicsClientId=self._server_id)
+
+        self._skel_arm = arm_1dof.robot_arm_muscle_1dof.RobotArmMuscle1Dof(bullet_ctrl=self, bullet_body_id=self._skel_arm_id)
+        self._skel_muscles = self._skel_arm.muscles
+
+        return self._skel_arm
+
+    def Simulate(self, sim_time) -> None:
+        t = 0
+        while t < sim_time:
+            pybullet.stepSimulation(physicsClientId=self._server_id)
+            
+            if self._skel_muscles:
+                # Update muscle joint states
+                self._skel_arm.UpdateStats()
+                self._skel_muscles.SetJointPos(JOINT_INDEX(0), self._skel_arm.JointPos(0))
+                self._skel_muscles.SetJointPos(JOINT_INDEX(1), self._skel_arm.JointPos(1))
+                self._skel_muscles.SetJointVel(JOINT_INDEX(0), self._skel_arm.JointVel(0))
+                self._skel_muscles.SetJointVel(JOINT_INDEX(1), self._skel_arm.JointVel(1))
+
+                self._skel_muscles.Integrate(self._timestep)
+
+                # Apply muscle torques to robot
+                torques = [ self._skel_muscles.GetTorque(JOINT_INDEX(0)), self._skel_muscles.GetTorque(JOINT_INDEX(1))]
+                self._skel_arm.SetJointTorques(joint_ids=[0,1], torques=torques)
+
+            t += self._timestep
+            

arm_1dof/robot_arm_1dof.py

@@ -0,0 +1,81 @@
+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]
+
+        # Disable position and velocity based joint motor controls
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=[RobotArm1Dof.SHOULDER_A_JOINT_ID, RobotArm1Dof.ELBOW_JOINT_ID], \
+            controlMode=pybullet.POSITION_CONTROL, forces=[0.0, 0.0], \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=[RobotArm1Dof.SHOULDER_A_JOINT_ID, RobotArm1Dof.ELBOW_JOINT_ID], \
+            controlMode=pybullet.VELOCITY_CONTROL, forces=[0.0, 0.0], \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        # Initialize joint/link state storage
+        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, \
+            computeLinkVelocity=True,
+            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 SetJointTorques(self, joint_ids, torques) -> None:
+        """Sets torques applied to specified joints."""
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=joint_ids, \
+            controlMode=pybullet.TORQUE_CONTROL, forces=torques, \
+            physicsClientId=self._bullet_ctrl.server_id)

arm_1dof/robot_arm_muscle_1dof.py

@@ -0,0 +1,89 @@
+from ntpath import join
+import numpy as np
+from pandas import array
+import pybullet
+from MuscleModelPython import TwoMuscleModel, MUSCLE_INDEX
+
+class RobotArmMuscle1Dof:
+    """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)
+        
+        # Init muscles
+        self.muscles = TwoMuscleModel()
+
+        # Set mass
+        dyn_info = pybullet.getDynamicsInfo(self._body_id, RobotArmMuscle1Dof.UPPER_ARM_LINK_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+        self._upper_arm_mass = dyn_info[0]
+
+        dyn_info = pybullet.getDynamicsInfo(self._body_id, RobotArmMuscle1Dof.FOREARM_LINK_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+        self._lower_arm_mass = dyn_info[1]
+
+        # Disable position and velocity based joint motor controls
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=[RobotArmMuscle1Dof.SHOULDER_A_JOINT_ID, RobotArmMuscle1Dof.ELBOW_JOINT_ID], \
+            controlMode=pybullet.POSITION_CONTROL, forces=[0.0, 0.0], \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=[RobotArmMuscle1Dof.SHOULDER_A_JOINT_ID, RobotArmMuscle1Dof.ELBOW_JOINT_ID], \
+            controlMode=pybullet.VELOCITY_CONTROL, forces=[0.0, 0.0], \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        # Initialize joint/link state storage
+        self._joint_stats = {}
+        self._ee_link_state = None
+
+        self.UpdateStats()
+
+    def UpdateStats(self):
+        """Query state info from bullet"""
+        # Retrieve joint info
+        self._joint_stats[RobotArmMuscle1Dof.SHOULDER_A_JOINT_ID] = pybullet.getJointState(self._body_id, RobotArmMuscle1Dof.SHOULDER_A_JOINT_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        self._joint_stats[RobotArmMuscle1Dof.ELBOW_JOINT_ID] = pybullet.getJointState(self._body_id, RobotArmMuscle1Dof.ELBOW_JOINT_ID, \
+            physicsClientId=self._bullet_ctrl.server_id)
+
+        self._ee_link_state = pybullet.getLinkState(self._body_id, RobotArmMuscle1Dof.HAND_LINK_ID, \
+            computeLinkVelocity=True,
+            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 SetMuscleActivations(self, muscle_activations) -> None:
+        """Set muscle activations. Takes a 2-dimensional array with activations"""
+
+
+    def SetJointTorques(self, joint_ids, torques) -> None:
+        """Sets torques applied to specified joints."""
+        pybullet.setJointMotorControlArray(self._body_id, jointIndices=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 arm_1dof.bullet_arm_1dof import BulletArm1Dof
+from arm_1dof.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()