Add more simulation on the quadrotor

hongkai-dai · Sep 10, 2024 · bc7fad5 · bc7fad5
1 parent 3b9f10b
commit bc7fad5
Show file tree

Hide file tree

Showing 6 changed files with 133 additions and 55 deletions.
diff --git a/compatible_clf_cbf/controller.py b/compatible_clf_cbf/controller.py
@@ -87,10 +87,12 @@ def calc_V(self, context: pydrake.systems.framework.Context, output):
         x_val: np.ndarray = self.get_input_port(0).Eval(context)
         env = {self.x[i]: x_val[i] for i in range(self.x.size)}
         V_val = self.V.Evaluate(env)
+        print(f"time={context.get_time()}, V_val={V_val}")
         output.set_value(np.array([V_val]))
 
     def calc_b(self, context: pydrake.systems.framework.Context, output):
         x_val: np.ndarray = self.get_input_port(0).Eval(context)
         env = {self.x[i]: x_val[i] for i in range(self.x.size)}
         b_val = np.array([b_i.Evaluate(env) for b_i in self.b])
+        print(f"time={context.get_time()}, b_val={b_val}")
         output.set_value(b_val)
diff --git a/examples/nonlinear_toy/demo_trigpoly.py b/examples/nonlinear_toy/demo_trigpoly.py
@@ -276,7 +276,7 @@ def visualize():
     fig = plt.figure()
     ax = fig.add_subplot()
     ax.set_xlabel(r"$\theta$ (rad)", fontsize=16)
-    ax.set_ylabel(r"$\dot{\theta} (rad/s)$", fontsize=16)
+    ax.set_ylabel(r"$\gamma$", fontsize=16)
     ax.set_xticks([-np.pi, -np.pi / 2, 0, np.pi / 2, np.pi])
     ax.set_xticklabels(
         [r"$-\pi$", r"$-\frac{\pi}{2}$", r"0", r"$\frac{\pi}{2}$", r"$\pi$"]
@@ -302,9 +302,10 @@ def visualize():
             h_V_init.legend_elements()[0][0],
             h_b_init.legend_elements()[0][0],
         ],
-        [r"$V(x)=1$", r"$b(x)=0$", r"$V_{init}(x)=1$", r"$b_{init}(x)=0$"],
+        [r"$V(x)=1$", r"$h(x)=0$", r"$V_{initial}(x)=1$", r"$h_{initial}(x)=0$"],
         prop={"size": 12},
     )
+    ax.set_title("Compatible CLF/CBF")
     fig.show()
     for fig_extension in (".png", ".pdf"):
         fig.savefig(

diff --git a/examples/nonlinear_toy/demo_trigpoly_incompatible.py b/examples/nonlinear_toy/demo_trigpoly_incompatible.py
@@ -167,9 +167,10 @@ def visualize():
             h_V.legend_elements()[0][0],
             h_b.legend_elements()[0][0],
         ],
-        [r"$V(x)=1$", r"$b(x)=0$"],
+        [r"$V(x)=1$", r"$h(x)=0$"],
         prop={"size": 12},
     )
+    ax.set_title("Incompatible CLF/CBF")
     fig.show()
     for fig_extension in (".png", ".pdf"):
         fig.savefig(

diff --git a/examples/quadrotor/demo.py b/examples/quadrotor/demo.py
@@ -24,6 +24,7 @@
 import pydrake.symbolic as sym
 import pydrake.systems.analysis
 from pydrake.systems.framework import Diagram, DiagramBuilder
+import pydrake.systems.framework
 from pydrake.systems.primitives import LogVectorOutput, VectorLogSink
 
 
@@ -84,18 +85,18 @@ def search(use_y_squared: bool, with_u_bound: bool):
 
     b_init = np.array([1 - V_init])
 
-    load_clf_cbf = False
+    load_clf_cbf = True
     if load_clf_cbf:
         data = clf_cbf.load_clf_cbf(
             os.path.join(
                 os.path.dirname(os.path.abspath(__file__)),
-                "../../data/quadrotor_clf_cbf7.pkl",
+                "../../data/quadrotor_clf_cbf8.pkl",
             ),
             x_set,
         )
         V_init = data["V"]
         b_init = data["b"]
-    kappa_V_sequences = [0.05, 0.075, 0.1, 0.125, 0.15, 0.175, 0.2, 0.2]
+    kappa_V_sequences = [0.05, 0.075, 0.1, 0.125, 0.15, 0.175, 0.2, 0.2, 0.25, 0.625]
     kappa_b_sequences = [np.array([0.2]) for _ in range(len(kappa_V_sequences))]
 
     compatible_lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees(
@@ -211,6 +212,30 @@ def search(use_y_squared: bool, with_u_bound: bool):
     candidate_compatible_states_sequences[-1][3, 6] = -0.35
     candidate_compatible_states_sequences[-1][4, 4] = -1
     candidate_compatible_states_sequences[-1][4, 6] = -0.35
+
+    candidate_compatible_states_sequences.append(np.zeros((7, 13)))
+    candidate_compatible_states_sequences[-1][1, 5] = -1
+    candidate_compatible_states_sequences[-1][1, 6] = -0.35
+    candidate_compatible_states_sequences[-1][2, 5] = 1
+    candidate_compatible_states_sequences[-1][2, 6] = -0.35
+    candidate_compatible_states_sequences[-1][3, 4] = 1
+    candidate_compatible_states_sequences[-1][3, 6] = -0.35
+    candidate_compatible_states_sequences[-1][4, 4] = -1
+    candidate_compatible_states_sequences[-1][4, 6] = -0.35
+    candidate_compatible_states_sequences[-1][5, 5] = -1
+    candidate_compatible_states_sequences[-1][5, :4] = np.array(
+        [np.cos(np.pi / 36) - 1, np.sin(np.pi / 36), 0, 0]
+    )
+    candidate_compatible_states_sequences[-1][6, 5] = 1
+    candidate_compatible_states_sequences[-1][6, :4] = np.array(
+        [np.cos(np.pi / 18) - 1, np.sin(np.pi / 18), 0, 0]
+    )
+
+    candidate_compatible_states_sequences.append(np.zeros((7, 13)))
+    candidate_compatible_states_sequences[-1] = candidate_compatible_states_sequences[
+        -2
+    ]
+
     lagrangian_sos_types = [
         solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
         solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
@@ -220,17 +245,10 @@ def search(use_y_squared: bool, with_u_bound: bool):
         solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
         solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
         solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
+        solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
+        solvers.MathematicalProgram.NonnegativePolynomial.kSdsos,
     ]
-    V_margin_sequence = [
-        None,
-        None,
-        None,
-        None,
-        None,
-        0.01,
-        0.01,
-        0.01,
-    ]
+    V_margin_sequence = [None, None, None, None, None, 0.01, 0.01, 0.01, 0.05, 0.05]
     b_margins_sequence = [
         None,
         None,
@@ -240,12 +258,14 @@ def search(use_y_squared: bool, with_u_bound: bool):
         np.array([0.02]),
         np.array([0.03]),
         np.array([0.03]),
+        np.array([0.02]),
+        np.array([0.02]),
     ]
     solver_options = solvers.SolverOptions()
     solver_options.SetOption(solvers.CommonSolverOption.kPrintToConsole, True)
     V = V_init
     b = b_init
-    max_iter_sequence = [1, 1, 1, 2, 1, 1, 2, 2]
+    max_iter_sequence = [1, 1, 1, 2, 1, 1, 2, 2, 3, 1, 1]
     backoff_scale_sequence = [
         BackoffScale(rel=None, abs=0.2),
         BackoffScale(rel=None, abs=0.2),
@@ -255,8 +275,10 @@ def search(use_y_squared: bool, with_u_bound: bool):
         BackoffScale(rel=None, abs=0.1),
         BackoffScale(rel=None, abs=0.1),
         BackoffScale(rel=None, abs=0.15),
+        BackoffScale(rel=None, abs=0.1),
+        BackoffScale(rel=None, abs=0.1),
     ]
-    for i in range(len(candidate_compatible_states_sequences)):
+    for i in range(9, 10):  # len(candidate_compatible_states_sequences)):
         compatible_states_options = clf_cbf.CompatibleStatesOptions(
             candidate_compatible_states=candidate_compatible_states_sequences[i],
             anchor_states=np.zeros((1, 13)),
@@ -307,7 +329,7 @@ def build_diagram() -> Tuple[
     quadrotor = builder.AddSystem(QuadrotorPolyPlant())
     scene_graph = builder.AddSystem(SceneGraph())
     QuadrotorPolyGeometry.AddToBuilder(
-        builder, quadrotor.get_output_port(0), scene_graph
+        builder, quadrotor.get_output_port(0), "quadrotor", scene_graph
     )
 
     meshcat = StartMeshcat()
@@ -328,7 +350,7 @@ def build_diagram() -> Tuple[
 
     pickle_path = os.path.join(
         os.path.dirname(os.path.abspath(__file__)),
-        "../../data/quadrotor_clf_cbf7.pkl",
+        "../../data/quadrotor_clf_cbf9.pkl",
     )
     clf_cbf_data = clf_cbf.load_clf_cbf(pickle_path, x_set)
     V = clf_cbf_data["V"]
@@ -387,37 +409,82 @@ def simulate(x0: np.ndarray, duration: float):
     ) = build_diagram()
     simulator = pydrake.systems.analysis.Simulator(diagram)
     simulator.get_mutable_context().SetContinuousState(x0)
+
+    def monitor(context):
+        quadrotor_context = quadrotor.GetMyContextFromRoot(context)
+        x_val = quadrotor.get_output_port(0).Eval(quadrotor_context)
+        if np.linalg.norm(x_val) < 1e-2:
+            return pydrake.systems.framework.EventStatus.ReachedTermination(
+                diagram, "reach_goal"
+            )
+        return pydrake.systems.framework.EventStatus.Succeeded()
+
+    simulator.set_monitor(monitor)
+    simulator_config = pydrake.systems.analysis.SimulatorConfig(
+        integration_scheme="runge_kutta3"
+    )
+    pydrake.systems.analysis.ApplySimulatorConfig(simulator_config, simulator)
     simulator.AdvanceTo(duration)
 
     state_data = state_logger.FindLog(simulator.get_context()).data()
     action_data = action_logger.FindLog(simulator.get_context()).data()
     clf_data = clf_logger.FindLog(simulator.get_context()).data()
     cbf_data = cbf_logger.FindLog(simulator.get_context()).data()
-    return state_data, action_data, clf_data, cbf_data
+    time_data = state_logger.FindLog(simulator.get_context()).sample_times()
+    return state_data, action_data, clf_data, cbf_data, time_data
 
 
 def run_simulations():
     x0_sequences = []
     x0_sequences.append(np.zeros((13,)))
-    x0_sequences[-1][4:7] = np.array([1, 0, -0.3])
+    x0_sequences[0][4:7] = np.array([1, 0, -0.3])
     x0_sequences.append(np.zeros((13,)))
-    x0_sequences[-1][4:7] = np.array([-1, 0, -0.3])
+    x0_sequences[1][4:7] = np.array([-1, 0, -0.3])
     x0_sequences.append(np.zeros((13,)))
-    x0_sequences[-1][4:7] = np.array([0, 1, -0.3])
+    x0_sequences[2][4:7] = np.array([0, 1, -0.3])
     x0_sequences.append(np.zeros((13,)))
-    x0_sequences[-1][4:7] = np.array([0, -1, -0.3])
-    for i, x0 in enumerate(x0_sequences):
-        state_data, action_data, clf_data, cbf_data = simulate(x0, duration=40)
+    x0_sequences[3][4:7] = np.array([0, -1, -0.3])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[4][4:7] = np.array([0.2, -0.8, 0.1])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[5][4:7] = np.array([0.1, -0.95, 0])
+    x0_sequences[5][:4] = np.array([np.cos(np.pi / 36) - 1, np.sin(np.pi / 36), 0, 0])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[6][4:7] = np.array([0.5, 0.7, 0.1])
+    x0_sequences[6][:4] = np.array([np.cos(np.pi / 20) - 1, np.sin(np.pi / 20), 0, 0])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[7][4:7] = np.array([-0.4, 0.75, -0.1])
+    x0_sequences[7][:4] = np.array([np.cos(np.pi / 15) - 1, np.sin(np.pi / 15), 0, 0])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[8][4:7] = np.array([-0.6, -0.55, -0.2])
+    x0_sequences[8][:4] = np.array([np.cos(np.pi / 18) - 1, 0, np.sin(np.pi / 18), 0])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[9][4:7] = np.array([0.6, -0.15, 0.05])
+    x0_sequences[9][:4] = np.array([np.cos(np.pi / 15) - 1, 0, np.sin(-np.pi / 15), 0])
+    x0_sequences.append(np.zeros((13,)))
+    x0_sequences[10][4:7] = np.array([-0.8, -0.05, -0.4])
+    x0_sequences[10][:4] = np.array([np.cos(np.pi / 30) - 1, 0, np.sin(np.pi / 30), 0])
+    for i in range(len(x0_sequences)):
+        state_data, action_data, clf_data, cbf_data, time_data = simulate(
+            x0_sequences[i], duration=60
+        )
         path = os.path.join(
             os.path.dirname(os.path.abspath(__file__)),
             f"../../data/quadrotor_sim{i}.npz",
         )
-        np.savez(path, state_data, action_data, clf_data, cbf_data)
+        np.savez(
+            path,
+            state_data=state_data,
+            action_data=action_data,
+            clf_data=clf_data,
+            cbf_data=cbf_data,
+            time_data=time_data,
+        )
 
 
 def main():
-    # search(use_y_squared=True, with_u_bound=False)
-    run_simulations()
+    search(use_y_squared=True, with_u_bound=False)
+    # run_simulations()
 
 
 if __name__ == "__main__":

diff --git a/examples/quadrotor/plant.py b/examples/quadrotor/plant.py
@@ -131,52 +131,56 @@ def affine_dynamics(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         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.
         """
+        if x.dtype == object:
+            data_type = sym.Polynomial
+        else:
+            data_type = float
         assert x.shape == (13,)
-        f = np.empty((13,), dtype=object)
-        g = np.empty((13, 4), dtype=object)
+        f = np.empty((13,), dtype=x.dtype)
+        g = np.empty((13, 4), dtype=x.dtype)
         quat = x[:4] + np.array([1, 0, 0, 0])
         w_NB_B = x[-3:]
-        f[0] = sym.Polynomial(
+        f[0] = data_type(
             0.5 * (-w_NB_B[0] * quat[1] - w_NB_B[1] * quat[2] - w_NB_B[2] * quat[3])
         )
-        f[1] = sym.Polynomial(
+        f[1] = data_type(
             0.5 * (w_NB_B[0] * quat[0] + w_NB_B[2] * quat[2] - w_NB_B[1] * quat[3])
         )
-        f[2] = sym.Polynomial(
+        f[2] = data_type(
             0.5 * (w_NB_B[1] * quat[0] - w_NB_B[2] * quat[1] + w_NB_B[0] * quat[3])
         )
-        f[3] = sym.Polynomial(
+        f[3] = data_type(
             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])
+        f[4] = data_type(x[7])
+        f[5] = data_type(x[8])
+        f[6] = data_type(x[9])
         for i in range(7):
             for j in range(4):
-                g[i][j] = sym.Polynomial()
+                g[i][j] = data_type(0)
 
-        f[7] = sym.Polynomial()
-        f[8] = sym.Polynomial()
-        f[9] = sym.Polynomial(-self.g)
+        f[7] = data_type(0)
+        f[8] = data_type(0)
+        f[9] = data_type(-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)
+                g[i + 7][j] = data_type(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()
+            f[10 + i] = data_type(-wIw[i] / self.I[i, i])
+        g[10, 0] = data_type(0)
+        g[10, 1] = data_type(self.kF * self.l / self.I[0, 0])
+        g[10, 2] = data_type(0)
         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, 0] = data_type(-self.kF * self.l / self.I[1, 1])
+        g[11, 1] = data_type(0)
         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[11, 3] = data_type(0)
+        g[12, 0] = data_type(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]
@@ -217,7 +221,9 @@ def __init__(
     ):
         super().__init__()
         mbp = pydrake.multibody.plant.MultibodyPlant(0.0)
-        parser = pydrake.multibody.parsing.Parser(mbp, scene_graph)
+        parser = pydrake.multibody.parsing.Parser(
+            mbp, scene_graph, "quadrotor" if name is None else name
+        )
         model_instance_indices = parser.AddModelsFromUrl(
             "package://drake_models/skydio_2/quadrotor.urdf"
         )
@@ -256,9 +262,10 @@ def output_geometry_pose(
     def AddToBuilder(
         builder: pydrake.systems.framework.DiagramBuilder,
         quadrotor_state_port: pydrake.systems.framework.OutputPort,
+        name: str,
         scene_graph: pydrake.geometry.SceneGraph,
     ) -> Self:
-        quadrotor_geometry = builder.AddSystem(QuadrotorPolyGeometry(scene_graph))
+        quadrotor_geometry = builder.AddSystem(QuadrotorPolyGeometry(scene_graph, name))
         builder.Connect(quadrotor_state_port, quadrotor_geometry.get_input_port(0))
         builder.Connect(
             quadrotor_geometry.get_output_port(0),

diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 black
 black[jupyter]
-drake
+drake==1.31.0
 flake8
 ipykernel
 jax