From 3f67e90ff284e7407eddc60c8907d32c07edb2f4 Mon Sep 17 00:00:00 2001 From: Hongkai Dai Date: Mon, 12 Aug 2024 17:36:13 -0700 Subject: [PATCH] Add a demo on quadrotor. Use polynomial dynamics for quadrotor. The numerics isn't very good. --- compatible_clf_cbf/clf.py | 39 +++- compatible_clf_cbf/clf_cbf.py | 60 ++++-- compatible_clf_cbf/utils.py | 67 +++++-- examples/nonlinear_toy/synthesize_demo.py | 6 +- examples/quadrotor/demo.py | 158 +++++++++++++++ examples/quadrotor/demo_clf.py | 148 ++++++++++++++ examples/quadrotor/plant.py | 188 ++++++++++++++++++ examples/quadrotor2d/demo.py | 18 +- examples/quadrotor2d/demo_taylor.py | 5 +- tests/examples/quadrotor/test_plant.py | 107 ++++++++++ .../{test_plant.py => test_plant_2d.py} | 0 tests/test_utils.py | 8 +- 12 files changed, 752 insertions(+), 52 deletions(-) create mode 100644 examples/quadrotor/demo.py create mode 100644 examples/quadrotor/demo_clf.py create mode 100644 examples/quadrotor/plant.py create mode 100644 tests/examples/quadrotor/test_plant.py rename tests/examples/quadrotor2d/{test_plant.py => test_plant_2d.py} (100%) diff --git a/compatible_clf_cbf/clf.py b/compatible_clf_cbf/clf.py index 2475db5..e535f9d 100644 --- a/compatible_clf_cbf/clf.py +++ b/compatible_clf_cbf/clf.py @@ -3,6 +3,8 @@ """ from dataclasses import dataclass +import os +import pickle from typing import List, Optional, Tuple, Union from typing_extensions import Self @@ -20,6 +22,7 @@ new_sos_polynomial, solve_with_id, ) +import compatible_clf_cbf.utils @dataclass @@ -571,4 +574,38 @@ def find_candidate_regional_lyapunov( prog.AddSosConstraint(derivative_sos_condition) return prog, V - pass + +def save_clf(V: sym.Polynomial, x_set: sym.Variables, kappa: float, pickle_path: str): + """ + Save the CLF to a pickle file. + """ + _, file_extension = os.path.splitext(pickle_path) + assert file_extension in (".pkl", ".pickle"), f"File extension is {file_extension}" + data = {} + data["V"] = compatible_clf_cbf.utils.serialize_polynomial(V, x_set) + data["kappa"] = kappa + + if os.path.exists(pickle_path): + overwrite_cmd = input( + f"File {pickle_path} already exists. Overwrite the file? Press [Y/n]:" + ) + if overwrite_cmd in ("Y", "y"): + save_cmd = True + else: + save_cmd = False + else: + save_cmd = True + + if save_cmd: + with open(pickle_path, "wb") as handle: + pickle.dump(data, handle) + + +def load_clf(pickle_path: str, x_set: sym.Variables) -> dict: + ret = {} + with open(pickle_path, "rb") as handle: + data = pickle.load(handle) + + ret["V"] = compatible_clf_cbf.utils.deserialize_polynomial(data["V"], x_set) + ret["kappa"] = data["kappa"] + return ret diff --git a/compatible_clf_cbf/clf_cbf.py b/compatible_clf_cbf/clf_cbf.py index ca1d25c..e0888d8 100644 --- a/compatible_clf_cbf/clf_cbf.py +++ b/compatible_clf_cbf/clf_cbf.py @@ -688,7 +688,8 @@ def search_clf_cbf_given_lagrangian( compatible_states_options: Optional[CompatibleStatesOptions] = None, solver_id: Optional[solvers.SolverId] = None, solver_options: Optional[solvers.SolverOptions] = None, - backoff_scale: Optional[float] = None, + backoff_rel_scale: Optional[float] = None, + backoff_abs_scale: Optional[float] = None, ) -> Tuple[ Optional[sym.Polynomial], Optional[np.ndarray], @@ -721,7 +722,9 @@ def search_clf_cbf_given_lagrangian( elif compatible_states_options is not None: self._add_compatible_states_options(prog, V, b, compatible_states_options) - result = solve_with_id(prog, solver_id, solver_options, backoff_scale) + result = solve_with_id( + prog, solver_id, solver_options, backoff_rel_scale, backoff_abs_scale + ) if result.is_success(): V_sol = None if V is None else result.GetSolution(V) b_sol = np.array([result.GetSolution(b_i) for b_i in b]) @@ -874,7 +877,7 @@ def bilinear_alternation( inner_ellipsoid_options: Optional[InnerEllipsoidOptions] = None, binary_search_scale_options: Optional[BinarySearchOptions] = None, compatible_states_options: Optional[CompatibleStatesOptions] = None, - backoff_scale: Optional[float] = None, + backoff_scales: Optional[List[compatible_clf_cbf.utils.BackoffScale]] = None, ) -> Tuple[Optional[sym.Polynomial], np.ndarray]: """ Synthesize the compatible CLF and CBF through bilinear alternation. We @@ -915,8 +918,26 @@ def bilinear_alternation( self.unsafe_regions ) + def evaluate_compatible_states(clf_fun, cbf_funs, x_val): + if clf_fun is not None: + V_candidates = clf_fun.EvaluateIndeterminates(self.x, x_val.T) + print(f"V(candidate_compatible_states)={V_candidates}") + b_candidates = [ + b_i.EvaluateIndeterminates( + self.x, + x_val.T, + ) + for b_i in cbf_funs + ] + for i, b_candidates_val in enumerate(b_candidates): + print(f"b[{i}](candidate_compatible_states)={b_candidates_val}") + for iteration in range(max_iter): print(f"iteration {iteration}") + if compatible_states_options is not None: + evaluate_compatible_states( + clf, cbf, compatible_states_options.candidate_compatible_states + ) # Search for the Lagrangians. ( compatible_lagrangians, @@ -996,23 +1017,22 @@ def bilinear_alternation( compatible_states_options=compatible_states_options, solver_id=solver_id, solver_options=solver_options, - backoff_scale=backoff_scale, + backoff_rel_scale=( + None + if backoff_scales is None + else backoff_scales[iteration].rel + ), + backoff_abs_scale=( + None + if backoff_scales is None + else backoff_scales[iteration].abs + ), ) assert cbf is not None - if clf is not None: - V_candidates = clf.EvaluateIndeterminates( - self.x, compatible_states_options.candidate_compatible_states.T - ) - b_candidates = [ - b_i.EvaluateIndeterminates( - self.x, - compatible_states_options.candidate_compatible_states.T, - ) - for b_i in cbf - ] - print(f"V(candidate_compatible_states)={V_candidates}") - for i, b_candidates_val in enumerate(b_candidates): - print(f"b[{i}](candidate_compatible_states)={b_candidates_val}") + if compatible_states_options is not None: + evaluate_compatible_states( + clf, cbf, compatible_states_options.candidate_compatible_states + ) return clf, cbf def check_compatible_at_state( @@ -1559,7 +1579,9 @@ def save_clf_cbf( kappa_b: np.ndarray, pickle_path: str, ): - """ """ + """ + Save the CLF and CBF to a pickle file. + """ _, file_extension = os.path.splitext(pickle_path) assert file_extension in (".pkl", ".pickle"), f"File extension is {file_extension}" data = {} diff --git a/compatible_clf_cbf/utils.py b/compatible_clf_cbf/utils.py index 1c7c422..c40958b 100644 --- a/compatible_clf_cbf/utils.py +++ b/compatible_clf_cbf/utils.py @@ -226,51 +226,76 @@ def to_lower_triangular_columns(mat: np.ndarray) -> np.ndarray: return ret +@dataclasses.dataclass +class BackoffScale: + rel: Optional[float] + abs: Optional[float] + + def solve_with_id( prog: solvers.MathematicalProgram, solver_id: Optional[solvers.SolverId] = None, solver_options: Optional[solvers.SolverOptions] = None, - backoff_scale: Optional[float] = None, + backoff_rel_scale: Optional[float] = None, + backoff_abs_scale: Optional[float] = None, ) -> solvers.MathematicalProgramResult: """ Args: - backoff_scale: when solving an optimization problem with an objective function, - we first solve the problem to optimality, and then "back off" a little bit to find - a sub-optimal but strictly feasible solution. backoff_scale=0 corresponds to no - backoff. Note that during backing off, we will modify the original `prog`. + backoff_rel_scale: when solving an optimization problem with an objective + function, we first solve the problem to optimality, and then "back off" a + little bit to find a sub-optimal but strictly feasible solution. + backoff_rel_scale=0 corresponds to no backoff. Note that during backing + off, we will modify the original `prog`. + backoff_abs_scale: The absolute scale to back off. """ if solver_id is None: result = solvers.Solve(prog, None, solver_options) else: solver = solvers.MakeSolver(solver_id) result = solver.Solve(prog, None, solver_options) - if ( - len(prog.linear_costs()) > 0 or len(prog.quadratic_costs()) > 0 - ) and backoff_scale is not None: + if (len(prog.linear_costs()) > 0 or len(prog.quadratic_costs()) > 0) and ( + backoff_rel_scale is not None or backoff_abs_scale is not None + ): assert ( len(prog.linear_costs()) == 1 ), "TODO(hongkai.dai): support program with multiple LinearCost objects." assert ( len(prog.quadratic_costs()) == 0 ), "TODO(hongkai.dai): we currently only support program with linear costs." - assert backoff_scale >= 0, "backoff_scale should be non-negative." + if backoff_rel_scale is not None: + assert backoff_rel_scale >= 0, "backoff_rel_scale should be non-negative." + if backoff_abs_scale is not None: + assert backoff_abs_scale >= 0, "backoff_abs_scale should be non-negative." + # Cannot handle both backoff_rel_scale and backoff_abs_scale + assert ( + backoff_rel_scale is None or backoff_abs_scale is None + ), "backoff_rel_scale and backoff_abs_scale cannot both be set." + optimal_cost = result.get_optimal_cost() coeff_cost = prog.linear_costs()[0].evaluator().a() var_cost = prog.linear_costs()[0].variables() constant_cost = prog.linear_costs()[0].evaluator().b() - prog.RemoveCost(prog.linear_costs()[0]) - cost_upper_bound = ( - optimal_cost * (1 + backoff_scale) - if optimal_cost > 0 - else optimal_cost * (1 - backoff_scale) - ) - prog.AddLinearConstraint( - coeff_cost, -np.inf, cost_upper_bound - constant_cost, var_cost - ) - if solver_id is None: - result = solvers.Solve(prog, None, solver_options) + if backoff_rel_scale is not None: + cost_upper_bound = ( + optimal_cost * (1 + backoff_rel_scale) + if optimal_cost > 0 + else optimal_cost * (1 - backoff_rel_scale) + ) + elif backoff_abs_scale is not None: + cost_upper_bound = optimal_cost + backoff_abs_scale else: - result = solver.Solve(prog, None, solver_options) + assert Exception("backoff_rel_scale or backoff_abs_scale should be set.") + if (backoff_rel_scale is not None and backoff_rel_scale > 0) or ( + backoff_abs_scale is not None and backoff_abs_scale + ) > 0: + prog.RemoveCost(prog.linear_costs()[0]) + prog.AddLinearConstraint( + coeff_cost, -np.inf, cost_upper_bound - constant_cost, var_cost + ) + if solver_id is None: + result = solvers.Solve(prog, None, solver_options) + else: + result = solver.Solve(prog, None, solver_options) return result diff --git a/examples/nonlinear_toy/synthesize_demo.py b/examples/nonlinear_toy/synthesize_demo.py index af06a5c..a6422ce 100644 --- a/examples/nonlinear_toy/synthesize_demo.py +++ b/examples/nonlinear_toy/synthesize_demo.py @@ -9,6 +9,7 @@ import pydrake.symbolic as sym from compatible_clf_cbf import clf_cbf +import compatible_clf_cbf.utils from examples.nonlinear_toy import toy_system @@ -84,7 +85,10 @@ def main(with_u_bound: bool): binary_search_scale_options=None, compatible_states_options=compatible_states_options, solver_options=solver_options, - backoff_scale=0.02, + backoff_scales=[ + compatible_clf_cbf.utils.BackoffScale(rel=0.02, abs=None) + for _ in range(max_iter) + ], ) diff --git a/examples/quadrotor/demo.py b/examples/quadrotor/demo.py new file mode 100644 index 0000000..9a39e1e --- /dev/null +++ b/examples/quadrotor/demo.py @@ -0,0 +1,158 @@ +""" +Certify the compatible CLF/CBF with input limits. +This uses the polynomial dynamics of the quadrotor (with quaternion as part of +the state), and equality constraint to enforce the unit-length constraint on the +quaternion. +""" + +import os + +import numpy as np +import pydrake.solvers as solvers +import pydrake.symbolic as sym + +import compatible_clf_cbf.clf as clf +from compatible_clf_cbf import clf_cbf +from examples.quadrotor.plant import QuadrotorPolyPlant +import examples.quadrotor.demo_clf +from compatible_clf_cbf.utils import BackoffScale + + +def main(use_y_squared: bool, with_u_bound: bool): + x = sym.MakeVectorContinuousVariable(13, "x") + quadrotor = QuadrotorPolyPlant() + f, g = quadrotor.affine_dynamics(x) + + if with_u_bound: + Au = np.concatenate((np.eye(4), -np.eye(4)), axis=0) + u_bound = quadrotor.m * quadrotor.g + bu = np.concatenate((np.full((4,), u_bound), np.zeros((4,)))) + else: + Au, bu = None, None + + # Ground as the unsafe region. + unsafe_regions = [np.array([sym.Polynomial(x[6] + 0.5)])] + state_eq_constraints = quadrotor.equality_constraint(x) + compatible = clf_cbf.CompatibleClfCbf( + f=f, + g=g, + x=x, + unsafe_regions=unsafe_regions, + Au=Au, + bu=bu, + with_clf=True, + use_y_squared=use_y_squared, + state_eq_constraints=state_eq_constraints, + ) + + load_V_init: bool = True + x_set = sym.Variables(x) + V_degree = 2 + b_degrees = [2] + if load_V_init: + data = clf.load_clf( + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "../../data/quadrotor_V_init.pkl", + ), + x_set, + ) + V_init = 3 * data["V"].RemoveTermsWithSmallCoefficients(1e-6) + else: + V_init = examples.quadrotor.demo_clf.find_trig_regional_clf( + V_degree=2, x=x, save_pickle_filename="quadrotor_V_init1.pkl" + ) + + b_init = np.array([1 - V_init]) + kappa_V = 0 + kappa_b = np.array([kappa_V]) + + compatible_lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees( + lambda_y=[ + clf_cbf.CompatibleLagrangianDegrees.Degree(x=1, y=0 if use_y_squared else 1) + for _ in range(4) + ], + xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree( + x=1, y=0 if use_y_squared else 1 + ), + y=( + None + if use_y_squared + else [ + clf_cbf.CompatibleLagrangianDegrees.Degree(x=4, y=0) + for _ in range(compatible.y.size) + ] + ), + rho_minus_V=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2), + b_plus_eps=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2)], + state_eq_constraints=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2)], + ) + unsafe_regions_lagrangian_degrees = [ + clf_cbf.UnsafeRegionLagrangianDegrees( + cbf=0, unsafe_region=[0], state_eq_constraints=[0] + ) + ] + barrier_eps = np.array([0.000]) + x_equilibrium = np.zeros((13,)) + candidate_compatible_states = np.zeros((6, 13)) + candidate_compatible_states[1, 6] = -0.3 + candidate_compatible_states[2, 5] = -0.3 + candidate_compatible_states[2, 6] = -0.3 + candidate_compatible_states[3, 5] = 0.3 + candidate_compatible_states[3, 6] = -0.3 + candidate_compatible_states[4, 4] = 0.3 + candidate_compatible_states[4, 6] = -0.3 + candidate_compatible_states[5, 4] = -0.3 + candidate_compatible_states[5, 6] = -0.3 + compatible_states_options = clf_cbf.CompatibleStatesOptions( + candidate_compatible_states=candidate_compatible_states, + anchor_states=np.zeros((1, 13)), + b_anchor_bounds=[(np.array([0.0]), np.array([1.0]))], + weight_V=1, + weight_b=np.array([1]), + ) + solver_options = solvers.SolverOptions() + solver_options.SetOption(solvers.CommonSolverOption.kPrintToConsole, True) + + max_iter = 1 + backoff_scales = [ + BackoffScale(rel=None, abs=0.2), + ] + V, b = compatible.bilinear_alternation( + V_init, + b_init, + compatible_lagrangian_degrees, + unsafe_regions_lagrangian_degrees, + kappa_V, + kappa_b, + barrier_eps, + x_equilibrium, + V_degree, + b_degrees, + max_iter=max_iter, + solver_options=solver_options, + lagrangian_coefficient_tol=None, + compatible_states_options=compatible_states_options, + backoff_scales=backoff_scales, + ) + pickle_path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "../../data/quadrotor_clf_cbf.pkl" + ) + clf_cbf.save_clf_cbf(V, b, x_set, kappa_V, kappa_b, pickle_path) + compatible_prog, compatible_lagrangians = ( + compatible.construct_search_compatible_lagrangians( + V, + b, + kappa_V=1e-3, + kappa_b=np.array([1e-3]), + lagrangian_degrees=compatible_lagrangian_degrees, + barrier_eps=barrier_eps, + local_clf=True, + ) + ) + compatible_result = solvers.Solve(compatible_prog, None, solver_options) + assert compatible_result.is_success() + + +if __name__ == "__main__": + main(use_y_squared=True, with_u_bound=False) diff --git a/examples/quadrotor/demo_clf.py b/examples/quadrotor/demo_clf.py new file mode 100644 index 0000000..f7f9a81 --- /dev/null +++ b/examples/quadrotor/demo_clf.py @@ -0,0 +1,148 @@ +""" +Search for a CLF. +This uses the polynomial dynamics of the 3D quadrotor (with quaternion as part +of the state), and the equality constraint to enforce the unit length constraint +on the quaternion. +""" + +import os +from typing import Tuple + +import numpy as np +import pydrake.solvers as solvers +import pydrake.symbolic as sym +import pydrake.systems.controllers + +from compatible_clf_cbf import clf +from examples.quadrotor.plant import QuadrotorPolyPlant + + +def lqr(quadrotor: QuadrotorPolyPlant) -> Tuple[np.ndarray, np.ndarray]: + x_des = np.zeros((13,)) + u_des = np.full((4,), quadrotor.m * quadrotor.g / 4) + xdot_des = quadrotor.dynamics(x_des, u_des) + np.testing.assert_allclose(xdot_des, np.zeros((13,)), atol=1e-8) + + A, B = quadrotor.linearize_dynamics(x_des, u_des) + Q = np.diag([1, 1, 1, 1, 1, 1, 1, 10, 10, 10, 10, 10, 10]) + R = np.diag([10, 10, 10, 10]) + # Gradient of the constraint quaternion.norm() = 1 + F = np.array([[2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]) + K, S = pydrake.systems.controllers.LinearQuadraticRegulator( + A, B, Q, R, N=np.empty((0, 4)), F=F + ) + return K, S + + +def find_trig_regional_clf( + V_degree: int, x: np.ndarray, save_pickle_filename="quadrotor_V_init.pkl" +) -> sym.Polynomial: + quadrotor = QuadrotorPolyPlant() + K_lqr, _ = lqr(quadrotor) + u_lqr = -K_lqr @ x + np.full((4,), quadrotor.m * quadrotor.g / 4) + dynamics_expr = quadrotor.dynamics(x, u_lqr) + dynamics = np.array([sym.Polynomial(dynamics_expr[i]) for i in range(13)]) + positivity_eps = 0.01 + d = int(V_degree / 2) + kappa = 0.1 + state_eq_constraints = quadrotor.equality_constraint(x) + positivity_ceq_lagrangian_degrees = [V_degree - 2] + derivative_ceq_lagrangian_degrees = [int(np.ceil((V_degree + 3) / 2) * 2 - 2)] + state_ineq_constraints = np.array([sym.Polynomial(x.dot(x) - 1e-3)]) + positivity_cin_lagrangian_degrees = [V_degree - 2] + derivative_cin_lagrangian_degrees = derivative_ceq_lagrangian_degrees + + prog, V = clf.find_candidate_regional_lyapunov( + x, + dynamics, + V_degree, + positivity_eps, + d, + kappa, + state_eq_constraints, + positivity_ceq_lagrangian_degrees, + derivative_ceq_lagrangian_degrees, + state_ineq_constraints, + positivity_cin_lagrangian_degrees, + derivative_cin_lagrangian_degrees, + ) + result = solvers.Solve(prog) + assert result.is_success() + V_sol = result.GetSolution(V) + x_set = sym.Variables(x) + pickle_path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "../../data/", save_pickle_filename + ) + clf.save_clf(V_sol, x_set, kappa, pickle_path) + return V_sol + + +def main(with_u_bound: bool): + x = sym.MakeVectorContinuousVariable(13, "x") + quadrotor = QuadrotorPolyPlant() + f, g = quadrotor.affine_dynamics(x) + + if with_u_bound: + thrust_max = quadrotor.m * quadrotor.g + u_vertices = thrust_max * np.array( + [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 1, 0, 0], + [0, 1, 0, 1], + [0, 1, 1, 0], + [0, 1, 1, 1], + [1, 0, 0, 0], + [1, 0, 0, 1], + [1, 0, 1, 0], + [1, 0, 1, 1], + [1, 1, 0, 0], + [1, 1, 0, 1], + [1, 1, 1, 0], + [1, 1, 1, 1], + ] + ) + else: + u_vertices = None + + state_eq_constraints = quadrotor.equality_constraint(x) + + V_degree = 2 + V_init = find_trig_regional_clf(V_degree, x) + kappa_V = 0.1 + solver_options = solvers.SolverOptions() + solver_options.SetOption(solvers.CommonSolverOption.kPrintToConsole, True) + + clf_search = clf.ControlLyapunov( + f=f, + g=g, + x=x, + x_equilibrium=np.zeros((13,)), + u_vertices=u_vertices, + state_eq_constraints=state_eq_constraints, + ) + if with_u_bound: + clf_lagrangian_degrees = clf.ClfWInputLimitLagrangianDegrees( + V_minus_rho=2, Vdot=[2] * 16, state_eq_constraints=[4] + ) + else: + clf_lagrangian_degrees = clf.ClfWoInputLimitLagrangianDegrees( + dVdx_times_f=4, + dVdx_times_g=[3, 3, 3, 3], + rho_minus_V=4, + state_eq_constraints=[4], + ) + clf_search.search_lagrangian_given_clf( + V_init, + rho=1, + kappa=kappa_V, + lagrangian_degrees=clf_lagrangian_degrees, + solver_options=solver_options, + ) + return + + +if __name__ == "__main__": + main(with_u_bound=True) diff --git a/examples/quadrotor/plant.py b/examples/quadrotor/plant.py new file mode 100644 index 0000000..9c1ae4d --- /dev/null +++ b/examples/quadrotor/plant.py @@ -0,0 +1,188 @@ +from typing import Tuple, Union + +import numpy as np +import jax +import jax.numpy as jnp +import pydrake.symbolic as sym + + +def quat2rotmat(quat: Union[np.ndarray, jnp.ndarray]) -> Union[np.ndarray, jnp.ndarray]: + if isinstance(quat, np.ndarray): + R = np.empty((3, 3), dtype=quat.dtype) + numpy_type = np + elif isinstance(quat, jnp.ndarray): + R = jnp.empty((3, 3)) + numpy_type = jnp + + def set_val(i, j, val): + if isinstance(R, np.ndarray): + R[i, j] = val + else: + R.at[i, j].set(val) + + R00 = 1 - 2 * quat[2] ** 2 - 2 * quat[3] ** 2 + R01 = 2 * quat[1] * quat[2] - 2 * quat[0] * quat[3] + R02 = 2 * quat[1] * quat[3] + 2 * quat[0] * quat[2] + R10 = 2 * quat[1] * quat[2] + 2 * quat[0] * quat[3] + R11 = 1 - 2 * quat[1] ** 2 - 2 * quat[3] ** 2 + R12 = 2 * quat[2] * quat[3] - 2 * quat[0] * quat[1] + R20 = 2 * quat[1] * quat[3] - 2 * quat[0] * quat[2] + R21 = 2 * quat[2] * quat[3] + 2 * quat[0] * quat[1] + R22 = 1 - 2 * quat[1] ** 2 - 2 * quat[2] ** 2 + R = numpy_type.stack((R00, R01, R02, R10, R11, R12, R20, R21, R22)).reshape((3, 3)) + return R + + +class QuadrotorPolyPlant: + """ + The state is x=(quat - [1, 0, 0, 0], pos, pos_dot, omega_WB_B) + Note that the equilibrium is at x=0. + The dynamics is a polynomial function of the state/action. + """ + + m: float + I: np.ndarray + l: float + g: float + kF: float + kM: float + + def __init__(self): + self.m = 0.775 + self.I = np.array([[0.0015, 0, 0], [0, 0.0025, 0], [0, 0, 0.0035]]) + self.g = 9.81 + self.l = 0.15 + self.kF = 1.0 + self.kM = 0.0245 + self.I_inv = np.linalg.inv(self.I) + + def dynamics( + self, x: Union[np.ndarray, jnp.ndarray], u: Union[np.ndarray, jnp.ndarray] + ) -> Union[np.ndarray, jnp.ndarray]: + """ + Dynamics xdot = f(x, u) + """ + if isinstance(x, np.ndarray): + numpy_type = np + np_array = np.array + elif isinstance(x, jnp.ndarray): + numpy_type = jnp + np_array = jnp.array + + uF_Bz = self.kF * u + + Faero_B = uF_Bz.sum() * np_array([0, 0, 1]) + Mx = self.l * (uF_Bz[1] - uF_Bz[3]) + My = self.l * (uF_Bz[2] - uF_Bz[0]) + uTau_Bz = self.kM * u + Mz = uTau_Bz[0] - uTau_Bz[1] + uTau_Bz[2] - uTau_Bz[3] + + tau_B = numpy_type.stack((Mx, My, Mz)) + Fgravity_N = np_array([0, 0, -self.m * self.g]) + + quat = x[:4] + np_array([1, 0, 0, 0]) + w_NB_B = x[-3:] + quat_dot = numpy_type.stack( + [ + 0.5 + * (-w_NB_B[0] * quat[1] - w_NB_B[1] * quat[2] - w_NB_B[2] * quat[3]), + 0.5 * (w_NB_B[0] * quat[0] + w_NB_B[2] * quat[2] - w_NB_B[1] * quat[3]), + 0.5 * (w_NB_B[1] * quat[0] - w_NB_B[2] * quat[1] + w_NB_B[0] * quat[3]), + 0.5 * (w_NB_B[2] * quat[0] + w_NB_B[1] * quat[1] - w_NB_B[0] * quat[2]), + ] + ) + R_NB = quat2rotmat(quat) + + xyzDDt = (Fgravity_N + R_NB @ Faero_B) / self.m + + wIw = numpy_type.cross(w_NB_B, self.I @ w_NB_B) + alpha_NB_B = self.I_inv @ (tau_B - wIw) + + xDt = numpy_type.concatenate((quat_dot, x[7:10], xyzDDt, alpha_NB_B)) + + return xDt + + def affine_dynamics(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Given the state x (as a vector of symbolic variables), return the affine + dynamics xdot = f(x) + g(x) * u, where f(x) is a vector of symbolic + polynomials, and g(x) is a matrix of symbolic polynomials. + """ + assert x.shape == (13,) + f = np.empty((13,), dtype=object) + g = np.empty((13, 4), dtype=object) + quat = x[:4] + np.array([1, 0, 0, 0]) + w_NB_B = x[-3:] + f[0] = sym.Polynomial( + 0.5 * (-w_NB_B[0] * quat[1] - w_NB_B[1] * quat[2] - w_NB_B[2] * quat[3]) + ) + f[1] = sym.Polynomial( + 0.5 * (w_NB_B[0] * quat[0] + w_NB_B[2] * quat[2] - w_NB_B[1] * quat[3]) + ) + f[2] = sym.Polynomial( + 0.5 * (w_NB_B[1] * quat[0] - w_NB_B[2] * quat[1] + w_NB_B[0] * quat[3]) + ) + f[3] = sym.Polynomial( + 0.5 * (w_NB_B[2] * quat[0] + w_NB_B[1] * quat[1] - w_NB_B[0] * quat[2]) + ) + f[4] = sym.Polynomial(x[7]) + f[5] = sym.Polynomial(x[8]) + f[6] = sym.Polynomial(x[9]) + for i in range(7): + for j in range(4): + g[i][j] = sym.Polynomial() + + f[7] = sym.Polynomial() + f[8] = sym.Polynomial() + f[9] = sym.Polynomial(-self.g) + + R_NB = quat2rotmat(quat) + + for i in range(3): + for j in range(4): + g[i + 7][j] = sym.Polynomial(R_NB[i, 2] * self.kF / self.m) + + wIw = np.cross(w_NB_B, self.I @ w_NB_B) + for i in range(3): + f[10 + i] = sym.Polynomial(-wIw[i] / self.I[i, i]) + g[10, 0] = sym.Polynomial() + g[10, 1] = sym.Polynomial(self.kF * self.l / self.I[0, 0]) + g[10, 2] = sym.Polynomial() + g[10, 3] = -g[10, 1] + g[11, 0] = sym.Polynomial(-self.kF * self.l / self.I[1, 1]) + g[11, 1] = sym.Polynomial() + g[11, 2] = -g[11, 0] + g[11, 3] = sym.Polynomial() + g[12, 0] = sym.Polynomial(self.kF * self.kM / self.I[2, 2]) + g[12, 1] = -g[12, 0] + g[12, 2] = -g[12, 1] + g[12, 3] = -g[12, 0] + return (f, g) + + def equality_constraint(self, x: np.ndarray) -> np.ndarray: + """ + Returns the left hand side of the equality constraint + quaternion.squared_norm() - 1 = 0 + """ + return np.array( + [sym.Polynomial(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + x[3] ** 2 + 2 * x[0])] + ) + + def linearize_dynamics( + self, x_val: np.ndarray, u_val: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + For the dynamics function xdot = f(x, u), compute the linearization as + A = ∂f/∂x, B = ∂f/∂u + """ + + x_jnp = jnp.array(x_val) + u_jnp = jnp.array(u_val) + + def f(xu): + return self.dynamics(xu[:13], xu[-4:]) + + xu_jnp = jnp.concatenate((x_jnp, u_jnp)) + AB_jnp = jax.jacfwd(f)(xu_jnp) + AB = np.array(AB_jnp) + return AB[:, :13], AB[:, 13:] diff --git a/examples/quadrotor2d/demo.py b/examples/quadrotor2d/demo.py index 15286a3..2d34452 100644 --- a/examples/quadrotor2d/demo.py +++ b/examples/quadrotor2d/demo.py @@ -14,6 +14,7 @@ from compatible_clf_cbf import clf_cbf from examples.quadrotor2d.plant import Quadrotor2dTrigPlant import examples.quadrotor2d.demo_clf as demo_clf +import compatible_clf_cbf.utils def main(use_y_squared: bool, with_u_bound: bool): @@ -54,14 +55,17 @@ def main(use_y_squared: bool, with_u_bound: bool): compatible_lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees( lambda_y=[ - clf_cbf.CompatibleLagrangianDegrees.Degree(x=1, y=1) for _ in range(2) + clf_cbf.CompatibleLagrangianDegrees.Degree(x=1, y=0 if use_y_squared else 1) + for _ in range(2) ], - xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=1), + xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree( + x=2, y=0 if use_y_squared else 1 + ), y=( None if use_y_squared else [ - clf_cbf.CompatibleLagrangianDegrees.Degree(x=4, y=1) + clf_cbf.CompatibleLagrangianDegrees.Degree(x=4, y=0) for _ in range(compatible.y.size) ] ), @@ -91,6 +95,7 @@ def main(use_y_squared: bool, with_u_bound: bool): weight_b=np.array([1]), ) + max_iter = 5 V, b = compatible.bilinear_alternation( V_init, b_init, @@ -102,12 +107,15 @@ def main(use_y_squared: bool, with_u_bound: bool): x_equilibrium, V_degree, b_degrees, - max_iter=0, + max_iter=max_iter, solver_options=solver_options, # solver_id = solvers.ClarabelSolver().id(), lagrangian_coefficient_tol=None, compatible_states_options=compatible_states_options, - backoff_scale=0.02, + backoff_scales=[ + compatible_clf_cbf.utils.BackoffScale(rel=0.02, abs=None) + for _ in range(max_iter) + ], ) x_set = sym.Variables(x) pickle_path = os.path.join( diff --git a/examples/quadrotor2d/demo_taylor.py b/examples/quadrotor2d/demo_taylor.py index f788e36..584a5d2 100644 --- a/examples/quadrotor2d/demo_taylor.py +++ b/examples/quadrotor2d/demo_taylor.py @@ -142,7 +142,10 @@ def search_clf_cbf( inner_ellipsoid_options=inner_ellipsoid_options, binary_search_scale_options=binary_search_scale_options, compatible_states_options=compatible_states_options, - backoff_scale=0.02, + backoff_scales=[ + compatible_clf_cbf.utils.BackoffScale(rel=0.02, abs=None) + for _ in range(max_iter) + ], ) return V, b diff --git a/tests/examples/quadrotor/test_plant.py b/tests/examples/quadrotor/test_plant.py new file mode 100644 index 0000000..7f0a5fc --- /dev/null +++ b/tests/examples/quadrotor/test_plant.py @@ -0,0 +1,107 @@ +import examples.quadrotor.plant as mut + +import numpy as np +import pytest # noqa +import pydrake.symbolic as sym +import pydrake.math + + +def test_quat2rotmat(): + quat = np.array([1, 2, 3, 4.0]) + quat = quat / np.linalg.norm(quat) + R = mut.quat2rotmat(quat) + R_expected = pydrake.math.RotationMatrix( + pydrake.common.eigen_geometry.Quaternion(quat) + ).matrix() + np.testing.assert_allclose(R, R_expected) + + +class TestQuadrotorPolyPlant: + def test_dynamics(self): + dut = mut.QuadrotorPolyPlant() + quat = np.array([0.3, -0.4, 0.5, -np.sqrt(2) / 2]) + pos = np.array([1, 2, 3.0]) + pos_dot = np.array([0.5, -0.3, 0.2]) + omega_WB_B = np.array([0.3, 0.4, 0.5]) + u = np.array([1, 2, 3, 4.0]) + x = np.concatenate((quat - np.array([1, 0, 0, 0]), pos, pos_dot, omega_WB_B)) + xdot = dut.dynamics(x, u) + assert xdot.shape == (13,) + quat_dot = xdot[:4] + np.testing.assert_almost_equal(quat.dot(quat_dot), 0, 6) + + def test_affine_dynamics(self): + dut = mut.QuadrotorPolyPlant() + x = sym.MakeVectorContinuousVariable(13, "x") + f, g = dut.affine_dynamics(x) + assert f.shape == (13,) + assert g.shape == (13, 4) + + def compare_xdot(quat_val, pos_val, pos_dot_val, omega_val, u_val): + x_val = np.concatenate( + (quat_val - np.array([1, 0, 0, 0.0]), pos_val, pos_dot_val, omega_val) + ) + xdot = dut.dynamics(x_val, u_val) + env = {x[i]: x_val[i] for i in range(13)} + f_val = np.array([f[i].Evaluate(env) for i in range(13)]) + g_val = np.array( + [[g[i, j].Evaluate(env) for j in range(4)] for i in range(13)] + ) + xdot_expected = f_val + g_val @ u_val + np.testing.assert_allclose(xdot, xdot_expected) + + compare_xdot( + np.array([1, 0, 0, 0.0]), + np.array([1, 2, 3.0]), + np.array([0.5, -1, 2]), + np.array([2, 3, 1]), + np.array([1, 2, 3, 4]), + ) + compare_xdot( + np.array([0.3, 0.4, 0.5, np.sqrt(2) / 2]), + np.array([1, 2, 3.0]), + np.array([0.5, -1, 2]), + np.array([2, 3, 1]), + np.array([1, 2, 3, 4]), + ) + + def test_equality_constraint(self): + dut = mut.QuadrotorPolyPlant() + x = sym.MakeVectorContinuousVariable(13, "x") + state_eq_constraints = dut.equality_constraint(x) + assert state_eq_constraints.shape == (1,) + + def check_constraint(quat_val, pos_val, pos_dot_val, omega_val): + x_val = np.concatenate( + (quat_val - np.array([1, 0, 0, 0]), pos_val, pos_dot_val, omega_val) + ) + env = {x[i]: x_val[i] for i in range(13)} + val = state_eq_constraints[0].Evaluate(env) + np.testing.assert_allclose(val, 0, atol=1e-5) + + check_constraint( + np.array([0.3, 0.4, 0.5, np.sqrt(2) / 2]), + np.array([1, 2, 3.0]), + np.array([0.5, -1, 2]), + np.array([2, 3, 1]), + ) + + def test_linearize_dynamics(self): + dut = mut.QuadrotorPolyPlant() + x = sym.MakeVectorContinuousVariable(13, "x") + f, g = dut.affine_dynamics(x) + + def check(x_val, u_val): + A, B = dut.linearize_dynamics(x_val, u_val) + env = {x[i]: x_val[i] for i in range(13)} + g_val = np.array( + [[g[i, j].Evaluate(env) for j in range(4)] for i in range(13)] + ) + np.testing.assert_allclose(B, g_val, atol=1e-5) + + x_val = np.array( + [-0.7, 0.4, 0.5, np.sqrt(2) / 2, 0.2, 1, 4, -2, 3, 1, 0.5, 1, 2] + ) + u_val = np.array([1, 2, 3, 4.0]) + A, B = dut.linearize_dynamics(x_val, u_val) + check(x_val, u_val) diff --git a/tests/examples/quadrotor2d/test_plant.py b/tests/examples/quadrotor2d/test_plant_2d.py similarity index 100% rename from tests/examples/quadrotor2d/test_plant.py rename to tests/examples/quadrotor2d/test_plant_2d.py diff --git a/tests/test_utils.py b/tests/test_utils.py index 69fc454..daf3d8a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -121,7 +121,7 @@ def test_solve_w_id(): prog.AddBoundingBoxConstraint(-1, 1, x) prog.AddLinearCost(x[0] + x[1] + 1) result = mut.solve_with_id( - prog, solver_id=None, solver_options=None, backoff_scale=0.1 + prog, solver_id=None, solver_options=None, backoff_rel_scale=0.1 ) assert result.is_success() # I know the optimal solution is obtained at (-1, -1), with the optimal cost being @@ -132,7 +132,7 @@ def test_solve_w_id(): # x[0] + x[1] = -1.9 prog.AddLinearCost(-x[0] - x[1]) result = mut.solve_with_id( - prog, solver_id=None, solver_options=None, backoff_scale=None + prog, solver_id=None, solver_options=None, backoff_rel_scale=None ) x_sol = result.GetSolution(x) np.testing.assert_allclose(x_sol[0] + x_sol[1], -1.9, atol=1e-5) @@ -143,7 +143,7 @@ def test_solve_w_id(): prog.AddBoundingBoxConstraint(-1, 1, x) prog.AddLinearCost(x[0] + x[1] + 3) result = mut.solve_with_id( - prog, solver_id=None, solver_options=None, backoff_scale=0.1 + prog, solver_id=None, solver_options=None, backoff_rel_scale=0.1 ) assert result.is_success() # I know the optimal solution is obtained at (-1, -1), with the optimal cost being @@ -154,7 +154,7 @@ def test_solve_w_id(): # x[0] + x[1] = -1.9 prog.AddLinearCost(-x[0] - x[1]) result = mut.solve_with_id( - prog, solver_id=None, solver_options=None, backoff_scale=None + prog, solver_id=None, solver_options=None, backoff_rel_scale=None ) x_sol = result.GetSolution(x) np.testing.assert_allclose(x_sol[0] + x_sol[1], -1.9, atol=1e-5)