[examples] re-tune some problems, plot convergence

bindings/python/aligator/utils/plotting.py

@@ -94,7 +94,7 @@ def plot_controls_traj(
             ax.set_ylim(*ylim)
         if joint_names is not None:
             joint_name = joint_names[i].lower()
-            ax.set_ylabel(joint_name)
+            ax.set_title(joint_name, fontsize=8)
     if nu > 1:
         fig.supxlabel(xlabel)
         fig.suptitle("Control trajectories")
@@ -113,6 +113,7 @@ def plot_velocity_traj(
     ncols=2,
     vel_limit=None,
     figsize=(6.4, 6.4),
+    xlabel="Time (s)",
 ) -> tuple[plt.Figure, list[plt.Axes]]:
     vs = np.asarray(vs)
     nv = rmodel.nv
@@ -132,7 +133,7 @@ def plot_velocity_traj(
     tf = times[-1]
 
     if axes is None:
-        fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
+        fig, axes = plt.subplots(nrows, ncols, sharex=True, figsize=figsize)
         fig: plt.Figure
     else:
         fig = axes.flat[0].get_figure()
@@ -148,9 +149,9 @@ def plot_velocity_traj(
             ax.hlines(-vel_limit[i], t0, tf, colors="k", linestyles="--")
             ax.hlines(+vel_limit[i], t0, tf, colors="r", linestyles="dashdot")
             ax.set_ylim(*ylim)
-        ax.set_ylabel(joint_name)
+        ax.set_title(joint_name, fontsize=8)
 
-    fig.supxlabel("Time $t$")
+    fig.supxlabel(xlabel)
     fig.suptitle("Velocity trajectories")
     fig.tight_layout()
     return fig, axes

examples/acrobot.py

@@ -99,22 +99,25 @@ class Args(ArgsBase):
     fig1, axes = plot_controls_traj(
         times, res.us, ncols=1, rmodel=rmodel, figsize=(6.4, 3.2)
     )
+    plt.title("Controls (N/m)")
     fig1.tight_layout()
     xs = np.stack(res.xs)
     vs = xs[:, nq:]
 
     theta_s = np.zeros((nsteps + 1, 2))
     theta_s0 = space.difference(space.neutral(), x0)[:2]
     theta_s = theta_s0 + np.cumsum(vs * timestep, axis=0)
-    fig2 = plt.figure(figsize=(6.4, 6.4))
-    plt.subplot(211)
+    fig2 = plt.figure(figsize=(6.4, 3.2))
+    plt.subplot(1, 2, 1)
     plt.plot(times, theta_s, label=("$\\theta_0$", "$\\theta_1$"))
-    plt.title("Joint angles")
+    plt.title("Joint angles (rad)")
+    plt.xlabel("Time (s)")
     plt.legend()
-    plt.subplot(212)
+    plt.subplot(1, 2, 2)
     plt.plot(times, xs[:, nq:], label=("$\\dot\\theta_0$", "$\\dot\\theta_1$"))
     plt.legend()
-    plt.title("Joint velocities")
+    plt.title("Joint velocities (rad/s)")
+    plt.xlabel("Time (s)")
     fig2.tight_layout()
 
     _fig_dict = {"controls": fig1, "velocities": fig2}

examples/cartpole.py

@@ -45,31 +45,27 @@ class Args(ArgsBase):
 
 # running cost regularizes the control input
 rcost = aligator.CostStack(space, nu)
-wu = np.ones(nu) * 1e-2
+wu = np.ones(nu) * 1e-3
+wx_reg = np.eye(ndx) * 1e-3
 rcost.addCost(
     "ureg", aligator.QuadraticControlCost(space, np.zeros(nu), np.diag(wu) * dt)
 )
 rcost.addCost(
-    "xreg",
-    aligator.QuadraticStateCost(space, nu, space.neutral(), 1e-2 * np.eye(ndx) * dt),
+    "xreg", aligator.QuadraticStateCost(space, nu, space.neutral(), wx_reg * dt)
 )
-frame_place_target = pin.SE3.Identity()
-frame_place_target.translation[:] = target_pos
-frame_err = aligator.FramePlacementResidual(
+frame_err = aligator.FrameTranslationResidual(
     ndx,
     nu,
     model,
-    frame_place_target,
+    target_pos,
     frame_id,
 )
-weights_frame_place = np.zeros(6)
-weights_frame_place[:3] = 1.0
-weights_frame_place = np.diag(weights_frame_place)
 term_cost = aligator.CostStack(space, nu)
+term_cost.addCost("xreg", aligator.QuadraticStateCost(space, nu, x0, wx_reg))
 
 # box constraint on control
-u_min = -5.0 * np.ones(nu)
-u_max = +5.0 * np.ones(nu)
+u_min = -2.0 * np.ones(nu)
+u_max = +2.0 * np.ones(nu)
 
 
 def get_box_cstr():
@@ -79,6 +75,7 @@ def get_box_cstr():
 
 nsteps = 500
 Tf = nsteps * dt
+print(f"Tf = {Tf}")
 problem = aligator.TrajOptProblem(x0, nu, space, term_cost)
 
 for i in range(nsteps):
@@ -87,11 +84,13 @@ def get_box_cstr():
         stage.addConstraint(*get_box_cstr())
     problem.addStage(stage)
 
-term_fun = aligator.FrameTranslationResidual(ndx, nu, model, target_pos, frame_id)
+term_fun = frame_err
 
 if args.term_cstr:
     problem.addTerminalConstraint(term_fun, constraints.EqualityConstraintSet())
 else:
+    weights_frame_place = 5 * np.ones(3)
+    weights_frame_place = np.diag(weights_frame_place)
     term_cost.addCost(
         aligator.QuadraticResidualCost(space, frame_err, weights_frame_place)
     )
@@ -139,7 +138,7 @@ def get_box_cstr():
     ax2.set_title("Angle $\\theta(t)$")
     ax2.legend()
 
-    plt.xlabel("Time $t$")
+    plt.xlabel("Time (s)")
 
     gs1 = gs[1].subgridspec(1, 2, width_ratios=[1, 2])
     ax3 = plt.subplot(gs1[0])
@@ -170,28 +169,8 @@ def get_box_cstr():
     fig2 = plt.figure(figsize=(7.2, 4))
     ax: plt.Axes = plt.subplot(111)
     ax.hlines(TOL, 0, res.num_iters, lw=2.2, alpha=0.8, colors="k")
-    plot_convergence(callback, ax, res)
-    prim_tols = np.array(callback.prim_tols)
-    al_iters = np.array(callback.al_index)
-
-    itrange = np.arange(len(al_iters))
-    legends_ = [
-        "$\\epsilon_\\mathrm{tol}$",
-        "Prim. err $p$",
-        "Dual err $d$",
-    ]
-    if len(itrange) > 0:
-        ax.step(itrange, prim_tols, c="green", alpha=0.9, lw=1.1)
-        al_change = al_iters[1:] - al_iters[:-1]
-        al_change_idx = itrange[:-1][al_change > 0]
-        legends_.extend(
-            [
-                "$\\eta_k$",
-                "AL iters",
-            ]
-        )
+    plot_convergence(callback, ax, res, show_al_iters=True)
 
-        ax.vlines(al_change_idx, *ax.get_ylim(), colors="gray", lw=4.0, alpha=0.5)
     ax.legend(
         [
             "$\\epsilon_\\mathrm{tol}$",
@@ -226,7 +205,9 @@ def get_box_cstr():
     qs = [x[:nq] for x in res.xs.tolist()]
     vs = [x[nq:] for x in res.xs.tolist()]
 
-    obj = pin.GeometryObject("objective", 0, hppfcl.Sphere(0.05), frame_place_target)
+    obj = pin.GeometryObject(
+        "objective", 0, hppfcl.Sphere(0.05), pin.SE3(np.eye(3), target_pos)
+    )
     color = [255, 20, 83, 255]
     obj.meshColor[:] = color
     obj.meshColor /= 255

examples/lqr.py

@@ -23,6 +23,12 @@ class Args(tap.Tap):
 args = Args().parse_args()
 np.random.seed(42)
 
+TAG = "LQR"
+if args.bounds:
+    TAG += "_bounded"
+if args.term_cstr:
+    TAG += "_cstr"
+
 nx = 3  # dimension of the state manifold
 nu = 3  # dimension of the input
 space = manifolds.VectorSpace(nx)
@@ -98,6 +104,8 @@ class Args(tap.Tap):
 
 plt.subplot(121)
 fig1: plt.Figure = plt.gcf()
+fig1.set_figwidth(6.4)
+fig1.set_figheight(3.6)
 
 lstyle = {"lw": 0.9, "marker": ".", "markersize": 5}
 trange = np.arange(nsteps + 1)
@@ -144,7 +152,7 @@ class Args(tap.Tap):
 plt.tight_layout()
 
 
-fig2: plt.Figure = plt.figure()
+fig2: plt.Figure = plt.figure(figsize=(6.4, 3.6))
 ax: plt.Axes = fig2.add_subplot()
 niter = res.num_iters
 ax.hlines(
@@ -155,15 +163,22 @@ class Args(tap.Tap):
     linestyles="-",
     linewidth=2.0,
 )
-plot_convergence(his_cb, ax, res)
+plot_convergence(his_cb, ax, res, show_al_iters=True)
 ax.set_title("Convergence (constrained LQR)")
 ax.legend(
     [
         "Tolerance $\\epsilon_\\mathrm{tol}$",
         "Primal error $p$",
         "Dual error $d$",
+        "$\\eta_k$",
+        "AL iters",
     ]
 )
 fig2.tight_layout()
+fig_dicts = {"traj": fig1, "conv": fig2}
+
+for name, _fig in fig_dicts.items():
+    _fig: plt.Figure
+    _fig.savefig(f"assets/{TAG}_{name}.pdf")
 
 plt.show()

examples/quadrotor_obstacles.py

@@ -375,7 +375,7 @@ def plot_costate_value() -> plt.Figure:
             plot_costate_value()
 
         nplot = 3
-        fig: plt.Figure = plt.figure(figsize=(7.2, 4.0))
+        fig: plt.Figure = plt.figure(figsize=(7.2, 3.6))
         ax0: plt.Axes = fig.add_subplot(1, nplot, 1)
         ax0.plot(times[:-1], us_opt)
         ax0.hlines((u_min[0], u_max[0]), *times[[0, -1]], colors="k", alpha=0.3, lw=1.4)

examples/solo_jump.py

@@ -83,23 +83,15 @@ def test():
 w_x[:6] = 0.0
 w_x[nv : nv + 6] = 0.0
 w_x = np.diag(w_x)
-w_u = np.eye(nu) * 1e-4
-
-
-def add_fly_high_cost(costs: aligator.CostStack, slope):
-    fly_high_w = 1.0
-    for fname, fid in FOOT_FRAME_IDS.items():
-        fn = aligator.FlyHighResidual(space.ndx, rmodel, fid, slope, nu)
-        fl_cost = aligator.QuadraticResidualCost(space, fn, np.eye(2) * dt)
-        costs.addCost(fl_cost, fly_high_w / len(FOOT_FRAME_IDS))
+w_u = np.eye(nu) * 1e-1
 
 
 def create_land_fns():
     out = {}
     for fname, fid in FOOT_FRAME_IDS.items():
         p_ref = rdata.oMf[fid].translation
         fn = aligator.FrameTranslationResidual(space.ndx, nu, rmodel, p_ref, fid)
-        out[fid] = fn[2]
+        out[fid] = fn  # [2]
     return out
 
 
@@ -129,14 +121,13 @@ def create_land_cost(costs, w):
     wxlocal_k = w_x * dt
     if mask[k]:
         vf = ode2
-    xref = x0_ref
 
-    xreg_cost = aligator.QuadraticStateCost(space, nu, xref, weights=wxlocal_k)
+    xref = x0_ref.copy()
+    xreg_cost = aligator.QuadraticStateCost(space, nu, x0_ref, weights=wxlocal_k)
     ureg_cost = aligator.QuadraticControlCost(space, nu, weights=w_u * dt)
     cost = aligator.CostStack(space, nu)
     cost.addCost(xreg_cost)
     cost.addCost(ureg_cost)
-    add_fly_high_cost(cost, slope=50)
 
     dyn_model = dynamics.IntegratorSemiImplEuler(vf, dt)
     stm = aligator.StageModel(cost, dyn_model)
@@ -162,13 +153,11 @@ def create_land_cost(costs, w):
 problem = aligator.TrajOptProblem(x0_ref, stages, term_cost)
 mu_init = 1e-5
 tol = 1e-4
-solver = aligator.SolverProxDDP(tol, mu_init, verbose=aligator.VERBOSE, max_iters=100)
+solver = aligator.SolverProxDDP(tol, mu_init, verbose=aligator.VERBOSE, max_iters=300)
 solver.rollout_type = aligator.ROLLOUT_LINEAR
-solver.linear_solver_choice = aligator.LQ_SOLVER_PARALLEL
-# solver.sa_strategy = aligator.SA_FILTER
 solver.setNumThreads(args.num_threads)
 
-cb_ = aligator.HistoryCallback()
+cb_ = aligator.HistoryCallback(solver)
 solver.registerCallback("his", cb_)
 solver.setup(problem)
 
@@ -207,7 +196,7 @@ def make_plots(res: aligator.Results):
 
     fig4 = plt.figure()
     ax = fig4.add_subplot(111)
-    plot_convergence(cb_, ax, res=res)
+    plot_convergence(cb_, ax, res, show_al_iters=True)
     fig4.tight_layout()
 
     _fig_dict = {"controls": fig1, "velocities": fig2, "foot_traj": fig3, "conv": fig4}

examples/ur10_ballistic.py

@@ -9,7 +9,11 @@
 from pathlib import Path
 from typing import Tuple
 from pinocchio.visualize import MeshcatVisualizer
-from aligator.utils.plotting import plot_controls_traj, plot_velocity_traj
+from aligator.utils.plotting import (
+    plot_controls_traj,
+    plot_convergence,
+    plot_velocity_traj,
+)
 from utils import (
     add_namespace_prefix_to_models,
     ArgsBase,
@@ -269,7 +273,7 @@ def create_stage(contact: bool):
 problem.addTerminalConstraint(*create_term_constraint(target_pos))
 problem.addTerminalConstraint(*get_velocity_limit_constraint())
 tol = 1e-4
-mu_init = 1e-1
+mu_init = 2e-1
 solver = aligator.SolverProxDDP(tol, mu_init, max_iters=300, verbose=aligator.VERBOSE)
 # solver.linear_solver_choice = aligator.LQ_SOLVER_PARALLEL
 # solver.rollout_type = aligator.ROLLOUT_LINEAR
@@ -332,23 +336,43 @@ def viz_callback(i: int):
     _joint_names = robot.model.names
     _efflims = robot.model.effortLimit
     _vlims = robot.model.velocityLimit
+    figsize = (6.4, 4.0)
     fig1, _ = plot_controls_traj(
-        times, us, joint_names=_joint_names, effort_limit=_efflims
+        times, us, rmodel=rmodel, effort_limit=_efflims, figsize=figsize
     )
+    fig1.suptitle("Controls (N/m)")
     fig2, _ = plot_velocity_traj(
-        times, vs[:, :-6], rmodel=robot.model, vel_limit=_vlims
+        times, vs[:, :-6], rmodel=robot.model, vel_limit=_vlims, figsize=figsize
     )
 
-    for fig, name in [(fig1, "controls"), (fig2, "velocity")]:
-        PLOTDIR = Path("assets")
-        for ext in [".png", ".pdf"]:
-            figpath: Path = PLOTDIR / f"{EXPERIMENT_NAME}_{name}"
-            fig.savefig(figpath.with_suffix(ext))
+    PLOTDIR = Path("assets")
 
     fig3 = plt.figure()
     ax: plt.Axes = fig3.add_subplot(111)
     ax.plot(dyn_slackn_slacks)
     ax.set_yscale("log")
     ax.set_title("Dynamic slack errors $\\|s\\|_\\infty$")
 
+    fig4 = plt.figure(figsize=(6.4, 3.6))
+    ax = fig4.add_subplot(111)
+    plot_convergence(his_cb, ax, res=solver.results, show_al_iters=True)
+    ax.set_title("Convergence")
+    ax.legend(
+        [
+            "Tolerance $\\epsilon_\\mathrm{tol}$",
+            "Primal error $p$",
+            "Dual error $d$",
+            "$\\eta_k$",
+            "AL iters",
+        ],
+        fontsize=8,
+    )
+    fig4.tight_layout()
+
+    _fig_dict = {"controls": fig1, "velocity": fig2, "conv": fig4}
+    for name, fig in _fig_dict.items():
+        for ext in [".png", ".pdf"]:
+            figpath: Path = PLOTDIR / f"{EXPERIMENT_NAME}_{name}"
+            fig.savefig(figpath.with_suffix(ext))
+
     plt.show()