plotting.py : get legend first before adding other labels

Author: ManifoldFR

bindings/python/aligator/utils/plotting.py

@@ -21,8 +21,8 @@ def plot_convergence(
         dual_infeas.append(res.dual_infeas)
     plot_pd_errs(ax, prim_infeas, dual_infeas)
     ax.grid(axis="y", which="major")
-    labels = [
-        "$\\epsilon_\\mathrm{tol}$",
+    handles, labels = ax.get_legend_handles_labels()
+    labels += [
         "Prim. err $p$",
         "Dual err $d$",
     ]
@@ -42,7 +42,7 @@ def plot_convergence(
             ax.vlines(al_change_idx, *ax.get_ylim(), colors="gray", lw=4.0, alpha=0.5)
 
     ax.legend(labels=labels, **legend_kwargs)
-    return
+    return labels
 
 
 def plot_se2_pose(

examples/solo_jump.py

@@ -65,7 +65,7 @@ def test():
 test()
 
 
-dt = 10e-3  # 20 ms
+dt = 5e-3  # 20 ms
 tf = 1.2  # in seconds
 nsteps = int(tf / dt)
 print("Num steps: {:d}".format(nsteps))
@@ -79,7 +79,7 @@ def test():
 mask = (switch_t0 <= times) & (times < switch_t1)
 
 q1 = q0.copy()
-q1[3:7] = pin.exp3_quat(np.array([0.0, 0.0, np.pi / 3]))
+# q1[3:7] = pin.exp3_quat(np.array([0.0, 0.0, np.pi / 3]))
 v0 = np.zeros(nv)
 x0_ref = np.concatenate((q0, v0))
 w_x = np.ones(space.ndx) * 1e-2
@@ -157,7 +157,7 @@ def create_land_cost(costs, w):
 
 problem = aligator.TrajOptProblem(x0_ref, stages, term_cost)
 mu_init = 1e-5
-tol = 1e-4
+tol = 1e-5
 solver = aligator.SolverProxDDP(tol, mu_init, verbose=aligator.VERBOSE, max_iters=300)
 solver.rollout_type = aligator.ROLLOUT_LINEAR
 solver.setNumThreads(args.num_threads)
@@ -201,6 +201,7 @@ def make_plots(res: aligator.Results):
 
     fig4 = plt.figure()
     ax = fig4.add_subplot(111)
+    ax.hlines(tol, 0, res.num_iters, lw=2.2, alpha=0.8, colors="k")
     plot_convergence(cb_, ax, res, show_al_iters=True)
     fig4.tight_layout()