refactor to make minitaur example more general

examples/pybullet/minitaur_evaluate.py

@@ -1,12 +1,15 @@
 from minitaur import Minitaur
-import time
-import numpy as np
 import pybullet as p
-import math
+import numpy as np
+import time
 import sys
+import math
 
 minitaur = None
 
+evaluate_func_map = dict()
+
+
 def current_position():
   global minitaur
   position = minitaur.getBasePosition()
@@ -19,8 +22,43 @@ def is_fallen():
   localUp = rotMat[6:]
   return np.dot(np.asarray([0, 0, 1]), np.asarray(localUp)) < 0
 
+def evaluate_desired_motorAngle_8Amplitude8Phase(i, params):
+  nMotors = 8
+  speed = 0.35
+  for jthMotor in range(nMotors):
+    joint_values[jthMotor] = math.sin(i*speed + params[nMotors + jthMotor])*params[jthMotor]*+1.57
+  return joint_values
 
-def evaluate_params_hop(params, urdfRoot='', timeStep=0.01, maxNumSteps=1000, sleepTime=0):
+def evaluate_desired_motorAngle_2Amplitude4Phase(i, params):
+  speed = 0.35
+  phaseDiff = params[2]
+  a0 = math.sin(i * speed) * params[0] + 1.57
+  a1 = math.sin(i * speed + phaseDiff) * params[1] + 1.57
+  a2 = math.sin(i * speed + params[3]) * params[0] + 1.57
+  a3 = math.sin(i * speed + params[3] + phaseDiff) * params[1] + 1.57
+  a4 = math.sin(i * speed + params[4] + phaseDiff) * params[1] + 1.57
+  a5 = math.sin(i * speed + params[4]) * params[0] + 1.57
+  a6 = math.sin(i * speed + params[5] + phaseDiff) * params[1] + 1.57
+  a7 = math.sin(i * speed + params[5]) * params[0] + 1.57
+  joint_values = [a0, a1, a2, a3, a4, a5, a6, a7]
+  return joint_values
+
+def evaluate_desired_motorAngle_hop(i, params):
+  amplitude = params[0]
+  speed = params[1]
+  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]
+  return joint_values
+
+
+evaluate_func_map['evaluate_desired_motorAngle_8Amplitude8Phase'] = evaluate_desired_motorAngle_8Amplitude8Phase
+evaluate_func_map['evaluate_desired_motorAngle_2Amplitude4Phase'] = evaluate_desired_motorAngle_2Amplitude4Phase
+evaluate_func_map['evaluate_desired_motorAngle_hop'] = evaluate_desired_motorAngle_hop
+
+
+
+def evaluate_params(evaluateFunc, params, objectiveParams, urdfRoot='', timeStep=0.01, maxNumSteps=1000, sleepTime=0):
   print('start evaluation')
   beforeTime = time.time()
   p.resetSimulation()
@@ -29,18 +67,18 @@ def evaluate_params_hop(params, urdfRoot='', timeStep=0.01, maxNumSteps=1000, sl
   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
+  total_energy = 0
 
   for i in range(maxNumSteps):
-    a1 = math.sin(i*speed)*amplitude+1.57
+    torques = minitaur.getMotorTorques()
-    a2 = math.sin(i*speed+3.14)*amplitude+1.57
+    velocities = minitaur.getMotorVelocities()
-    joint_values = [a1, 1.57, a2, 1.57, 1.57, a1, 1.57, a2]
+    total_energy += np.dot(np.fabs(torques), np.fabs(velocities)) * timeStep
+
+    joint_values = evaluate_func_map[evaluateFunc](i, params)
     minitaur.applyAction(joint_values)
     p.stepSimulation()
     if (is_fallen()):
@@ -53,7 +91,9 @@ def evaluate_params_hop(params, urdfRoot='', timeStep=0.01, maxNumSteps=1000, sl
 
   print(' ')
 
+  alpha = objectiveParams[0]
   final_distance = np.linalg.norm(start_position - current_position())
+  finalReturn = final_distance - alpha * total_energy
   elapsedTime = time.time() - beforeTime
-  print ("trial for amplitude", amplitude, "speed", speed, "final_distance", final_distance, "elapsed_time", elapsedTime)
+  print ("trial for ", params, " final_distance", final_distance, "total_energy", total_energy, "finalReturn", finalReturn, "elapsed_time", elapsedTime)
-  return final_distance
+  return finalReturn

examples/pybullet/minitaur_test.py

@@ -1,7 +1,7 @@
 import pybullet as p
-
 from minitaur import Minitaur
-import minitaur_evaluate
+from minitaur_evaluate import *
+
 import time
 import math
 import numpy as np
@@ -12,29 +12,12 @@ def main(unused_args):
   if (c<0):
       c = p.connect(p.GUI)
 
-  amplitude = 0.24795664427
+  params = [0.1903581461951056, 0.0006732219568880068, 0.05018085615283363, 3.219916795483583, 6.2406418167980595, 4.189869754607539]
-  speed = 0.2860877729434
+  evaluate_func = 'evaluate_desired_motorAngle_2Amplitude4Phase'
+  energy_weight = 0.01
 
-  final_distance = minitaur_evaluate.evaluate_params_hop(params=[amplitude, speed], timeStep=timeStep, sleepTime=timeStep)
+  finalReturn = evaluate_params(evaluateFunc = evaluate_func, params=params, objectiveParams=[energy_weight], timeStep=timeStep, sleepTime=timeStep)
-  print(final_distance)
 
-  # p.resetSimulation()
+  print(finalReturn)
-  # 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)