Clean up formatting. Add test for symbolic and sim model mismatch

Author: amacati

crazyflow/sim/symbolic.py

@@ -19,8 +19,7 @@
 from casadi import MX
 from numpy.typing import NDArray
 
-from crazyflow.constants import ARM_LEN, GRAVITY, SIGN_MIX_MATRIX
-from crazyflow.control.control import KF, KM
+from crazyflow.constants import GRAVITY
 from crazyflow.sim import Sim
 
 
@@ -156,22 +155,22 @@ def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
     # # Set up the dynamics model for a 3D quadrotor.
     nx, nu = 12, 4
     # Define states.
-    x = cs.MX.sym('x')
-    x_dot = cs.MX.sym('x_dot')
-    y = cs.MX.sym('y')
-    y_dot = cs.MX.sym('y_dot')
-    phi = cs.MX.sym('phi')  # roll angle [rad]
-    phi_dot = cs.MX.sym('phi_dot')
-    theta = cs.MX.sym('theta')  # pitch angle [rad]
-    theta_dot = cs.MX.sym('theta_dot')
-    psi = cs.MX.sym('psi')  # yaw angle [rad]
-    psi_dot = cs.MX.sym('psi_dot')
+    x = cs.MX.sym("x")
+    x_dot = cs.MX.sym("x_dot")
+    y = cs.MX.sym("y")
+    y_dot = cs.MX.sym("y_dot")
+    phi = cs.MX.sym("phi")  # roll angle [rad]
+    phi_dot = cs.MX.sym("phi_dot")
+    theta = cs.MX.sym("theta")  # pitch angle [rad]
+    theta_dot = cs.MX.sym("theta_dot")
+    psi = cs.MX.sym("psi")  # yaw angle [rad]
+    psi_dot = cs.MX.sym("psi_dot")
     X = cs.vertcat(x, x_dot, y, y_dot, z, z_dot, phi, theta, psi, phi_dot, theta_dot, psi_dot)
     # Define input collective thrust and theta.
-    T = cs.MX.sym('T_c')  # normalized thrust [N]
-    R = cs.MX.sym('R_c')  # desired roll angle [rad]
-    P = cs.MX.sym('P_c')  # desired pitch angle [rad]
-    Y = cs.MX.sym('Y_c')  # desired yaw angle [rad]
+    T = cs.MX.sym("T_c")  # normalized thrust [N]
+    R = cs.MX.sym("R_c")  # desired roll angle [rad]
+    P = cs.MX.sym("P_c")  # desired pitch angle [rad]
+    Y = cs.MX.sym("Y_c")  # desired yaw angle [rad]
     U = cs.vertcat(T, R, P, Y)
     # The thrust in PWM is converted from the normalized thrust.
     # With the formulat F_desired = b_F * T + a_F
@@ -185,24 +184,25 @@ def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
 
     # Define dynamics equations.
     # TODO: create a parameter for the new quad model
-    X_dot = cs.vertcat(x_dot,
-                        (params_acc[0] * T + params_acc[1]) * (
-                            cs.cos(phi) * cs.sin(theta) * cs.cos(psi) + cs.sin(phi) * cs.sin(psi)),
-                        y_dot,
-                        (params_acc[0] * T + params_acc[1]) * (
-                            cs.cos(phi) * cs.sin(theta) * cs.sin(psi) - cs.sin(phi) * cs.cos(psi)),
-                        z_dot,
-                        (params_acc[0] * T + params_acc[1]) * cs.cos(phi) * cs.cos(theta) - g,
-                        phi_dot,
-                        theta_dot,
-                        psi_dot,
-                        params_roll_rate[0] * phi + params_roll_rate[1] * phi_dot + params_roll_rate[2] * R,
-                        params_pitch_rate[0] * theta + params_pitch_rate[1] * theta_dot + params_pitch_rate[2] * P,
-                        params_yaw_rate[0] * psi + params_yaw_rate[1] * psi_dot + params_yaw_rate[2] * Y)
+    X_dot = cs.vertcat(
+        x_dot,
+        (params_acc[0] * T + params_acc[1])
+        * (cs.cos(phi) * cs.sin(theta) * cs.cos(psi) + cs.sin(phi) * cs.sin(psi)),
+        y_dot,
+        (params_acc[0] * T + params_acc[1])
+        * (cs.cos(phi) * cs.sin(theta) * cs.sin(psi) - cs.sin(phi) * cs.cos(psi)),
+        z_dot,
+        (params_acc[0] * T + params_acc[1]) * cs.cos(phi) * cs.cos(theta) - g,
+        phi_dot,
+        theta_dot,
+        psi_dot,
+        params_roll_rate[0] * phi + params_roll_rate[1] * phi_dot + params_roll_rate[2] * R,
+        params_pitch_rate[0] * theta + params_pitch_rate[1] * theta_dot + params_pitch_rate[2] * P,
+        params_yaw_rate[0] * psi + params_yaw_rate[1] * psi_dot + params_yaw_rate[2] * Y,
+    )
     # Define observation.
     Y = cs.vertcat(x, x_dot, y, y_dot, z, z_dot, phi, theta, psi, phi_dot, theta_dot, psi_dot)
 
-
     # Define cost (quadratic form).
     Q, R = MX.sym("Q", nx, nx), MX.sym("R", nu, nu)
     Xr, Ur = MX.sym("Xr", nx, 1), MX.sym("Ur", nu, 1)

tests/integration/test_symbolic.py

@@ -0,0 +1,65 @@
+import numpy as np
+import pytest
+from numpy.typing import NDArray
+from scipy.spatial.transform import Rotation as R
+
+from crazyflow.control.control import MAX_THRUST, MIN_THRUST
+from crazyflow.sim import Sim
+from crazyflow.sim.structs import SimState
+from crazyflow.sim.symbolic import symbolic_from_sim
+
+
+def sim_state2symbolic_state(state: SimState) -> NDArray[np.float32]:
+    """Convert the simulation state to the symbolic state vector."""
+    pos = state.pos.squeeze()  # shape: (3,)
+    vel = state.vel.squeeze()  # shape: (3,)
+    euler = R.from_quat(state.quat.squeeze()).as_euler("xyz")  # shape: (3,), Euler angles
+    rpy_rates = state.rpy_rates.squeeze()  # shape: (3,)
+    return np.array([pos[0], vel[0], pos[1], vel[1], pos[2], vel[2], *euler, *rpy_rates])
+
+
+@pytest.mark.integration
+def test_attitude_symbolic():
+    sim = Sim(physics="sys_id")
+    sym = symbolic_from_sim(sim)
+
+    x0 = np.zeros(12)
+
+    # Simulate with both models for 0.5 seconds
+    t_end = 0.5
+    dt = 1 / sim.freq
+    steps = int(t_end / dt)
+
+    # Track states over time
+    x_sym_log = []
+    x_sim_log = []
+
+    # Initialize logs with initial state
+    x_sym = x0.copy()
+    x_sim = x0.copy()
+    x_sym_log.append(x_sym)
+    x_sim_log.append(x_sim)
+
+    u_low = np.array([4 * MIN_THRUST, -np.pi, -np.pi, -np.pi])
+    u_high = np.array([4 * MAX_THRUST, np.pi, np.pi, np.pi])
+
+    # Run simulation
+    for _ in range(steps):
+        u_rand = (np.random.rand(4) * (u_high - u_low) + u_low).astype(np.float32)
+        # Simulate with symbolic model
+        res = sym.fd_func(x0=x_sym, p=u_rand)
+        x_sym = res["xf"].full().flatten()
+        x_sym_log.append(x_sym)
+        # Simulate with attitude controller
+        sim.attitude_control(u_rand.reshape(1, 1, 4))
+        sim.step(sim.freq // sim.control_freq)
+        x_sim_log.append(sim_state2symbolic_state(sim.data.states))
+
+    # Convert logs to arrays. Do not record the rpy rates (deviate easily).
+    x_sym_log = np.array(x_sym_log)[..., :-3]
+    x_sim_log = np.array(x_sim_log)[..., :-3]
+
+    # Check if states match throughout simulation
+    err_msg = "Symbolic and simulation prediction do not match approximately"
+    assert np.allclose(x_sym_log, x_sim_log, rtol=1e-2, atol=1e-3), err_msg
+    sim.close()

tutorials/LQR_ILQR.ipynb

@@ -33,10 +33,9 @@
     "from scipy.spatial.transform import Rotation as R\n",
     "\n",
     "from crazyflow.constants import GRAVITY, MASS, J\n",
-    "from crazyflow.sim.physics import Physics\n",
     "from crazyflow.control import Control\n",
-    "from crazyflow.control.control import MAX_THRUST, MIN_THRUST\n",
-    "from crazyflow.sim.symbolic import symbolic\n"
+    "from crazyflow.sim.physics import Physics\n",
+    "from crazyflow.sim.symbolic import symbolic"
    ]
   },
   {
@@ -175,9 +174,7 @@
     "A, B = df[0].toarray(), df[1].toarray()\n",
     "\n",
     "print(\"A shape:\", A.shape)  # Should be (12, 12)\n",
-    "print(\"B shape:\", B.shape)  # Should be (12, 4)\n",
-    "# print(\"A :\\n\", A)\n",
-    "# print(\"B :\\n\", B)"
+    "print(\"B shape:\", B.shape)  # Should be (12, 4)"
    ]
   },
   {
@@ -284,7 +281,7 @@
     "# gain_lqr = np.dot(np.linalg.inv(R_lqr), np.dot(B.T, P))\n",
     "\n",
     "print(\"gain:\\n\", gain_lqr)\n",
-    "print(\"shape of gain:\", gain_lqr.shape)\n"
+    "print(\"shape of gain:\", gain_lqr.shape)"
    ]
   },
   {
@@ -322,12 +319,10 @@
     "\n",
     "    control_input = np.clip(control_input, envs.action_space.low, envs.action_space.high)\n",
     "    action = control_input.reshape(1,4).astype(np.float32)\n",
-    "    # print(action)\n",
     "    thrust_log.append(action.flatten())\n",
     "    obs, reward, terminated, truncated, info = envs.step(action)\n",
     "\n",
     "    state = obs_to_state(obs)\n",
-    "    # print('state:',state)\n",
     "    x_log.append(state[0])\n",
     "    y_log.append(state[2])\n",
     "    z_log.append(state[4])\n",
@@ -377,7 +372,7 @@
     "plt.title(\"position vs Time\")\n",
     "plt.legend()\n",
     "plt.grid()\n",
-    "plt.show()\n"
+    "plt.show()"
    ]
   },
   {

tutorials/compare_sim_and_symbolic.ipynb

@@ -16,7 +16,7 @@
     "from crazyflow.constants import MASS, J\n",
     "from crazyflow.control import Control\n",
     "from crazyflow.sim.physics import Physics\n",
-    "from crazyflow.sim.symbolic import symbolic\n"
+    "from crazyflow.sim.symbolic import symbolic"
    ]
   },
   {
@@ -128,7 +128,6 @@
    "source": [
     "obs, info = envs.reset()\n",
     "state = obs_to_state(obs)\n",
-    "print(state)\n",
     "x_sim_log.append(state[0])\n",
     "y_sim_log.append(state[2])\n",
     "z_sim_log.append(state[4])\n",
@@ -173,26 +172,19 @@
    "source": [
     "MIN_THRUST = 0.028161688\n",
     "MAX_THRUST = 0.14834145\n",
-    "# shape = (1, 4)\n",
     "xf = state\n",
     "\n",
     "for i in range(200):\n",
-    "    # random_input = np.random.uniform(low=low, high=high, size=shape).astype(np.float32)\n",
-    "\n",
     "    random_input = np.array([\n",
     "        np.random.uniform(4 * MIN_THRUST, 4 * MAX_THRUST),  # thrust\n",
     "        np.random.uniform(-np.pi, np.pi),  # \n",
     "        np.random.uniform(-np.pi, np.pi),  # \n",
     "        np.random.uniform(-np.pi, np.pi)   # \n",
     "    ], dtype=np.float32)   \n",
-    "\n",
     "    random_input = random_input.reshape((1, 4))\n",
     "\n",
-    "    print(\"input:\", random_input)\n",
-    "\n",
     "    res = symbolic_model.fd_func(x0=xf, p=random_input)\n",
     "    xf = res[\"xf\"].full().flatten()\n",
-    "    print(\"xf from symbolic:\", xf)\n",
     "    x_sym_log.append(xf[0])\n",
     "    y_sym_log.append(xf[2])\n",
     "    z_sym_log.append(xf[4])\n",
@@ -209,7 +201,6 @@
     "    obs, reward, terminated, truncated, info = envs.step(random_input)\n",
     "    state = obs_to_state(obs)\n",
     "\n",
-    "    print(\"state from sim:\", state)\n",
     "    x_sim_log.append(state[0])\n",
     "    y_sim_log.append(state[2])\n",
     "    z_sim_log.append(state[4])\n",
@@ -222,7 +213,6 @@
     "    roll_dot_sim_log.append(state[9])\n",
     "    pitch_dot_sim_log.append(state[10])\n",
     "    yaw_dot_sim_log.append(state[11])\n",
-    "    envs.render()\n",
     "envs.sim.close()\n",
     "envs.close()"
    ]