diff --git a/README.md b/README.md
index 7dc325b1..e14c11ad 100755
--- a/README.md
+++ b/README.md
@@ -137,43 +137,63 @@ python tutorials/multi_robot_3d_reachavoid.py
- Ellipsoidal-obstacle CBF
Unicycle reach-goal with linear class-K
+ Ellipsoidal-obstacle CBF ๐
Unicycle reach-goal with linear class-K
|
- Stochastic CBF (SDE)
Safety under Brownian disturbance
+ Stochastic CBF (SDE) ๐
Safety under Brownian disturbance
|
- Robust CBF
Worst-case bounded disturbance
+ Robust CBF ๐
Worst-case bounded disturbance
|
- MPPI rollout sampling
Sampling-based planning
+ MPPI rollout sampling ๐
Sampling-based planning
|
- MPPI reach-avoid
Sampling-based planning with goal + obstacle cost
+ MPPI reach-avoid ๐
Sampling-based planning with goal + obstacle cost
|
- Multi-robot 2D
Coordination via shared CBFs
+ Multi-robot 2D ๐
Coordination via shared CBFs
|
- Fixed-wing aerial 3D
UAV reach-drop-point in 3D
+ Fixed-wing aerial 3D ๐
UAV reach-drop-point in 3D
|
- Pedestrian head-on
Dynamic-agent avoidance
+ Pedestrian head-on ๐
Dynamic-agent avoidance
|
- EKF state estimation
CBF over noisy estimates
+ EKF state estimation ๐
Unicycle reach-goal under measurement noise
+ |
+
+
+
+ 
+ Van der Pol (CLF) ๐
Nonlinear regulation to the origin
+ |
+
+ 
+ Model Predictive Control ๐
Receding-horizon LTI tracking
+ |
+
+ 
+ Quadrotor 6-DOF ๐
Geometric SE(3) tracking + altitude CBF
+ |
+
+
+
+ 
+ Monte Carlo safety verification ๐
200 stochastic CBF rollouts (jax.vmap), live empirical risk
|
diff --git a/examples/unicycle/reach_goal/ekf.py b/examples/unicycle/reach_goal/ekf.py
index 381f0e96..96564da0 100755
--- a/examples/unicycle/reach_goal/ekf.py
+++ b/examples/unicycle/reach_goal/ekf.py
@@ -104,7 +104,8 @@ def dhdx(x):
x_lim=(-5, 5),
y_lim=(-5, 5),
dt=dt,
- title="System Behavior",
+ title="EKF State Estimation",
save_animation=save,
animation_filename="examples/unicycle/reach_goal/results/ekf_estimation.gif",
+ backend="matplotlib",
)
diff --git a/examples/unicycle/reach_goal/mppi_cbf.py b/examples/unicycle/reach_goal/mppi_cbf.py
index 85856588..3ee68b50 100644
--- a/examples/unicycle/reach_goal/mppi_cbf.py
+++ b/examples/unicycle/reach_goal/mppi_cbf.py
@@ -219,8 +219,8 @@ def terminal_cost(state: Array, action: Array) -> Array:
title="System Behavior (MPPI + CBF)",
obstacles=obstacles,
ellipsoids=ellipsoids,
- save_animation=False,
- animation_filename=file_path + "bh_mppi_cbf_control",
+ save_animation=save,
+ animation_filename=file_path + "mppi_rollouts.gif",
)
final_pos = x[-1, :2]
diff --git a/media/showcase/ekf_estimation.gif b/media/showcase/ekf_estimation.gif
index be08a74a..9818451d 100644
Binary files a/media/showcase/ekf_estimation.gif and b/media/showcase/ekf_estimation.gif differ
diff --git a/media/showcase/monte_carlo_safety.gif b/media/showcase/monte_carlo_safety.gif
new file mode 100644
index 00000000..af7a34f7
Binary files /dev/null and b/media/showcase/monte_carlo_safety.gif differ
diff --git a/media/showcase/mpc_double_integrator.gif b/media/showcase/mpc_double_integrator.gif
new file mode 100644
index 00000000..5e9f6768
Binary files /dev/null and b/media/showcase/mpc_double_integrator.gif differ
diff --git a/media/showcase/mppi_rollouts.gif b/media/showcase/mppi_rollouts.gif
index 0ba2fcf0..4e395497 100644
Binary files a/media/showcase/mppi_rollouts.gif and b/media/showcase/mppi_rollouts.gif differ
diff --git a/media/showcase/pedestrian_head_on.gif b/media/showcase/pedestrian_head_on.gif
index cfb47c29..a8592bbb 100644
Binary files a/media/showcase/pedestrian_head_on.gif and b/media/showcase/pedestrian_head_on.gif differ
diff --git a/media/showcase/quadrotor_6dof.gif b/media/showcase/quadrotor_6dof.gif
new file mode 100644
index 00000000..bcea4043
Binary files /dev/null and b/media/showcase/quadrotor_6dof.gif differ
diff --git a/media/showcase/van_der_pol_clf.gif b/media/showcase/van_der_pol_clf.gif
new file mode 100644
index 00000000..f8740511
Binary files /dev/null and b/media/showcase/van_der_pol_clf.gif differ
diff --git a/scripts/render_showcase.py b/scripts/render_showcase.py
index eec41655..989047e8 100644
--- a/scripts/render_showcase.py
+++ b/scripts/render_showcase.py
@@ -1023,5 +1023,507 @@ def update(i):
return str(out)
+@register("van_der_pol_clf")
+def render_van_der_pol_clf() -> str:
+ """Van der Pol: Lyapunov-based regulation of a nonlinear oscillator to the origin."""
+ import jax.numpy as jnp
+ import matplotlib.pyplot as plt
+ import numpy as np
+ from jax import jit
+ from matplotlib.animation import FuncAnimation, PillowWriter
+
+ import cbfkit.simulation.simulator as sim
+ from cbfkit.estimators import naive as estimator
+ from cbfkit.integration import runge_kutta_4 as integrator
+ from cbfkit.sensors import perfect as sensor
+ from cbfkit.systems import van_der_pol
+ from cbfkit.utils.user_types import ControllerData
+
+ epsilon = 0.2
+ dyn = van_der_pol.reverse_van_der_pol_oscillator(epsilon=epsilon)
+
+ # Lyapunov-based regulation law. The plant's input matrix g = [0, 1/x2] is singular,
+ # so the control is formed as u = x2 * (...) to cancel the 1/x2 amplification โ which is
+ # exactly why the packaged closed-form FxTS law cannot be dropped onto this model directly.
+ def regulation_controller(eps, k1=4.0, k2=4.0):
+ @jit
+ def controller(_t, x, _key, _xd=None):
+ x1, x2 = x
+ u = x2 * ((k1 - 1.0) * x1 - k2 * x2 + eps * (1.0 - x1**2) * x2)
+ return jnp.array([u]), ControllerData()
+
+ return controller
+
+ x0 = jnp.array([2.0, 2.0])
+ dt, tf = 1e-3, 5.0
+ n = int(tf / dt)
+ res = sim.execute(
+ x0=x0,
+ dt=dt,
+ num_steps=n,
+ dynamics=dyn,
+ integrator=integrator,
+ nominal_controller=regulation_controller(epsilon),
+ sensor=sensor,
+ estimator=estimator,
+ use_jit=True,
+ )
+ states = np.asarray(res["states"])
+
+ fig, ax = plt.subplots(figsize=(6, 6))
+ # Open-loop Van der Pol vector field: shows the nonlinearity the Lyapunov law tames.
+ gx = np.linspace(-3.0, 3.0, 22)
+ GX, GY = np.meshgrid(gx, gx)
+ FX = -GY
+ FY = GX - epsilon * (1.0 - GX**2) * GY
+ mag = np.hypot(FX, FY) + 1e-9
+ ax.quiver(GX, GY, FX / mag, FY / mag, color="gray", alpha=0.35, width=0.003)
+ ax.add_patch(plt.Circle((0, 0), 0.1, color="green", alpha=0.3))
+ ax.plot(0, 0, "g*", markersize=18, label="Origin (goal)")
+ (line,) = ax.plot([], [], "b-", lw=2)
+ dot = ax.scatter([], [], s=80, color="blue", zorder=5)
+ ax.set_xlim(-3, 3)
+ ax.set_ylim(-3, 3)
+ ax.set_aspect("equal")
+ ax.set_xlabel("$x_1$")
+ ax.set_ylabel("$x_2$")
+ ax.set_title("Van der Pol โ Lyapunov regulation to the origin", fontsize=10)
+ ax.legend(loc="upper right", fontsize=9)
+ ax.grid(True, alpha=0.3)
+
+ def update(i):
+ line.set_data(states[: i + 1, 0], states[: i + 1, 1])
+ dot.set_offsets([[states[i, 0], states[i, 1]]])
+ return line, dot
+
+ stride = max(1, len(states) // 70)
+ anim = FuncAnimation(fig, update, frames=range(0, len(states), stride), interval=100, blit=True)
+ out = OUT / "van_der_pol_clf.gif"
+ anim.save(out, writer=PillowWriter(fps=10))
+ plt.close(fig)
+ return str(out)
+
+
+@register("mpc_double_integrator")
+def render_mpc_double_integrator() -> str:
+ """Classical receding-horizon MPC: LTI double-integrator tracking to a goal.
+
+ Honest framing: this solver carries only equality (dynamics) constraints, so it is
+ reference tracking, not a safety filter. Driven as a standalone receding-horizon loop.
+ """
+ import jax.numpy as jnp
+ import matplotlib.pyplot as plt
+ import numpy as np
+ from matplotlib.animation import FuncAnimation, PillowWriter
+
+ from cbfkit.optimization.mpc.quadratic_cost_linear_dynamics import (
+ generate_mpc_solver_quadratic_cost_linear_dynamics,
+ )
+
+ dt = 0.1
+ # Discrete-time double integrator: state [px, py, vx, vy], control [ax, ay].
+ A = jnp.array(
+ [[1.0, 0.0, dt, 0.0], [0.0, 1.0, 0.0, dt], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+ )
+ B = jnp.array([[0.0, 0.0], [0.0, 0.0], [dt, 0.0], [0.0, dt]])
+ Q = jnp.diag(jnp.array([10.0, 10.0, 1.0, 1.0]))
+ R = 0.1 * jnp.eye(2)
+ Qn = 50.0 * Q
+ N = 20
+ solve = generate_mpc_solver_quadratic_cost_linear_dynamics(A, B, Q, R, Qn, N)
+
+ goal = jnp.array([4.0, 4.0, 0.0, 0.0])
+ ref_horizon = jnp.tile(goal, (N, 1)) # (N, 4) constant reference over the horizon
+ x = jnp.array([0.0, 0.0, 0.0, 0.0])
+ n_steps = 40
+
+ xs = [np.asarray(x)]
+ preds = []
+ for _ in range(n_steps):
+ concatenated_x_xr = jnp.vstack([x.reshape(1, -1), ref_horizon]) # (N+1, 4)
+ x_opt, u_opt = solve(concatenated_x_xr) # x_opt (4, N+1), u_opt (2, N)
+ u = u_opt[:, 0]
+ x = A @ x + B @ u
+ xs.append(np.asarray(x))
+ preds.append(np.asarray(x_opt.T)) # (N+1, 4) predicted state horizon
+ xs = np.array(xs)
+
+ fig, ax = plt.subplots(figsize=(6, 6))
+ ax.plot(float(goal[0]), float(goal[1]), "g*", markersize=18, label="Goal")
+ (realized,) = ax.plot([], [], "b-", lw=2, label="Realized")
+ (pred,) = ax.plot(
+ [], [], color="orange", ls="--", lw=1.5, alpha=0.85, label="Predicted horizon"
+ )
+ dot = ax.scatter([], [], s=80, color="blue", zorder=5)
+ ax.set_xlim(-0.5, 4.5)
+ ax.set_ylim(-0.5, 4.5)
+ ax.set_aspect("equal")
+ ax.set_xlabel("x")
+ ax.set_ylabel("y")
+ ax.set_title("Model Predictive Control โ receding-horizon LTI tracking", fontsize=10)
+ ax.legend(loc="lower right", fontsize=9)
+ ax.grid(True, alpha=0.3)
+
+ def update(i):
+ realized.set_data(xs[: i + 1, 0], xs[: i + 1, 1])
+ dot.set_offsets([[xs[i, 0], xs[i, 1]]])
+ p = preds[min(i, len(preds) - 1)]
+ pred.set_data(p[:, 0], p[:, 1])
+ return realized, pred, dot
+
+ anim = FuncAnimation(fig, update, frames=len(xs), interval=100, blit=True)
+ out = OUT / "mpc_double_integrator.gif"
+ anim.save(out, writer=PillowWriter(fps=10))
+ plt.close(fig)
+ return str(out)
+
+
+@register("quadrotor_6dof")
+def render_quadrotor_6dof() -> str:
+ """6-DOF quadrotor: geometric SE(3) tracking + live CBF altitude-envelope value.
+
+ Honest framing: we drive the quadrotor through 3D space with the geometric
+ controller (Lee-Leok-McClamroch) and display the altitude-CBF barrier value
+ h(z) alongside, demonstrating the available CBF certificate without claiming
+ an active barrier-projection filter (which the geometric controller doesn't
+ natively expose).
+ """
+ import jax.numpy as jnp
+ import matplotlib.pyplot as plt
+ import numpy as np
+ from matplotlib.animation import FuncAnimation, PillowWriter
+ from mpl_toolkits.mplot3d import Axes3D # noqa: F401 โ registers 3D projection
+
+ import cbfkit.simulation.simulator as sim
+ from cbfkit.estimators import naive as estimator
+ from cbfkit.integration import runge_kutta_4 as integrator
+ from cbfkit.sensors import perfect as sensor
+ from cbfkit.systems.quadrotor_6dof.certificates.barrier_functions import h_alt
+ from cbfkit.systems.quadrotor_6dof.controllers.geometric import geometric_controller
+ from cbfkit.systems.quadrotor_6dof.models.quadrotor_6dof_dynamics import (
+ quadrotor_6dof_dynamics,
+ )
+
+ # Mass/inertia must be consistent between plant and controller: geometric_controller's
+ # default gains are tuned for mโ4.34 kg, while quadrotor_6dof_dynamics defaults to
+ # m=0.25 kg. Mismatch -> instant integration NaN. Use the heavier plant.
+ m, jx, jy, jz = 4.34, 0.0820, 0.0845, 0.1377
+ three_tuple = quadrotor_6dof_dynamics(m=m, jx=jx, jy=jy, jz=jz)
+
+ def dyn(x):
+ f, g, _s = three_tuple(x)
+ return f, g
+
+ desired = jnp.array([2.0, 1.5, 3.0]) # target (pn, pe, h)
+ dt = 0.01
+ tf = 6.0
+ n = int(tf / dt)
+
+ # state layout: [pn, pe, h, u, v, w, phi, theta, psi, p, q, r]
+ x0 = jnp.array([0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+
+ nominal = geometric_controller(
+ dynamics=dyn, desired_state=desired, dt=dt, m=m, jx=jx, jy=jy, jz=jz
+ )
+
+ res = sim.execute(
+ x0=x0,
+ dt=dt,
+ num_steps=n,
+ dynamics=dyn,
+ integrator=integrator,
+ nominal_controller=nominal,
+ sensor=sensor,
+ estimator=estimator,
+ use_jit=True,
+ )
+ states = np.asarray(res["states"]) # (n+1, 12)
+
+ # Altitude-CBF barrier value h_alt(z, alt_limit). z = hstack([x, t]).
+ # alt_limit must comfortably exceed our setpoint altitude (3 m) โ pick 5 m.
+ alt_limit = 5.0
+ n_states_full = states.shape[0]
+ ts = np.linspace(0.0, tf, n_states_full)
+ h_vals = np.array(
+ [
+ float(h_alt(jnp.hstack([jnp.asarray(states[i]), jnp.asarray(ts[i])]), alt_limit))
+ for i in range(n_states_full)
+ ]
+ )
+
+ # Subsample frames for a compact GIF.
+ stride = max(1, n_states_full // 80)
+ idx = np.arange(0, n_states_full, stride)
+ pn, pe, h_alt_traj = states[idx, 0], states[idx, 1], states[idx, 2]
+
+ fig = plt.figure(figsize=(10, 5))
+ ax3d = fig.add_subplot(1, 2, 1, projection="3d")
+ ax_h = fig.add_subplot(1, 2, 2)
+
+ ax3d.scatter(
+ [float(desired[0])],
+ [float(desired[1])],
+ [float(desired[2])],
+ color="green",
+ s=120,
+ marker="*",
+ label="Goal",
+ zorder=10,
+ )
+ (line3d,) = ax3d.plot([], [], [], "b-", lw=2, label="Quadrotor")
+ dot3d = ax3d.scatter([], [], [], s=60, color="blue", zorder=11)
+ pad = 0.5
+ ax3d.set_xlim(min(pn.min(), float(desired[0])) - pad, max(pn.max(), float(desired[0])) + pad)
+ ax3d.set_ylim(min(pe.min(), float(desired[1])) - pad, max(pe.max(), float(desired[1])) + pad)
+ ax3d.set_zlim(0, alt_limit + 0.5)
+ ax3d.set_xlabel("pn [m]")
+ ax3d.set_ylabel("pe [m]")
+ ax3d.set_zlabel("h [m]")
+ ax3d.set_title("Quadrotor 6-DOF โ geometric SE(3) tracking", fontsize=10)
+ ax3d.legend(loc="upper right", fontsize=8)
+ ax3d.view_init(elev=22, azim=-60)
+
+ # h(z) trace: stays >0 โ altitude envelope satisfied.
+ ax_h.plot(ts, h_vals, color="purple", lw=1.5)
+ (h_dot,) = ax_h.plot([], [], "o", color="purple", markersize=7)
+ ax_h.axhline(0.0, color="red", ls="--", lw=1, alpha=0.7, label="Safety boundary h=0")
+ ax_h.set_xlim(0, tf)
+ ax_h.set_ylim(min(0.0, float(h_vals.min())) - 0.1, max(1.0, float(h_vals.max())) + 0.1)
+ ax_h.set_xlabel("t [s]")
+ ax_h.set_ylabel("$h_{\\rm alt}(z)$")
+ ax_h.set_title("Altitude-CBF barrier value (positive โ safe)", fontsize=10)
+ ax_h.legend(loc="lower right", fontsize=8)
+ ax_h.grid(True, alpha=0.3)
+
+ def update(i):
+ line3d.set_data(pn[: i + 1], pe[: i + 1])
+ line3d.set_3d_properties(h_alt_traj[: i + 1])
+ dot3d._offsets3d = ([pn[i]], [pe[i]], [h_alt_traj[i]])
+ # Map subsampled index back to full-resolution h_vals index for the dot.
+ full_i = idx[i]
+ h_dot.set_data([ts[full_i]], [h_vals[full_i]])
+ return line3d, dot3d, h_dot
+
+ anim = FuncAnimation(fig, update, frames=len(idx), interval=100, blit=False)
+ out = OUT / "quadrotor_6dof.gif"
+ anim.save(out, writer=PillowWriter(fps=10))
+ plt.close(fig)
+ return str(out)
+
+
+@register("monte_carlo_safety")
+def render_monte_carlo_safety() -> str:
+ """GPU/vmap Monte Carlo safety funnel: N stochastic single-integrator rollouts kept
+ safe around an obstacle by a CBF-QP filter, with a live empirical violation-rate counter.
+
+ Each of the N trials gets its own initial state (Gaussian funnel-mouth) and its own
+ Brownian process noise (Euler-Maruyama), all executed as one ``jax.vmap`` kernel via
+ ``conduct_monte_carlo_gpu``. The empirical risk = fraction of trials that have entered
+ the obstacle by the current frame; the CBF holds it at ~0.
+ """
+ import contextlib
+ import os
+
+ import jax.numpy as jnp
+ import matplotlib.pyplot as plt
+ import numpy as np
+ from jax import random
+ from matplotlib.animation import FuncAnimation, PillowWriter
+ from matplotlib.collections import LineCollection
+
+ from cbfkit.controllers.cbf_clf.cbf_clf_qp_generator import cbf_clf_qp_generator
+ from cbfkit.controllers.cbf_clf.generate_constraints import (
+ generate_compute_vanilla_clf_constraints,
+ generate_compute_zeroing_cbf_constraints,
+ )
+ from cbfkit.integration import forward_euler
+ from cbfkit.modeling.additive_disturbances import generate_stochastic_perturbation
+ from cbfkit.simulation.monte_carlo_gpu import MonteCarloSetup, conduct_monte_carlo_gpu
+ from cbfkit.utils.user_types import CertificateCollection, ControllerData, PlannerData
+
+ # --- Scenario (verified-clean: low alpha keeps the jaxopt QP stable under vmap) ---
+ GOAL = jnp.array([4.0, 4.0])
+ OBS = jnp.array([2.0, 2.0])
+ R = 0.6
+ ALPHA = 1.0
+ NOISE = 0.4
+ DT, NSTEPS, N_TRIALS = 0.05, 100, 200
+
+ def dynamics(x):
+ return jnp.zeros(2), jnp.eye(2)
+
+ # h(x) = ||x - c||^2 - r^2, relative-degree-1 zeroing barrier (single integrator).
+ f_h = lambda _t, x: jnp.sum((x - OBS) ** 2) - R**2 # noqa: E731
+ j_h = lambda _t, x: 2.0 * (x - OBS) # noqa: E731
+ h_h = lambda _t, _x: 2.0 * jnp.eye(2) # noqa: E731
+ p_h = lambda _t, _x: 0.0 # noqa: E731
+ a_h = lambda h: ALPHA * h # noqa: E731
+ barriers = CertificateCollection([f_h], [j_h], [h_h], [p_h], [a_h])
+
+ controller = cbf_clf_qp_generator(
+ generate_compute_zeroing_cbf_constraints,
+ generate_compute_vanilla_clf_constraints,
+ )(
+ control_limits=jnp.array([8.0, 8.0]),
+ dynamics_func=dynamics,
+ barriers=barriers,
+ relaxable_cbf=False,
+ relaxable_clf=True,
+ )
+
+ def nominal_controller(t, x, _key, _ref):
+ return 2.0 * (GOAL - x), None
+
+ def initial_state_sampler(key):
+ return jnp.array([0.0, 0.0]) + 0.18 * random.normal(key, (2,))
+
+ def _sensor(t, x, *, sigma=None, key=None):
+ return x
+
+ def _estimator(t, y, z, u, c):
+ return y, (c if c is not None else jnp.zeros((len(y), len(y))))
+
+ # Pass the perturbation UNWRAPPED so its `.is_increment` flag survives (Euler-Maruyama).
+ perturbation = generate_stochastic_perturbation(sigma=lambda x: NOISE * jnp.eye(2), dt=DT)
+
+ _, c_data = controller(0.0, jnp.zeros(2), jnp.zeros(2), random.PRNGKey(0), ControllerData())
+ setup = MonteCarloSetup(
+ dt=DT,
+ num_steps=NSTEPS,
+ dynamics=dynamics,
+ integrator=forward_euler,
+ initial_state_sampler=initial_state_sampler,
+ nominal_controller=nominal_controller,
+ controller=controller,
+ sensor=_sensor,
+ estimator=_estimator,
+ perturbation=perturbation,
+ sigma=jnp.zeros(0),
+ controller_data=c_data,
+ planner=None,
+ planner_data=PlannerData(),
+ )
+
+ # The CBF-QP controller emits batched jax.debug.print spam under vmap (every branch of
+ # its status lax.switch fires); silence it at the fd level around the kernel run.
+ @contextlib.contextmanager
+ def _silence_fds():
+ saved = os.dup(1), os.dup(2)
+ devnull = os.open(os.devnull, os.O_WRONLY)
+ os.dup2(devnull, 1)
+ os.dup2(devnull, 2)
+ try:
+ yield
+ finally:
+ os.dup2(saved[0], 1)
+ os.dup2(saved[1], 2)
+ os.close(devnull)
+ os.close(saved[0])
+ os.close(saved[1])
+
+ print(f"[monte_carlo_safety] running {N_TRIALS} vmap'd stochastic rollouts...")
+ with _silence_fds():
+ results = conduct_monte_carlo_gpu(setup, n_trials=N_TRIALS, seed=0)
+ states = np.asarray(results.states) # (N_TRIALS, NSTEPS, 2)
+ print(f"[monte_carlo_safety] kernel wall time: {results.wall_time_s:.2f}s")
+
+ # Geometric safety check (independent of the controller's internal barrier bookkeeping).
+ dist = np.linalg.norm(states - np.asarray(OBS), axis=-1) # (N, NSTEPS)
+ inside = dist < R # (N, NSTEPS)
+ ever_inside = inside.any(axis=1) # (N,)
+ # Cumulative empirical violation rate up to each step.
+ cum_viol_rate = np.array([float(inside[:, : k + 1].any(axis=1).mean()) for k in range(NSTEPS)])
+ overall_rate = float(ever_inside.mean())
+ print(
+ f"[monte_carlo_safety] overall empirical violation rate: {overall_rate:.3f} "
+ f"(min dist to obstacle center {dist.min():.3f}, R={R})"
+ )
+
+ # Draw a representative subset to keep the GIF small (the full 200-line translucent tangle
+ # bloats the palette), but ALWAYS include every breaching trial so the red paths shown stay
+ # consistent with the empirical-risk counter, which is computed over ALL N_TRIALS.
+ from matplotlib.lines import Line2D
+
+ N_DRAW = 60
+ rng = np.random.default_rng(0)
+ viol_idx = np.flatnonzero(ever_inside)
+ safe_idx = np.flatnonzero(~ever_inside)
+ n_safe_draw = min(len(safe_idx), max(0, N_DRAW - len(viol_idx)))
+ safe_draw = rng.choice(safe_idx, size=n_safe_draw, replace=False)
+ draw_idx = np.concatenate([safe_draw, viol_idx]).astype(int)
+ draw_states = states[draw_idx] # (N_DRAW, NSTEPS, 2)
+ draw_colors = ["tab:red" if ever_inside[i] else "tab:blue" for i in draw_idx]
+
+ fig, ax = plt.subplots(figsize=(5.0, 5.0))
+ ax.add_patch(plt.Circle((float(OBS[0]), float(OBS[1])), R, color="red", alpha=0.3, zorder=1))
+ ax.add_patch(
+ plt.Circle((float(OBS[0]), float(OBS[1])), R, fill=False, color="red", lw=1.5, zorder=2)
+ )
+ ax.plot(float(GOAL[0]), float(GOAL[1]), "g*", markersize=18, zorder=6)
+ ax.plot(0.0, 0.0, "ks", markersize=6, zorder=6)
+
+ lc = LineCollection([], colors=draw_colors, linewidths=0.5, alpha=0.3, zorder=3)
+ ax.add_collection(lc)
+ dots = ax.scatter(
+ draw_states[:, 0, 0], draw_states[:, 0, 1], s=6, c=draw_colors, alpha=0.7, zorder=4
+ )
+ txt = ax.text(
+ 0.03,
+ 0.97,
+ "",
+ transform=ax.transAxes,
+ va="top",
+ ha="left",
+ fontsize=9,
+ family="monospace",
+ bbox=dict(boxstyle="round", facecolor="white", alpha=0.85),
+ )
+
+ legend_handles = [
+ Line2D([0], [0], color="tab:blue", lw=1.5, label="safe rollout"),
+ Line2D(
+ [0], [0], marker="*", color="w", markerfacecolor="g", markersize=12, lw=0, label="Goal"
+ ),
+ Line2D(
+ [0], [0], marker="s", color="w", markerfacecolor="k", markersize=7, lw=0, label="Start"
+ ),
+ ]
+ if len(viol_idx) > 0:
+ legend_handles.insert(1, Line2D([0], [0], color="tab:red", lw=1.5, label="breached"))
+
+ ax.set_xlim(-1.0, 5.0)
+ ax.set_ylim(-1.0, 5.0)
+ ax.set_aspect("equal")
+ ax.set_xlabel("x")
+ ax.set_ylabel("y")
+ ax.set_title(
+ f"Monte Carlo safety verification โ {N_TRIALS} stochastic CBF rollouts", fontsize=9
+ )
+ ax.legend(handles=legend_handles, loc="lower right", fontsize=8)
+ ax.grid(True, alpha=0.3)
+
+ stride = max(1, NSTEPS // 30)
+ frame_idx = list(range(0, NSTEPS, stride))
+
+ def update(k):
+ lc.set_segments([draw_states[i, : k + 1, :] for i in range(len(draw_idx))])
+ dots.set_offsets(draw_states[:, k, :])
+ rate = cum_viol_rate[k]
+ n_viol = int(round(rate * N_TRIALS))
+ txt.set_text(
+ f"step {k + 1:3d}/{NSTEPS}\n"
+ f"trials {N_TRIALS}\n"
+ f"violations {n_viol}\n"
+ f"empirical risk {rate * 100:4.1f}%"
+ )
+ return lc, dots, txt
+
+ anim = FuncAnimation(fig, update, frames=frame_idx, interval=100, blit=False)
+ out = OUT / "monte_carlo_safety.gif"
+ anim.save(out, writer=PillowWriter(fps=10), dpi=80)
+ plt.close(fig)
+ return str(out)
+
+
if __name__ == "__main__":
sys.exit(main())
diff --git a/src/cbfkit/optimization/mpc/quadratic_cost_linear_dynamics.py b/src/cbfkit/optimization/mpc/quadratic_cost_linear_dynamics.py
index 167faa78..8e16764c 100644
--- a/src/cbfkit/optimization/mpc/quadratic_cost_linear_dynamics.py
+++ b/src/cbfkit/optimization/mpc/quadratic_cost_linear_dynamics.py
@@ -116,8 +116,13 @@ def mpc_to_qp(
p_bar = jnp.hstack(
[
- -Q1a @ concatenated_x_xr[1:, :n_states].T.flatten(),
- -Q2a @ concatenated_x_xr[-1, :n_states].T.flatten(),
+ # Reference must be time-major ([state@t0, state@t1, ...]) to match the
+ # kron(I_N, Q) cost block and the time-major decision-vector layout.
+ # A prior `.T` here made it state-major, scrambling per-dimension weights.
+ # Factor of 2: the solver minimizes xแตHx + fแตx (no ยฝ), so tracking cost
+ # (xโr)แตQ(xโr) โ f = โ2Qr; without it the optimum collapses to r/2.
+ -2.0 * Q1a @ concatenated_x_xr[1:, :n_states].flatten(),
+ -2.0 * Q2a @ concatenated_x_xr[-1, :n_states].flatten(),
jnp.zeros(horizon * n_inputs),
]
)
diff --git a/src/cbfkit/systems/quadrotor_6dof/controllers/geometric.py b/src/cbfkit/systems/quadrotor_6dof/controllers/geometric.py
index 1e27b340..0e8c750e 100644
--- a/src/cbfkit/systems/quadrotor_6dof/controllers/geometric.py
+++ b/src/cbfkit/systems/quadrotor_6dof/controllers/geometric.py
@@ -4,7 +4,7 @@
from jax import Array, jit
from cbfkit.utils.lqr import compute_lqr_gain
-from cbfkit.utils.matrix_vector_operations import hat, normalize, vee
+from cbfkit.utils.matrix_vector_operations import normalize, vee
from cbfkit.utils.user_types import (
ControllerCallable,
ControllerCallableReturns,
@@ -12,47 +12,20 @@
DynamicsCallable,
)
-from ..certificates.lyapunov_functions import V_pv as V
from ..models.quadrotor_6dof_dynamics import g_accel as g
from ..utils.rotations import rotation_body_frame_to_inertial_frame
-def _extract_drift(dynamics_output: Tuple[Array, ...]) -> Array:
- """Extract drift vector from dynamics outputs supporting 2- or 3-tuples.
-
- Some dynamics call sites return ``(f, g)`` while older/stochastic variants
- return ``(f, g, s)``. Geometric control only uses the drift component.
- """
- if len(dynamics_output) == 2:
- f_val, _ = dynamics_output
- return f_val
- if len(dynamics_output) == 3:
- f_val, _, _ = dynamics_output
- return f_val
- raise ValueError(
- "Expected dynamics callable to return (f, g) or (f, g, s). "
- f"Received tuple of length {len(dynamics_output)}."
- )
-
-
def geometric_controller(
dynamics: DynamicsCallable,
desired_state: Array,
dt: float,
- # m: float = 0.5,
- # jx: float = 0.25,
- # jy: float = 0.25,
- # jz: float = 0.1,
- # kx: float = 1.0,
- # kv: float = 2.05,
- # kr: float = 0.35,
- # ko: float = 0.15,
m: float = 4.34,
jx: float = 0.0820,
jy: float = 0.0845,
jz: float = 0.1377,
- kx: float = 16.0,
- kv: float = 5.6,
+ kx: float = 8.0,
+ kv: float = 8.0,
kr: float = 8.81,
ko: float = 2.54,
) -> ControllerCallable:
@@ -84,18 +57,16 @@ def geometric_controller(
j_vec = jnp.array([jx, jy, jz])
_b1_d = jnp.array([1.0, 0.0, 0.0])
- # # Flatness-based control -- LQR
- # get_desired_pos_vel_acc = lqr_control(desired_state, dt)
-
- # Flatness-based control -- FxTS
- tg = 10.0
- c1, e1, e2 = 0.5, 0.5, 1.5
- c2 = 1 / ((e2 - 1) * (tg - 1 / (c1 * (1 - e1))))
-
- def fV(x):
- return -c1 * V(x, desired_state) ** e1 - c2 * V(x, desired_state) ** e2
-
- get_desired_pos_vel_acc = lyapunov_control(desired_state, dt, fV)
+ # The packaged body->inertial matrix is orthogonal but *improper* (det = -1): its
+ # third row is the standard ZYX rotation's third row negated, encoding the model's
+ # "body z-down, inertial h-up" convention. Geometric SE(3) tracking requires a
+ # proper rotation, so we left-multiply by S = diag(1, 1, -1) to flip that row back.
+ # The result is a true SO(3) matrix expressed in the z-down ("NED") inertial frame
+ # y = S @ p -- exactly the frame the plant's own velocity/gravity terms live in.
+ # (Left-multiplying S negates the 3rd row and preserves Rdot = R @ hat([p, q, r]);
+ # a right-multiply would silently remap the body rates to [-p, -q, r] and diverge.)
+ S = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]])
+ pos_d = jnp.matmul(S, desired_state[:3])
@jit
def controller(
@@ -116,74 +87,32 @@ def controller(
u (Array): computed control inputs
data (dict): requisite dictionary return
"""
- nonlocal _b1_d, e3, j_vec
-
- _, _, _, _, _, _, _, theta, psi, _, _, _ = x
- # dynamics returns (f, g), we unpack the first element which is f?
- # The original code was: dynamics(x)[0]. This seems to imply dynamics(x)
- # returns a tuple/list and the first element is f.
- # Assuming dynamics(x) -> (f, g) or similar.
- # And unpacking f: _, _, _, _, _, _, phi_dot, theta_dot, psi_dot, _, _, _ = f
- f_val = _extract_drift(dynamics(x))
- _, _, _, _, _, _, phi_dot, theta_dot, psi_dot, _, _, _ = f_val
-
- # Get rotation matrix
- body_to_inertial_rotation = rotation_body_frame_to_inertial_frame(x)
-
- # Compute desired position, velocity, acceleration
- pos_d, vel_d, acc_d = get_desired_pos_vel_acc(t, x)
- vel_d = jnp.zeros((3,))
- acc_d = jnp.zeros((3,))
-
- # Define tracking errors
- e_pos = x[:3] - pos_d
- e_vel = jnp.matmul(body_to_inertial_rotation, x[3:6]) - vel_d
-
- # Compute desired attitude and attitude tracking error
- b3_d = -normalize(-kx * e_pos - kv * e_vel + m * g * e3 + m * acc_d)
+ # Proper SO(3) rotation, body frame -> z-down ("NED") inertial frame.
+ rotation = jnp.matmul(S, rotation_body_frame_to_inertial_frame(x))
+
+ # Translational tracking errors in the z-down inertial frame (vel_d = 0).
+ e_pos = jnp.matmul(S, x[:3]) - pos_d
+ e_vel = jnp.matmul(rotation, x[3:6])
+
+ # Desired thrust direction (Lee et al. 2010, NED form: gravity acts along +e3).
+ thrust_vec = -kx * e_pos - kv * e_vel - m * g * e3
+ b3_d = -normalize(thrust_vec)
b2_d = normalize(jnp.cross(b3_d, _b1_d))
b1_d = normalize(jnp.cross(b2_d, b3_d))
rot_d = jnp.array([b1_d, b2_d, b3_d]).T
- e_rot = (
- 1
- / 2
- * vee(
- jnp.matmul(rot_d.T, body_to_inertial_rotation)
- - jnp.matmul(body_to_inertial_rotation.T, rot_d)
- )
- )
- # Compute angular velocity tracking error
- wx_b = phi_dot * jnp.sin(theta) * jnp.sin(psi) + theta_dot * jnp.cos(psi)
- wy_b = phi_dot * jnp.sin(theta) * jnp.cos(psi) - theta_dot * jnp.sin(psi)
- wz_b = phi_dot * jnp.cos(theta) + psi_dot
- omega = jnp.array([wx_b, wy_b, wz_b])
-
- # Compute rotation tracking error
- omega_d = jnp.zeros((3,))
- omega_d_dot = jnp.zeros((3,))
- e_ome = omega - jnp.matmul(body_to_inertial_rotation.T, jnp.matmul(rot_d, omega_d))
-
- # Compute force input
- f = -jnp.dot(
- -kx * e_pos - kv * e_vel + m * g * e3 + m * acc_d,
- jnp.matmul(body_to_inertial_rotation, e3),
- )
+ # Attitude error on SO(3): the vee map is only valid because `rotation` is proper.
+ e_rot = 1 / 2 * vee(jnp.matmul(rot_d.T, rotation) - jnp.matmul(rotation.T, rot_d))
- # Compute moment inputs
- #! need to double check this
- moments = (
- -kr * e_rot
- - ko * e_ome
- + jnp.cross(omega, j_vec * omega)
- - j_vec
- * (
- jnp.matmul(
- jnp.matmul(hat(omega), body_to_inertial_rotation.T), jnp.matmul(rot_d, omega_d)
- )
- - jnp.matmul(body_to_inertial_rotation.T, jnp.matmul(rot_d, omega_d_dot))
- )
- )
+ # Body angular velocity [p, q, r] is part of the state; omega_d = 0 for a setpoint.
+ omega = x[9:12]
+ e_ome = omega
+
+ # Thrust magnitude: project the desired force onto the body-down axis.
+ f = -jnp.dot(thrust_vec, jnp.matmul(rotation, e3))
+
+ # Moment inputs (omega_d = omega_d_dot = 0 for setpoint regulation).
+ moments = -kr * e_rot - ko * e_ome + jnp.cross(omega, j_vec * omega)
inputs = jnp.hstack([f, moments])
diff --git a/tests/test_optimization/test_mpc_tracking.py b/tests/test_optimization/test_mpc_tracking.py
new file mode 100644
index 00000000..73dfd501
--- /dev/null
+++ b/tests/test_optimization/test_mpc_tracking.py
@@ -0,0 +1,79 @@
+"""Regression tests for ``generate_mpc_solver_quadratic_cost_linear_dynamics``.
+
+History: this solver had two latent cost-formulation bugs (the module had no tests
+or examples exercising it):
+
+1. The reference horizon was flattened state-major (``.T.flatten()``) while the
+ Hessian block ``kron(I_N, Q)`` and the decision-vector unpack are time-major.
+ That scrambled per-state-dimension weights against per-state-dimension references.
+
+2. The linear cost term was ``-Q @ r`` instead of ``-2*Q @ r``. The jaxopt QP solver
+ minimizes ``xแตHx + fแตx`` with no implicit 1/2, so a tracking cost
+ ``(x-r)แตQ(x-r)`` produces ``f = -2Q r``. The missing factor of 2 made the QP
+ optimum collapse to ``r/2``: a goal of (4, 4) parked the system at (2, 2).
+
+These tests pin both corrections by driving a discrete-time double-integrator and
+asserting the realized trajectory actually reaches the goal at the correct location.
+"""
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from cbfkit.optimization.mpc.quadratic_cost_linear_dynamics import (
+ generate_mpc_solver_quadratic_cost_linear_dynamics,
+)
+
+
+def _double_integrator_2d(dt: float = 0.1):
+ A = jnp.array(
+ [[1.0, 0.0, dt, 0.0], [0.0, 1.0, 0.0, dt], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
+ )
+ B = jnp.array([[0.0, 0.0], [0.0, 0.0], [dt, 0.0], [0.0, dt]])
+ return A, B
+
+
+def _run_mpc(goal_xy, n_steps: int = 40, horizon: int = 20, dt: float = 0.1):
+ A, B = _double_integrator_2d(dt)
+ Q = jnp.diag(jnp.array([10.0, 10.0, 1.0, 1.0]))
+ R = 0.1 * jnp.eye(2)
+ Qn = 50.0 * Q
+ solve = generate_mpc_solver_quadratic_cost_linear_dynamics(A, B, Q, R, Qn, horizon)
+
+ goal = jnp.array([goal_xy[0], goal_xy[1], 0.0, 0.0])
+ ref_horizon = jnp.tile(goal, (horizon, 1))
+ x = jnp.zeros(4)
+ for _ in range(n_steps):
+ concatenated_x_xr = jnp.vstack([x.reshape(1, -1), ref_horizon])
+ _x_opt, u_opt = solve(concatenated_x_xr)
+ x = A @ x + B @ u_opt[:, 0]
+ return np.asarray(x)
+
+
+def test_mpc_reaches_goal_no_factor_of_two_collapse():
+ """Goal (4, 4) must produce x_T โ (4, 4) โ pre-fix it parked at (2, 2)."""
+ x_final = _run_mpc((4.0, 4.0))
+ assert np.allclose(x_final[:2], [4.0, 4.0], atol=1e-2), (
+ f"MPC final position {x_final[:2]} did not reach goal (4, 4). "
+ "If it landed near (2, 2) the linear-cost factor-of-2 regressed."
+ )
+ # Terminal velocity should also be ~0 (system is fully regulated).
+ assert np.allclose(
+ x_final[2:], [0.0, 0.0], atol=1e-2
+ ), f"Terminal velocity {x_final[2:]} non-zero; MPC failed to settle."
+
+
+@pytest.mark.parametrize("goal_xy", [(3.0, 1.0), (1.0, 3.0), (-2.0, 2.5)])
+def test_mpc_tracks_asymmetric_goal_no_axis_scrambling(goal_xy):
+ """Asymmetric goals catch state-major vs time-major reference ordering.
+
+ If the cost vector reverts to state-major (the old ``.T.flatten()``),
+ px-weights pair with py-references and the realized x/y will swap or smear.
+ Symmetric goals like (4, 4) cannot detect this; these asymmetric goals can.
+ """
+ x_final = _run_mpc(goal_xy, n_steps=60)
+ assert np.allclose(x_final[:2], goal_xy, atol=2e-2), (
+ f"MPC reached {x_final[:2]} for goal {goal_xy}. "
+ "If x and y are swapped/smeared, the reference time/state ordering regressed."
+ )
diff --git a/tests/test_quadrotor_geometric_tracking.py b/tests/test_quadrotor_geometric_tracking.py
new file mode 100644
index 00000000..1158ea82
--- /dev/null
+++ b/tests/test_quadrotor_geometric_tracking.py
@@ -0,0 +1,159 @@
+"""Closed-loop regression tests for the 6-DOF quadrotor ``geometric_controller``.
+
+History: the controller had two latent bugs that the pre-existing tests
+(``test_quadrotor_geometric_controller.py``) could not catch, because they only
+asserted that a *single* control evaluation was finite and shape ``(4,)`` โ never
+that a closed-loop simulation actually tracks a setpoint:
+
+1. **Improper rotation (det = -1).** Both ``rotation_body_frame_to_inertial_frame``
+ and the plant's ``rotation_body_to_inertial`` return an orthogonal but *improper*
+ matrix (a reflection): the standard ZYX 3rd row is negated to encode the model's
+ "body z-down, inertial h-up" convention. The geometric attitude error
+ ``vee(Rdแต R - Rแต Rd)`` assumes ``R โ SO(3)``; feeding it a reflection loads the
+ lateral attitude error onto the wrong axis, so the east position ``pe`` diverged
+ (to โ -8 m for a +1.5 m setpoint) while ``pn`` and altitude tracked. The fix builds
+ a proper rotation ``S @ R`` (``S = diag(1, 1, -1)``) inside the controller and runs
+ the Lee 2010 law in the z-down frame โ the plant itself is left untouched.
+
+2. **Wrong body angular velocity.** The controller reconstructed ``omega`` from
+ Euler-angle rates with a transform that is wrong even at hover (it placed the roll
+ rate into the yaw slot), leaving roll undamped โ integration NaN within ~1 s. The
+ body angular velocity ``[p, q, r]`` is already the state slice ``x[9:12]``; the fix
+ uses it directly.
+
+These tests pin both corrections: convergence to an asymmetric setpoint (catches the
+lateral-axis sign bug), absence of NaN/divergence, and that a pure roll rate produces
+a roll-damping moment (catches the omega-axis bug).
+"""
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+import cbfkit.simulation.simulator as sim
+from cbfkit.estimators import naive as estimator
+from cbfkit.integration import runge_kutta_4 as integrator
+from cbfkit.sensors import perfect as sensor
+from cbfkit.systems.quadrotor_6dof.controllers.geometric import geometric_controller
+from cbfkit.systems.quadrotor_6dof.models.quadrotor_6dof_dynamics import (
+ quadrotor_6dof_dynamics,
+ rotation_body_to_inertial,
+)
+from cbfkit.systems.quadrotor_6dof.utils.rotations import (
+ rotation_body_frame_to_inertial_frame,
+)
+
+# Mass/inertia consistent between plant and the controller's default gains.
+_M, _JX, _JY, _JZ = 4.34, 0.0820, 0.0845, 0.1377
+
+
+def _plant():
+ three_tuple = quadrotor_6dof_dynamics(m=_M, jx=_JX, jy=_JY, jz=_JZ)
+
+ def dyn(x):
+ f_val, g_mat, _s = three_tuple(x)
+ return f_val, g_mat
+
+ return dyn
+
+
+def _run(desired, x0=None, dt: float = 0.01, tf: float = 6.0, use_jit: bool = False):
+ dyn = _plant()
+ if x0 is None:
+ x0 = jnp.array([0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+ controller = geometric_controller(
+ dynamics=dyn, desired_state=jnp.asarray(desired), dt=dt, m=_M, jx=_JX, jy=_JY, jz=_JZ
+ )
+ res = sim.execute(
+ x0=x0,
+ dt=dt,
+ num_steps=int(tf / dt),
+ dynamics=dyn,
+ integrator=integrator,
+ nominal_controller=controller,
+ sensor=sensor,
+ estimator=estimator,
+ use_jit=use_jit,
+ )
+ return np.asarray(res["states"])
+
+
+def test_geometric_controller_converges_without_nan():
+ """Free-running setpoint regulation must reach the goal โ pre-fix it NaN'd at ~0.95 s."""
+ desired = jnp.array([2.0, 1.5, 3.0])
+ states = _run(desired)
+
+ assert not np.isnan(states).any(), "Integration produced NaN (omega-reconstruction bug?)."
+ final_err = np.linalg.norm(states[-1, :3] - np.asarray(desired))
+ assert (
+ final_err < 0.1
+ ), f"Final position error {final_err:.3f} m too large; controller did not track."
+
+
+def test_no_lateral_divergence_on_asymmetric_setpoint():
+ """An asymmetric (pn != pe) setpoint catches the improper-rotation lateral-axis sign bug.
+
+ Pre-fix, the east position ``pe`` diverged to โ -8 m (wrong sign, growing) for a
+ +1.5 m target while ``pn``/altitude tracked. A symmetric target cannot detect this.
+ """
+ desired = jnp.array([2.0, 1.5, 3.0])
+ states = _run(desired)
+
+ pe = states[:, 1]
+ # pe must stay bounded and finish near +1.5 (not run away to large negative).
+ assert np.abs(pe).max() < 3.0, f"|pe| peaked at {np.abs(pe).max():.2f} m โ lateral divergence."
+ assert (
+ abs(float(states[-1, 1]) - 1.5) < 0.1
+ ), f"Final pe={float(states[-1, 1]):.3f} did not reach +1.5; attitude-error axis sign regressed."
+ # Attitude stays well away from gimbal/flip throughout.
+ assert np.abs(states[:, 6:9]).max() < 1.0, "Attitude excursion too large; controller unstable."
+
+
+def test_jit_and_nonjit_paths_agree():
+ desired = jnp.array([2.0, 1.5, 3.0])
+ s_nojit = _run(desired, use_jit=False)
+ s_jit = _run(desired, use_jit=True)
+ assert np.allclose(s_nojit[-1], s_jit[-1], atol=1e-3), "JIT and non-JIT trajectories diverge."
+
+
+def test_pure_roll_rate_produces_roll_damping_moment():
+ """A pure body roll rate p>0 (at the setpoint, level attitude) must yield a *roll*-damping
+ moment M_x < 0 โ pinning ``omega = x[9:12]``.
+
+ Pre-fix, the Euler-rate reconstruction routed the roll rate into the yaw channel, so
+ M_x โ 0 and M_z absorbed the (mis-attributed) damping. This isolates that bug from the
+ rotation bug by placing the quadrotor exactly at its setpoint with level attitude.
+ """
+ desired = jnp.array([0.0, 0.0, 1.0])
+ # At the setpoint, level, with only a roll rate p = 0.3 rad/s.
+ x = jnp.array([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.3, 0.0, 0.0])
+ controller = geometric_controller(dynamics=_plant(), desired_state=desired, dt=0.01)
+ control, _ = controller(0.0, x)
+
+ m_x, m_y, m_z = float(control[1]), float(control[2]), float(control[3])
+ assert (
+ m_x < -0.1
+ ), f"M_x={m_x:.3f} should be negative (damp roll); omega axis-mapping regressed."
+ # The roll rate must NOT leak into the yaw moment.
+ assert abs(m_z) < 1e-6, f"M_z={m_z:.3e} nonzero for a pure roll rate; omega routed to yaw."
+
+
+def test_packaged_rotation_is_improper_but_controller_fix_is_proper():
+ """Document the root cause and pin the fix's premise.
+
+ The packaged body->inertial matrices are reflections (det = -1); the controller's
+ ``S @ R`` correction (S = diag(1, 1, -1)) restores a proper SO(3) rotation.
+ """
+ S = np.diag([1.0, 1.0, -1.0])
+ for phi, theta, psi in [(0.0, 0.0, 0.0), (0.3, -0.2, 0.5), (-0.4, 0.6, -0.7)]:
+ x = jnp.array([0, 0, 0, 0, 0, 0, phi, theta, psi, 0, 0, 0.0])
+ R = np.asarray(rotation_body_frame_to_inertial_frame(x))
+ R_dyn = np.asarray(rotation_body_to_inertial(phi, theta, psi))
+ assert np.allclose(R, R_dyn, atol=1e-9), "utils and dynamics rotation matrices must match."
+ assert np.isclose(
+ np.linalg.det(R), -1.0, atol=1e-6
+ ), "packaged rotation should be improper."
+ R_proper = S @ R
+ assert np.isclose(np.linalg.det(R_proper), 1.0, atol=1e-6), "S @ R must be proper SO(3)."
+ assert np.allclose(R_proper @ R_proper.T, np.eye(3), atol=1e-6), "S @ R must be orthogonal."