Improved test

Author: EnricoMingo

bindings/python/tests/pyopensot_hpipm_tests.py

@@ -7,7 +7,7 @@
 import unittest
 from collections import deque
 
-SHOW_PLOTS = False
+SHOW_PLOTS = True
 
 def plot_trajectory(Ns, x_value, u_value, zmp_refs, dt):
     plt.figure(figsize=(8, 4))
@@ -20,6 +20,7 @@ def plot_trajectory(Ns, x_value, u_value, zmp_refs, dt):
     plt.legend()
     plt.tight_layout()
     plt.show()
+
 def zmp_pattern(ns):
     zref = np.zeros((2, ns))
     for i in range(ns):
@@ -41,10 +42,8 @@ def zmp_pattern(ns):
 tf = 3.0 # final time
 
 vars = list()
-for i in range(Ns):
-    vars.append((f"x{i}", nx))
-    vars.append((f"u{i}", nu))
-vars.append((f"x{Ns}", nx))
+vars.append(("x", nx))
+vars.append(("u", nu))
 
 variables = OptvarHelper(vars)
 print(f"variables.getSize(): {variables.getSize()}")
@@ -58,63 +57,59 @@ def euler(x, xdot, dt):
 
 h = 0.83  # height of the CoM
 dt = tf/Ns
-integration = list()
 
-for i in range(Ns):
-    x0 = variables.getVariable(f"x{i}")
-    u0 = variables.getVariable(f"u{i}")
 
-    r = x0[0:2]
-    rdot = x0[2:]
-    rddot = lipm(r, u0, h)
+x = variables.getVariable(f"x")
+u = variables.getVariable(f"u")
 
-    xdot0 = AffineHelper.pile(rdot, rddot)
+r = x[0:2]
+rdot = x[2:]
+rddot = lipm(r, u, h)
 
-    integration_ = euler(x0, xdot0, dt)
-    integration.append(GenericTask(f"integration_{i}", integration_.getM(), -integration_.getq()))
+xdot = AffineHelper.pile(rdot, rddot)
 
+integration_ = euler(x, xdot, dt)
+integration = GenericTask(f"integration", integration_.getM(), -integration_.getq())
+
+
+print(f"integration.getA()\n: {integration.getA()}")
 
-print(f"integration[0].getA()\n: {integration[0].getA()}")
 # to retrieve A and B for each stage:
-print(f"integration[0].getA() @ variables.getVariable(x0).getM().T \n: {integration[0].getA() @ variables.getVariable("x0").getM().T}")
-print(f"integration[0].getA() @ variables.getVariable(u0).getM().T \n: {integration[0].getA() @ variables.getVariable("u0").getM().T}")
+print(f"integration.getA() @ x.getM().T \n: {integration.getA() @ x.getM().T}")
+print(f"integration.getA() @ u.getM().T \n: {integration.getA() @ u.getM().T}")
+
+
 
 solver = hpipmoc.hpipmOC(Ns)
 for i in range(Ns):
-    Ai = integration[i].getA() @ variables.getVariable(f"x{i}").getM().T
-    Bi = integration[i].getA() @ variables.getVariable(f"u{i}").getM().T
+    Ai = integration.getA() @ x.getM().T
+    Bi = integration.getA() @ u.getM().T
     bi = np.zeros((nx, 1))
     solver.setStageDynamics(i, Ai, Bi, bi)
 
-#xinit = variables.getVariable("x0") + np.array([0., 0., 0, 0.])
-#initial_state = GenericTask("initial_state", xinit.getM(), -xinit.getq())
-
-#print(f"initial_state.getA()\n: {initial_state.getA()}")
-#print(f"initial_state.getA() @ variables.getVariable(x0).getM().T\n: {initial_state.getA() @ variables.getVariable("x0").getM().T}")
 
-#id = [0, 1, 2, 3] # these are the indices of the first 4 state variables (x) at stage 0
-#solver.setBoxConstraintsX(0, id, initial_state.getb(), initial_state.getb())
 
 zmp_tasks = list()
 for i in range(Ns):
-    min_ui = GenericTask("zmp_tracking", variables.getVariable(f"u{i}").getM(), -variables.getVariable(f"u{i}").getq())
+    min_ui = GenericTask("zmp_tracking", u.getM(), -u.getq())
     min_ui.setWeight(1e4 * np.array([[1, 0], [0, 1]]))
     zmp_tasks.append(min_ui)
 
 print(f"zmp_tasks[2].getA()\n: {zmp_tasks[2].getA()}")
-print(f"zmp_tasks[2].getA() @ variables.getVariable(u2).getM().T\n: {zmp_tasks[2].getA() @ variables.getVariable("u2").getM().T}")
+print(f"zmp_tasks[2].getA() @ u.T\n: {zmp_tasks[2].getA() @ u.getM().T}")
 
 print(f"zmp_tasks[2].getb()\n: {zmp_tasks[2].getb()}")
 print(f"zmp_tasks[2].getWeight()\n: {zmp_tasks[2].getWeight()}")
 
+
 zmp_refs = zmp_pattern(Ns)
 for i in range(Ns):
     zmp_tasks[i].setb(zmp_refs[:, i])
     zmp_tasks[i].update()
 
 x_tasks = list()
 for i in range(Ns+1):
-    min_xi = GenericTask("min_x", variables.getVariable(f"x{i}").getM(), variables.getVariable(f"x{i}").getq())
+    min_xi = GenericTask("min_x", x.getM(), -x.getq())
     Q = 1e-3 * np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
     if i == Ns:
         Q = 1e6 * np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
@@ -123,10 +118,9 @@ def euler(x, xdot, dt):
 
 
 for i in range(Ns):
-    solver.setLSCost(i,
-                     x_tasks[i].getA() @ variables.getVariable(f"x{i}").getM().T, x_tasks[i].getWeight(), x_tasks[i].getb(),
-                     zmp_tasks[i].getA() @ variables.getVariable(f"u{i}").getM().T, zmp_tasks[i].getWeight(), zmp_tasks[i].getb())
-solver.setLSCost(Ns, x_tasks[Ns].getA() @ variables.getVariable(f"x{Ns}").getM().T, x_tasks[Ns].getWeight(), x_tasks[Ns].getb())
+    solver.setLSCost(i, x_tasks[i].getA() @ x.getM().T, x_tasks[i].getWeight(), x_tasks[i].getb(),
+                     zmp_tasks[i].getA() @ u.getM().T, zmp_tasks[i].getWeight(), zmp_tasks[i].getb())
+solver.setLSCost(Ns, x_tasks[Ns].getA() @ x.getM().T, x_tasks[Ns].getWeight(), x_tasks[Ns].getb())
 
 tic()
 success = solver.solve(np.array([0., 0., 0, 0.]))
@@ -154,6 +148,7 @@ def euler(x, xdot, dt):
 else:
     print("HPIPM could not solve the problem!")
 
+
 # zeroing references
 for i in range(Ns):
     zmp_tasks[i].setb(np.zeros((2, 1)))
@@ -182,8 +177,6 @@ def euler(x, xdot, dt):
     ry.append(x0[1,0])
     ryplot = ry.popleft()
 
-    #initial_state.setb(x0)
-    #initial_state.update()
 
     # shift reference to left
     for j in range(1, Ns):
@@ -197,10 +190,9 @@ def euler(x, xdot, dt):
     # update internal qp
     #solver.setBoxConstraintsX(0, id, initial_state.getb(), initial_state.getb())
     for i in range(Ns):
-        solver.setLSCost(i,
-                         x_tasks[i].getA() @ variables.getVariable(f"x{i}").getM().T, x_tasks[i].getWeight(), x_tasks[i].getb(),
-                         zmp_tasks[i].getA() @ variables.getVariable(f"u{i}").getM().T, zmp_tasks[i].getWeight(), zmp_tasks[i].getb())
-    solver.setLSCost(Ns, x_tasks[Ns].getA() @ variables.getVariable(f"x{Ns}").getM().T, x_tasks[Ns].getWeight(), x_tasks[Ns].getb())
+        solver.setLSCost(i, x_tasks[i].getA() @ x.getM().T, x_tasks[i].getWeight(), x_tasks[i].getb(),
+                         zmp_tasks[i].getA() @ u.getM().T, zmp_tasks[i].getWeight(), zmp_tasks[i].getb())
+    solver.setLSCost(Ns, x_tasks[Ns].getA() @ x.getM().T, x_tasks[Ns].getWeight(), x_tasks[Ns].getb())
 
     tic()
     success = solver.solve(x0)
@@ -213,8 +205,6 @@ def euler(x, xdot, dt):
         u_value[:, i] = solution[i].u
     x_value[:, Ns] = solution[Ns].x
 
-    utest.assertTrue(( np.linalg.norm(u_value[:,0] - zmp_tasks[0].getb(), ord=2) <= 1e-4).all())
-
     rdot = x_value[2:4, 0].flatten()
     rddot = lipm(x_value[0:2, 0], u_value[:, 0], h)