now minitaur class can output joint angles, velocities and torques. I also extract evaluate functions to a file

3 files changed, 137 insertions, 31 deletions

diff --git a/examples/pybullet/minitaur.py b/examples/pybullet/minitaur.py
index 5364d8124..58aab4a38 100644
--- a/examples/pybullet/minitaur.py
+++ b/examples/pybullet/minitaur.py
@@ -2,15 +2,11 @@ import pybullet as p
 import numpy as np
 
 class Minitaur:
-  def __init__(self):
+  def __init__(self, urdfRootPath=''):
+    self.urdfRootPath = urdfRootPath
     self.reset()
 
-  def reset(self):
-    self.quadruped = p.loadURDF("quadruped/quadruped.urdf",0,0,.3)
-    self.kp = 1
-    self.kd = 0.1
-    self.maxForce = 3.5
-    self.motorDir = [1, -1, 1, -1, -1, 1, -1, 1]
+  def buildJointNameToIdDict(self):
     nJoints = p.getNumJoints(self.quadruped)
     self.jointNameToId = {}
     for i in range(nJoints):
@@ -20,13 +16,39 @@ class Minitaur:
     for i in range(100):
       p.stepSimulation()
 
+  def buildMotorIdList(self):
+    self.motorIdList.append(self.jointNameToId['motor_front_leftR_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_front_leftL_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_back_leftR_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_back_leftL_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_front_rightL_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_front_rightR_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_back_rightL_joint'])
+    self.motorIdList.append(self.jointNameToId['motor_back_rightR_joint'])
+
+
+  def reset(self):
+    self.quadruped = p.loadURDF("%s/quadruped/quadruped.urdf" % self.urdfRootPath,0,0,.3)
+    self.kp = 1
+    self.kd = 0.1
+    self.maxForce = 3.5
+    self.nMotors = 8
+    self.motorIdList = []
+    self.motorDir = [1, -1, 1, -1, -1, 1, -1, 1]
+    self.buildJointNameToIdDict()
+    self.buildMotorIdList()
+
+
   def disableAllMotors(self):
     nJoints = p.getNumJoints(self.quadruped)
     for i in range(nJoints):
       p.setJointMotorControl2(bodyIndex=self.quadruped, jointIndex=i, controlMode=p.VELOCITY_CONTROL, force=0)
 
+  def setMotorAngleById(self, motorId, desiredAngle):
+    p.setJointMotorControl2(bodyIndex=self.quadruped, jointIndex=motorId, controlMode=p.POSITION_CONTROL, targetPosition=desiredAngle, positionGain=self.kp, velocityGain=self.kd, force=self.maxForce)
+
   def setMotorAngleByName(self, motorName, desiredAngle):
-    p.setJointMotorControl2(bodyIndex=self.quadruped, jointIndex=self.jointNameToId[motorName], controlMode=p.POSITION_CONTROL, targetPosition=desiredAngle, positionGain=self.kp, velocityGain=self.kd, force=self.maxForce)
+    self.setMotorAngleById(self.jointNameToId[motorName], desiredAngle)
 
   def resetPose(self):
     #right front leg
@@ -76,11 +98,29 @@ class Minitaur:
 
   def applyAction(self, motorCommands):
     motorCommandsWithDir = np.multiply(motorCommands, self.motorDir)
-    self.setMotorAngleByName('motor_front_leftR_joint', motorCommandsWithDir[0])
-    self.setMotorAngleByName('motor_front_leftL_joint', motorCommandsWithDir[1])
-    self.setMotorAngleByName('motor_back_leftR_joint', motorCommandsWithDir[2])
-    self.setMotorAngleByName('motor_back_leftL_joint', motorCommandsWithDir[3])
-    self.setMotorAngleByName('motor_front_rightL_joint', motorCommandsWithDir[4])
-    self.setMotorAngleByName('motor_front_rightR_joint', motorCommandsWithDir[5])
-    self.setMotorAngleByName('motor_back_rightL_joint', motorCommandsWithDir[6])
-    self.setMotorAngleByName('motor_back_rightR_joint', motorCommandsWithDir[7])
+    for i in range(self.nMotors):
+      self.setMotorAngleById(self.motorIdList[i], motorCommandsWithDir[i])
+
+  def getMotorAngles(self):
+    motorAngles = []
+    for i in range(self.nMotors):
+      jointState = p.getJointState(self.quadruped, self.motorIdList[i])
+      motorAngles.append(jointState[0])
+    motorAngles = np.multiply(motorAngles, self.motorDir)
+    return motorAngles
+
+  def getMotorVelocities(self):
+    motorVelocities = []
+    for i in range(self.nMotors):
+      jointState = p.getJointState(self.quadruped, self.motorIdList[i])
+      motorVelocities.append(jointState[1])
+    motorVelocities = np.multiply(motorVelocities, self.motorDir)
+    return motorVelocities
+
+  def getMotorTorques(self):
+    motorTorques = []
+    for i in range(self.nMotors):
+      jointState = p.getJointState(self.quadruped, self.motorIdList[i])
+      motorTorques.append(jointState[3])
+    motorTorques = np.multiply(motorTorques, self.motorDir)
+    return motorTorques
diff --git a/examples/pybullet/minitaur_evaluate.py b/examples/pybullet/minitaur_evaluate.py
new file mode 100644
index 000000000..73212f125
--- /dev/null
+++ b/examples/pybullet/minitaur_evaluate.py
@@ -0,0 +1,59 @@
+from minitaur import Minitaur
+import time
+import numpy as np
+import pybullet as p
+import math
+import sys
+
+minitaur = None
+
+def current_position():
+  global minitaur
+  position = minitaur.getBasePosition()
+  return np.asarray(position)
+
+def is_fallen():
+  global minitaur
+  orientation = minitaur.getBaseOrientation()
+  rotMat = p.getMatrixFromQuaterion(orientation)
+  localUp = rotMat[6:]
+  return np.dot(np.asarray([0, 0, 1]), np.asarray(localUp)) < 0
+
+
+def evaluate_params_hop(params, urdfRoot='', timeStep=0.01, maxNumSteps=1000, sleepTime=0):
+  print('start evaluation')
+  beforeTime = time.time()
+  p.resetSimulation()
+
+  p.setTimeStep(timeStep)
+  p.loadURDF("%s/plane.urdf" % urdfRoot)
+  p.setGravity(0,0,-10)
+
+  amplitude = params[0]
+  speed = params[1]
+
+  global minitaur
+  minitaur = Minitaur(urdfRoot)
+  start_position = current_position()
+  last_position = None  # for tracing line
+
+  for i in range(maxNumSteps):
+    a1 = math.sin(i*speed)*amplitude+1.57
+    a2 = math.sin(i*speed+3.14)*amplitude+1.57
+    joint_values = [a1, 1.57, a2, 1.57, 1.57, a1, 1.57, a2]
+    minitaur.applyAction(joint_values)
+    p.stepSimulation()
+    if (is_fallen()):
+      break
+
+    if i % 100 == 0:
+      sys.stdout.write('.')
+      sys.stdout.flush()
+    time.sleep(sleepTime)
+
+  print(' ')
+
+  final_distance = np.linalg.norm(start_position - current_position())
+  elapsedTime = time.time() - beforeTime
+  print ("trial for amplitude", amplitude, "speed", speed, "final_distance", final_distance, "elapsed_time", elapsedTime)
+  return final_distance
diff --git a/examples/pybullet/minitaur_test.py b/examples/pybullet/minitaur_test.py
index ddc2ac136..2b0726f6f 100644
--- a/examples/pybullet/minitaur_test.py
+++ b/examples/pybullet/minitaur_test.py
@@ -1,6 +1,7 @@
 import pybullet as p
 
 from minitaur import Minitaur
+import minitaur_evaluate
 import time
 import math
 import numpy as np
@@ -10,24 +11,30 @@ def main(unused_args):
   c = p.connect(p.SHARED_MEMORY)
   if (c<0):
       c = p.connect(p.GUI)
-  p.resetSimulation()
-  p.setTimeStep(timeStep)
-  p.loadURDF("plane.urdf")
-  p.setGravity(0,0,-10)
 
-  minitaur = Minitaur()
   amplitude = 0.24795664427
   speed = 0.2860877729434
-  for i in range(1000):
-    a1 = math.sin(i*speed)*amplitude+1.57
-    a2 = math.sin(i*speed+3.14)*amplitude+1.57
-    joint_values = [a1, 1.57, a2, 1.57, 1.57, a1, 1.57, a2]
-    minitaur.applyAction(joint_values)
-
-    p.stepSimulation()
-#    print(minitaur.getBasePosition())
-    time.sleep(timeStep)
-  final_distance = np.linalg.norm(np.asarray(minitaur.getBasePosition()))
+
+  final_distance = minitaur_evaluate.evaluate_params_hop(params=[amplitude, speed], timeStep=timeStep, sleepTime=timeStep)
   print(final_distance)
 
+  # p.resetSimulation()
+  # p.setTimeStep(timeStep)
+  # p.loadURDF("plane.urdf")
+  # p.setGravity(0,0,-10)
+
+  # minitaur = Minitaur()
+
+  # for i in range(1000):
+  #   a1 = math.sin(i*speed)*amplitude+1.57
+  #   a2 = math.sin(i*speed+3.14)*amplitude+1.57
+  #   joint_values = [a1, 1.57, a2, 1.57, 1.57, a1, 1.57, a2]
+  #   minitaur.applyAction(joint_values)
+  #   torques = minitaur.getMotorTorques()
+  #   print(torques)
+  #   p.stepSimulation()
+  #   time.sleep(timeStep)
+  # final_distance = np.linalg.norm(np.asarray(minitaur.getBasePosition()))
+  # print(final_distance)
+
 main(0)