From 54b89da217ba74dadfb83f4d6e257f4ab2644873 Mon Sep 17 00:00:00 2001 From: edmundwsy Date: Mon, 21 Mar 2022 23:16:52 +0100 Subject: [PATCH] Debug: fix animation bug --- high_mpc/run_mpc.py | 26 +++++++++++++------------- high_mpc/simulation/animation.py | 4 ++-- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/high_mpc/run_mpc.py b/high_mpc/run_mpc.py index 6311721..3bbc995 100644 --- a/high_mpc/run_mpc.py +++ b/high_mpc/run_mpc.py @@ -15,7 +15,7 @@ # def arg_parser(): parser = argparse.ArgumentParser() - parser.add_argument('--save_video', type=bool, default=False, + parser.add_argument('--save_video', type=bool, default=True, help="Save the animation as a video file") return parser @@ -36,7 +36,7 @@ def run_mpc(env): update = False if t>= env.sim_T: update = True - # yield [info, t, update] + yield [info, t, update] def main(): # @@ -53,20 +53,20 @@ def main(): sim_visual = SimVisual(env) # - run_mpc(env) + # run_mpc(env) - # run_frame = partial(run_mpc, env) - # ani = animation.FuncAnimation(sim_visual.fig, sim_visual.update, frames=run_frame, - # init_func=sim_visual.init_animate, interval=100, blit=True, repeat=False) + run_frame = partial(run_mpc, env) + ani = animation.FuncAnimation(sim_visual.fig, sim_visual.update, frames=run_frame, + init_func=sim_visual.init_animate, interval=100, blit=True, repeat=False) - # # - # if args.save_video: - # writer = animation.writers["ffmpeg"] - # writer = writer(fps=10, metadata=dict(artist='Me'), bitrate=1800) - # ani.save("MPC_0.mp4", writer=writer) + # + if args.save_video: + writer = animation.writers["ffmpeg"] + writer = writer(fps=10, metadata=dict(artist='Me'), bitrate=1800) + ani.save("MPC_0.mp4", writer=writer) - # plt.tight_layout() - # plt.show() + plt.tight_layout() + plt.show() if __name__ == "__main__": main() diff --git a/high_mpc/simulation/animation.py b/high_mpc/simulation/animation.py index 4890106..3a62891 100644 --- a/high_mpc/simulation/animation.py +++ b/high_mpc/simulation/animation.py @@ -282,8 +282,8 @@ def update(self, data_info): self.l_quad_pred_traj.set_3d_properties(pred_quad_traj[:, 2]) if quad_pos_arr[-1, 0] <= 2.0: # plot planner trajectory - self.l_pend_pred_traj.set_data(pred_pend_traj[opt_idx, 0], pred_pend_traj[opt_idx, 1]) - self.l_pend_pred_traj.set_3d_properties(pred_pend_traj[opt_idx, 2]) + self.l_pend_pred_traj.set_data(np.array([pred_pend_traj[opt_idx, 0]]), np.array([pred_pend_traj[opt_idx, 1]])) + self.l_pend_pred_traj.set_3d_properties(np.array([pred_pend_traj[opt_idx, 2]])) # self.l_pend.set_data([self.pivot_point[0], pend_obs[kPosX]], [self.pivot_point[1], pend_obs[kPosY]] )