diff --git a/python/quadruped_reactive_walking/WB_MPC_Wrapper.py b/python/quadruped_reactive_walking/WB_MPC_Wrapper.py index 71a9e2f1c1073774b3b6b7074f17ca5b4a42c8f0..b6c71a7a1384a5d1464d804b96ad6f0e0c2c9a90 100644 --- a/python/quadruped_reactive_walking/WB_MPC_Wrapper.py +++ b/python/quadruped_reactive_walking/WB_MPC_Wrapper.py @@ -46,7 +46,7 @@ class MPC_Wrapper: self.in_xs = Array("d", [0] * ((pd.T + 1) * pd.nx)) self.in_us = Array("d", [0] * (pd.T * pd.nu)) self.in_footstep = Array("d", [0] * 12) - self.in_gait = Array("d", [0] * ( (pd.T+1) * 4)) + self.in_gait = Array("d", [0] * ((pd.T + 1) * 4)) self.out_xs = Array("d", [0] * ((pd.T + 1) * pd.nx)) self.out_us = Array("d", [0] * (pd.T * pd.nu)) self.out_k = Array("d", [0] * (pd.T * pd.nu * pd.ndx)) @@ -66,8 +66,10 @@ class MPC_Wrapper: k (int): Number of inv dynamics iterations since the start of the simulation """ if self.multiprocessing: + print("-- RUNNING ASYNCHRONOUS --") self.run_MPC_asynchronous(k, x0, tasks, gait, xs, us) else: + print("-- RUNNING SYNCHRONOUS --") self.run_MPC_synchronous(x0, tasks, gait, guess) def get_latest_result(self): @@ -97,14 +99,81 @@ class MPC_Wrapper: """ Run the MPC (synchronous version) """ + self.ocp.solve(x0, tasks, gait, guess) - self.last_available_result.xs,\ - self.last_available_result.acs,\ - self.last_available_result.us,\ - self.last_available_result.fs,\ - self.last_available_result.solving_duration = self.ocp.get_results() + ( + self.last_available_result.xs, + self.last_available_result.acs, + self.last_available_result.us, + self.last_available_result.fs, + self.last_available_result.solving_duration, + ) = self.ocp.get_results() self.new_result.value = True + import pinocchio as pin + from matplotlib import pyplot as plt + + index6 = [1, 3, 5, 2, 4, 6] + index12 = [1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12] + + us = np.array(self.last_available_result.us) + xs_raw = np.array(self.last_available_result.xs) + xs = np.zeros((xs_raw.shape[0], xs_raw.shape[1] - 1)) + + xs[:, :3] = xs_raw[:, :3] + for i in range(xs.shape[0]): + xs[i, 3:6] = pin.rpy.matrixToRpy( + pin.Quaternion(xs_raw[i, 3:7]).toRotationMatrix() + ) + xs[:, 6:] = xs_raw[:, 7:] + + plt.figure() + N = min(10, len(self.ocp.log_vel_base)) + for k in range(N): + plt.plot(self.ocp.log_vel_base[k]) + # plt.plot(self.ocp.log_friction_cone[k]) + plt.legend([str(k) for k in range(N)]) + + plt.figure() + plt.suptitle("Linear components") + lgd = ["$x$", "$y$", "$z$", "$\dot{x}$", "$\dot{y}$", "$\dot{z}$"] + for i in range(6): + plt.subplot(3, 2, index6[i]) + if i < 3: + plt.plot(xs[:, i]) + else: + plt.plot(xs[:, 18 + i - 3]) + plt.legend([lgd[i]]) + + plt.figure() + plt.suptitle("Angular components") + lgd = [ + "$\phi$", + "$\theta$", + "$\psi$", + "$\dot{\omega_{x}}$", + "$\dot{\omega_{y}}$", + "$\dot{\omega_{z}}$", + ] + for i in range(6): + plt.subplot(3, 2, index6[i]) + if i < 3: + plt.plot(xs[:, 3 + i]) + else: + plt.plot(xs[:, 18 + i]) + plt.legend([lgd[i]]) + + plt.figure() + plt.suptitle("Torques") + for i in range(12): + plt.subplot(3, 4, index12[i]) + plt.plot(us[:, i]) + plt.show() + + from IPython import embed + + embed() + def run_MPC_asynchronous(self, k, x0, tasks, gait, xs, us): """ Run the MPC (asynchronous version) @@ -159,7 +228,9 @@ class MPC_Wrapper: with self.in_footstep.get_lock(): np.frombuffer(self.in_footstep.get_obj()).reshape((3, 4))[:, :] = tasks with self.in_gait.get_lock(): - np.frombuffer(self.in_gait.get_obj()).reshape((self.pd.T + 1, 4))[:, :] = gait + np.frombuffer(self.in_gait.get_obj()).reshape((self.pd.T + 1, 4))[ + :, : + ] = gait if xs is None or us is None: self.in_warm_start.value = False