Add symbolic thrust model. Improve symbolic_from_sim. Add tests

Author: amacati

crazyflow/sim/__init__.py

@@ -1,5 +1,5 @@
 from crazyflow.sim.physics import Physics
 from crazyflow.sim.sim import Sim
-from crazyflow.sim.symbolic import symbolic
+from crazyflow.sim.symbolic import symbolic_attitude, symbolic_from_sim, symbolic_thrust
 
-__all__ = ["Sim", "Physics", "symbolic"]
+__all__ = ["Sim", "Physics", "symbolic_attitude", "symbolic_from_sim", "symbolic_thrust"]

crazyflow/sim/physics.py

@@ -10,7 +10,7 @@
 from crazyflow.control.control import KF, KM
 
 SYS_ID_PARAMS = {
-    "acc": jnp.array([20.91, 3.65]),
+    "acc": jnp.array([20.907574256269616, 3.653687545690674]),
     "roll_acc": jnp.array([-130.3, -16.33, 119.3]),
     "pitch_acc": jnp.array([-99.94, -13.3, 84.73]),
     "yaw_acc": jnp.array([0.0, 0.0, 0.0]),

crazyflow/sim/symbolic.py

@@ -19,7 +19,8 @@
 from casadi import MX
 from numpy.typing import NDArray
 
-from crazyflow.constants import GRAVITY
+from crazyflow.constants import ARM_LEN, GRAVITY, SIGN_MIX_MATRIX
+from crazyflow.control.control import KF, KM, Control
 from crazyflow.sim import Sim
 
 
@@ -142,9 +143,11 @@ def setup_linearization(self):
         self.loss = cs.Function("loss", l_inputs, l_outputs, l_inputs_str, l_outputs_str)
 
 
-def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
+def symbolic_attitude(dt: float) -> SymbolicModel:
     """Create symbolic (CasADi) models for dynamics, observation, and cost of a quadcopter.
 
+    This model is based on the identified model derived from real-world data of the Crazyflie 2.1.
+
     Returns:
         The CasADi symbolic model of the environment.
     """
@@ -177,13 +180,15 @@ def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
     params_acc = [20.907574256269616, 3.653687545690674]
     params_roll_rate = [-130.3, -16.33, 119.3]
     params_pitch_rate = [-99.94, -13.3, 84.73]
-    # params_yaw_rate = [0, 0, 0], because we always keep yaw as 0 when we identified the parameters.
-    # We introduce a small negative offset here to make sure that we could get result of LQR ect..
-    # TODO: identify params_yaw_rate
+    # The identified model sets params_yaw_rate to [0, 0, 0], because the training data did not
+    # contain any data with yaw != 0. Therefore, it cannot infer the impact of setting the yaw
+    # attitude to a non-zero value on the dynamics. However, using a zero vector will make the
+    # system matrix ill-conditioned for control methods like LQR. Therefore, we introduce a small
+    # spring-like term to the yaw dynamics that leads to a non-singular system matrix.
+    # TODO: identify proper parameters for yaw_rate from real data.
     params_yaw_rate = [-0.01, 0, 0]
 
     # 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])
@@ -213,6 +218,76 @@ def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
     return SymbolicModel(dynamics=dynamics, cost=cost, dt=dt)
 
 
+def symbolic_thrust(mass: float, J: NDArray, dt: float) -> SymbolicModel:
+    """Create symbolic (CasADi) models for dynamics, observation, and cost of a quadcopter.
+
+    This model is based on the analytical model of Luis, Carlos, and Jérôme Le Ny. "Design of a
+    trajectory tracking controller for a nanoquadcopter." arXiv preprint arXiv:1608.05786 (2016).
+
+    Returns:
+        The CasADi symbolic model of the environment.
+    """
+    # Define states.
+    z = MX.sym("z")
+    z_dot = MX.sym("z_dot")
+
+    # Set up the dynamics model for a 3D quadrotor.
+    nx, nu = 12, 4
+    Ixx, Iyy, Izz = J.diagonal()
+    J = cs.blockcat([[Ixx, 0.0, 0.0], [0.0, Iyy, 0.0], [0.0, 0.0, Izz]])
+    Jinv = cs.blockcat([[1.0 / Ixx, 0.0, 0.0], [0.0, 1.0 / Iyy, 0.0], [0.0, 0.0, 1.0 / Izz]])
+    gamma = KM / KF
+    # System state variables
+    x, y = MX.sym("x"), MX.sym("y")
+    x_dot, y_dot = MX.sym("x_dot"), MX.sym("y_dot")
+    phi, theta, psi = MX.sym("phi"), MX.sym("theta"), MX.sym("psi")
+    p, q, r = MX.sym("p"), MX.sym("q"), MX.sym("r")
+    # Rotation matrix transforming a vector in the body frame to the world frame. PyBullet Euler
+    # angles use the SDFormat for rotation matrices.
+    Rob = csRotXYZ(phi, theta, psi)
+    # Define state variables.
+    X = cs.vertcat(x, x_dot, y, y_dot, z, z_dot, phi, theta, psi, p, q, r)
+    # Define inputs.
+    f1, f2, f3, f4 = MX.sym("f1"), MX.sym("f2"), MX.sym("f3"), MX.sym("f4")
+    U = cs.vertcat(f1, f2, f3, f4)
+
+    # Defining the dynamics function.
+    # We are using the velocity of the base wrt to the world frame expressed in the world frame.
+    # Note that the reference expresses this in the body frame.
+    oVdot_cg_o = Rob @ cs.vertcat(0, 0, f1 + f2 + f3 + f4) / mass - cs.vertcat(0, 0, GRAVITY)
+    pos_ddot = oVdot_cg_o
+    pos_dot = cs.vertcat(x_dot, y_dot, z_dot)
+    # We use the spin directions (signs) from the mix matrix used in the simulation.
+    sx, sy, sz = SIGN_MIX_MATRIX[..., 0], SIGN_MIX_MATRIX[..., 1], SIGN_MIX_MATRIX[..., 2]
+    Mb = cs.vertcat(
+        ARM_LEN / cs.sqrt(2.0) * (sx[0] * f1 + sx[1] * f2 + sx[2] * f3 + sx[3] * f4),
+        ARM_LEN / cs.sqrt(2.0) * (sy[0] * f1 + sy[1] * f2 + sy[2] * f3 + sy[3] * f4),
+        gamma * (sz[0] * f1 + sz[1] * f2 + sz[2] * f3 + sz[3] * f4),
+    )
+    rate_dot = Jinv @ (Mb - (cs.skew(cs.vertcat(p, q, r)) @ J @ cs.vertcat(p, q, r)))
+    ang_dot = cs.blockcat(
+        [
+            [1, cs.sin(phi) * cs.tan(theta), cs.cos(phi) * cs.tan(theta)],
+            [0, cs.cos(phi), -cs.sin(phi)],
+            [0, cs.sin(phi) / cs.cos(theta), cs.cos(phi) / cs.cos(theta)],
+        ]
+    ) @ cs.vertcat(p, q, r)
+    X_dot = cs.vertcat(
+        pos_dot[0], pos_ddot[0], pos_dot[1], pos_ddot[1], pos_dot[2], pos_ddot[2], ang_dot, rate_dot
+    )
+
+    Y = cs.vertcat(x, x_dot, y, y_dot, z, z_dot, phi, theta, psi, p, q, r)
+
+    # 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)
+    cost_func = 0.5 * (X - Xr).T @ Q @ (X - Xr) + 0.5 * (U - Ur).T @ R @ (U - Ur)
+    # Define dynamics and cost dictionaries.
+    dynamics = {"dyn_eqn": X_dot, "obs_eqn": Y, "vars": {"X": X, "U": U}}
+    cost = {"cost_func": cost_func, "vars": {"X": X, "U": U, "Xr": Xr, "Ur": Ur, "Q": Q, "R": R}}
+    return SymbolicModel(dynamics=dynamics, cost=cost, dt=dt)
+
+
 def symbolic_from_sim(sim: Sim) -> SymbolicModel:
     """Create a symbolic model from a Sim instance.
 
@@ -226,8 +301,14 @@ def symbolic_from_sim(sim: Sim) -> SymbolicModel:
         The model is expected to deviate from the true dynamics when the sim parameters are
         randomized.
     """
-    mass, J = sim.default_data.params.mass[0, 0], sim.default_data.params.J[0, 0]
-    return symbolic(mass, J, 1 / sim.freq)
+    match sim.control:
+        case Control.attitude:
+            return symbolic_attitude(1 / sim.freq)
+        case Control.thrust:
+            mass, J = sim.default_data.params.mass[0, 0], sim.default_data.params.J[0, 0]
+            return symbolic_thrust(mass, J, 1 / sim.freq)
+        case _:
+            raise ValueError(f"Unsupported control type for symbolic model: {sim.control}")
 
 
 def csRotXYZ(phi: float, theta: float, psi: float) -> MX:

examples/symbolic.py

@@ -1,15 +1,17 @@
 from crazyflow.constants import MASS, J
 from crazyflow.sim import Sim
-from crazyflow.sim.symbolic import symbolic, symbolic_from_sim
+from crazyflow.sim.symbolic import symbolic_attitude, symbolic_from_sim, symbolic_thrust
 
 
 def main():
-    # We can create a symbolic model without the simulation
+    # We can create a symbolic attitude control model without the simulation
     dt = 1 / 500
-    symbolic_model = symbolic(MASS, J, dt)
+    symbolic_model = symbolic_attitude(dt)
+    # We can also create a symbolic thrust control model
+    symbolic_model = symbolic_thrust(MASS, J, dt)
 
     # Or we can create a symbolic model directly from the simulation. Note that this will use the
-    # nominal parameters of the simulation.
+    # nominal parameters of the simulation and choose the control type based on the simulation.
     sim = Sim()
     symbolic_model = symbolic_from_sim(sim)
     assert symbolic_model.nx == 12  # 3 for pos, 3 for orient, 3 for vel, 3 for ang vel

tests/integration/test_symbolic.py

@@ -64,3 +64,48 @@ def test_attitude_symbolic():
     err_msg = "Symbolic and simulation prediction do not match approximately"
     assert np.allclose(x_sym_log, x_sim_log, rtol=1e-2, atol=1e-2), err_msg
     sim.close()
+
+
+@pytest.mark.integration
+def test_thrust_symbolic():
+    sim = Sim(physics="analytical", control="thrust")
+    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)
+
+    rng = np.random.default_rng(seed=42)
+
+    # Run simulation
+    for _ in range(steps):
+        u_rand = rng.uniform(MIN_THRUST, MIN_THRUST, 4)
+        # 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.thrust_control(u_rand.reshape(1, 1, 4))
+        sim.step(sim.freq // sim.control_freq)
+        x_sim_log.append(sim_state2symbolic_state(sim.data.states))
+
+    x_sym_log = np.array(x_sym_log)
+    x_sim_log = np.array(x_sim_log)
+
+    # 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()

tests/unit/test_symbolic.py

@@ -1,15 +1,28 @@
 import pytest
 
 from crazyflow.constants import MASS, J
+from crazyflow.control import Control
 from crazyflow.sim import Sim
-from crazyflow.sim.symbolic import symbolic, symbolic_from_sim
+from crazyflow.sim.symbolic import symbolic_attitude, symbolic_from_sim, symbolic_thrust
 
 
 @pytest.mark.unit
-def test_symbolic_model_creation():
-    """Test creating symbolic model directly."""
+def test_symbolic_attitude_model_creation():
+    """Test creating symbolic attitude model directly."""
     dt = 0.01
-    model = symbolic(mass=MASS, J=J, dt=dt)
+    model = symbolic_attitude(dt)
+
+    assert model.nx == 12  # State dimension
+    assert model.nu == 4  # Input dimension
+    assert model.ny == 12  # Output dimension
+    assert model.dt == dt
+
+
+@pytest.mark.unit
+def test_symbolic_thrust_model_creation():
+    """Test creating symbolic thrust model directly."""
+    dt = 0.01
+    model = symbolic_thrust(mass=MASS, J=J, dt=dt)
 
     assert model.nx == 12  # State dimension
     assert model.nu == 4  # Input dimension
@@ -19,9 +32,10 @@ def test_symbolic_model_creation():
 
 @pytest.mark.unit
 @pytest.mark.parametrize("n_worlds", [1, 2])
-def test_symbolic_from_sim(n_worlds: int):
+@pytest.mark.parametrize("control", [Control.attitude, Control.thrust])
+def test_symbolic_from_sim(n_worlds: int, control: Control):
     """Test creating symbolic model from sim instance."""
-    sim = Sim(n_worlds=n_worlds, n_drones=1)
+    sim = Sim(n_worlds=n_worlds, n_drones=1, control=control)
     model = symbolic_from_sim(sim)
 
     assert model.nx == 12

tutorials/LQR_ILQR.ipynb

@@ -32,10 +32,10 @@
     "from numpy.typing import NDArray\n",
     "from scipy.spatial.transform import Rotation as R\n",
     "\n",
-    "from crazyflow.constants import GRAVITY, MASS, J\n",
+    "from crazyflow.constants import GRAVITY, MASS\n",
     "from crazyflow.control import Control\n",
     "from crazyflow.sim.physics import Physics\n",
-    "from crazyflow.sim.symbolic import symbolic"
+    "from crazyflow.sim.symbolic import symbolic_from_sim"
    ]
   },
   {
@@ -161,7 +161,7 @@
    "outputs": [],
    "source": [
     "dt = 1 / envs.sim.freq\n",
-    "symbolic_model = symbolic(MASS, J, dt)\n",
+    "symbolic_model = symbolic_from_sim(envs.unwrapped.sim)\n",
     "\n",
     "nx = 12 # dimension of state vector\n",
     "nu = 4 # dimension of input vector\n",

tutorials/compare_sim_and_symbolic.ipynb

@@ -13,10 +13,9 @@
     "from numpy.typing import NDArray\n",
     "from scipy.spatial.transform import Rotation as R\n",
     "\n",
-    "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"
+    "from crazyflow.sim.symbolic import symbolic_from_sim"
    ]
   },
   {
@@ -82,7 +81,7 @@
    "outputs": [],
    "source": [
     "dt = 1/envs.sim.freq\n",
-    "symbolic_model = symbolic(MASS, J, dt)"
+    "symbolic_model = symbolic_from_sim(envs.unwrapped.sim)"
    ]
   },
   {