Add muscle-controlled robot

arm_1dof/bullet_arm_1dof.py

@@ -1,6 +1,8 @@
 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
 
@@ -16,6 +18,7 @@ class BulletArm1Dof:
         self._server_id = -1
         self._skel_arm = None
         self._skel_arm_id = -1
+        self._skel_muscles = None
 
     @property
     def server_id(self):
@@ -41,9 +44,33 @@ class BulletArm1Dof:
 
         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_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)