Merge branch 'dev'

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,14 +10,10 @@
 from crazyflow.control.control import KF, KM
 
 SYS_ID_PARAMS = {
-    "acc_k1": 20.91,
-    "acc_k2": 3.65,
-    "roll_alpha": -3.96,
-    "roll_beta": 4.08,
-    "pitch_alpha": -6.00,
-    "pitch_beta": 6.21,
-    "yaw_alpha": 0.00,
-    "yaw_beta": 0.00,
+    "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]),
 }
 
 
@@ -60,27 +56,24 @@ def surrogate_identified_collective_wrench(
     rot = R.from_quat(quat)
     thrust = rot.apply(jnp.array([0, 0, collective_thrust]))
     drift = rot.apply(jnp.array([0, 0, 1]))
-    prev_rpy_rates = rot.apply(rpy_rates, inverse=True)
-    a1, a2 = SYS_ID_PARAMS["acc_k1"], SYS_ID_PARAMS["acc_k2"]
-    acc = thrust * a1 + drift * a2
-    # rpy_rates_deriv have no real meaning in this context, since the identified dynamics set the
-    # rpy_rates to the commanded values directly. However, since we use a unified integration
-    # interface for all physics models, we cannot access states directly. Instead, we calculate
-    # which rpy_rates_deriv would have resulted in the desired rpy_rates, and return that.
-    roll_cmd, pitch_cmd, yaw_cmd = attitude
+    rpy_rates_local = rot.apply(rpy_rates, inverse=True)
+    k1, k2 = SYS_ID_PARAMS["acc"]
+    acc = thrust * k1 + drift * k2
     rpy = rot.as_euler("xyz")
-    roll_rate = SYS_ID_PARAMS["roll_alpha"] * rpy[0] + SYS_ID_PARAMS["roll_beta"] * roll_cmd
-    pitch_rate = SYS_ID_PARAMS["pitch_alpha"] * rpy[1] + SYS_ID_PARAMS["pitch_beta"] * pitch_cmd
-    yaw_rate = SYS_ID_PARAMS["yaw_alpha"] * rpy[2] + SYS_ID_PARAMS["yaw_beta"] * yaw_cmd
-    rpy_rates_local = jnp.array([roll_rate, pitch_rate, yaw_rate])
-    rpy_rates_local_deriv = (rpy_rates_local - prev_rpy_rates) / dt
+    k1, k2, k3 = SYS_ID_PARAMS["roll_acc"]
+    roll_rate_deriv = k1 * rpy[0] + k2 * rpy_rates_local[0] + k3 * attitude[0]
+    k1, k2, k3 = SYS_ID_PARAMS["pitch_acc"]
+    pitch_rate_deriv = k1 * rpy[1] + k2 * rpy_rates_local[1] + k3 * attitude[1]
+    k1, k2, k3 = SYS_ID_PARAMS["yaw_acc"]
+    yaw_rate_deriv = k1 * rpy[2] + k2 * rpy_rates_local[2] + k3 * attitude[2]
+    rpy_rates_deriv = jnp.array([roll_rate_deriv, pitch_rate_deriv, yaw_rate_deriv])
     # The identified dynamics model does not use forces or torques, because we assume no knowledge
     # of the drone's mass and inertia. However, to remain compatible with the physics pipeline, we
     # return surrogate forces and torques that result in the desired acceleration and rpy rates
     # derivative. When converting back to the state derivative, the mass and inertia will cancel
     # out, resulting in the correct acceleration and rpy rates derivative regardless of the model's
     # mass and inertia.
-    surrogate_torques = rot.apply(J @ rpy_rates_local_deriv)
+    surrogate_torques = rot.apply(J @ rpy_rates_deriv)
     surrogate_forces = acc * mass
     return surrogate_forces, surrogate_torques
 
@@ -94,7 +87,8 @@ def collective_force2acceleration(force: Array, mass: Array) -> Array:
 @partial(vectorize, signature="(3),(4),(3,3)->(3)")
 def collective_torque2rpy_rates_deriv(torque: Array, quat: Array, J_INV: Array) -> Array:
     """Convert torques to rpy_rates_deriv."""
-    return R.from_quat(quat).apply(J_INV @ torque)
+    rot = R.from_quat(quat)
+    return rot.apply(J_INV @ rot.apply(torque, inverse=True))
 
 
 @partial(vectorize, signature="(4),(4),(3),(3,3)->(3),(3)")

crazyflow/sim/sim.py

@@ -66,37 +66,22 @@ def __init__(
 
         # Initialize MuJoCo world and data
         self._xml_path = xml_path or self.default_path
-        self.spec, self.mj_model, self.mj_data, self.mjx_model, mjx_data = self.setup_mj()
+        self.spec = self.init_mjx_spec()
+        self.mj_model, self.mj_data, self.mjx_model, mjx_data = self.init_mjx_model(self.spec)
         self.viewer: MujocoRenderer | None = None
 
-        # Allocate internal states and controls
-        drone_ids = [self.mj_model.body(f"drone:{i}").id for i in range(n_drones)]
-        self.data = SimData(
-            states=SimState.create(n_worlds, n_drones, self.device),
-            states_deriv=SimStateDeriv.create(n_worlds, n_drones, self.device),
-            controls=SimControls.create(
-                n_worlds, n_drones, state_freq, attitude_freq, thrust_freq, self.device
-            ),
-            params=SimParams.create(n_worlds, n_drones, MASS, J, J_INV, self.device),
-            core=SimCore.create(freq, n_worlds, n_drones, drone_ids, rng_key, self.device),
-            mjx_data=mjx_data,
-            mjx_model=None,
+        self.data, self.default_data = self.init_data(
+            state_freq, attitude_freq, thrust_freq, rng_key, mjx_data
         )
-        if self.n_drones > 1:  # If multiple drones, arrange them in a grid
-            grid = grid_2d(self.n_drones)
-            states = self.data.states.replace(pos=self.data.states.pos.at[..., :2].set(grid))
-            self.data: SimData = self.data.replace(states=states)
-
-        self.data = self.sync_sim2mjx(self.data, self.mjx_model)
-        self.default_data = self.data.replace()  # TODO: Only save the data of one world
 
         # Default functions for the simulation pipeline
         self.disturbance_fn: Callable[[SimData], SimData] | None = None
 
         # Build the simulation pipeline and overwrite the default _step implementation with it
-        self.build()
+        self.init_step_fn()
 
-    def setup_mj(self) -> tuple[Any, Any, Any, Model, Data]:
+    def init_mjx_spec(self) -> mujoco.MjSpec:
+        """Build the MuJoCo model specification for the simulation."""
         assert self._xml_path.exists(), f"Model file {self._xml_path} does not exist"
         spec = mujoco.MjSpec.from_file(str(self._xml_path))
         spec.option.timestep = 1 / self.freq
@@ -110,7 +95,10 @@ def setup_mj(self) -> tuple[Any, Any, Any, Model, Data]:
         for i in range(self.n_drones):
             drone = frame.attach_body(drone_spec.find_body("drone"), "", f":{i}")
             drone.add_freejoint()
-        # Compile and create data structures
+        return spec
+
+    def init_mjx_model(self, spec: mujoco.MjSpec) -> tuple[Any, Any, Model, Data]:
+        """Build the MuJoCo model and data structures for the simulation."""
         mj_model = spec.compile()
         mj_data = mujoco.MjData(mj_model)
         mjx_model = mjx.put_model(mj_model, device=self.device)
@@ -120,9 +108,9 @@ def setup_mj(self) -> tuple[Any, Any, Any, Model, Data]:
         # https://github.com/jax-ml/jax/issues/4274#issuecomment-692406759
         # Tracking issue: https://github.com/google-deepmind/mujoco/issues/2306
         mjx_data = mjx_data.replace(time=jnp.float32(mjx_data.time))
-        return spec, mj_model, mj_data, mjx_model, mjx_data
+        return mj_model, mj_data, mjx_model, mjx_data
 
-    def build(self):
+    def init_step_fn(self):
         """Setup the chain of functions that are called in Sim.step().
 
         We know all the functions that are called in succession since the simulation is configured
@@ -179,6 +167,62 @@ def step(data: SimData, n_steps: int = 1) -> SimData:
 
         self._step = step
 
+    def init_data(
+        self, state_freq: int, attitude_freq: int, thrust_freq: int, rng_key: Array, mjx_data: Data
+    ) -> tuple[SimData, SimData]:
+        """Initialize the simulation data."""
+        drone_ids = [self.mj_model.body(f"drone:{i}").id for i in range(self.n_drones)]
+        N, D = self.n_worlds, self.n_drones
+        data = SimData(
+            states=SimState.create(N, D, self.device),
+            states_deriv=SimStateDeriv.create(N, D, self.device),
+            controls=SimControls.create(N, D, state_freq, attitude_freq, thrust_freq, self.device),
+            params=SimParams.create(N, D, MASS, J, J_INV, self.device),
+            core=SimCore.create(self.freq, N, D, drone_ids, rng_key, self.device),
+            mjx_data=mjx_data,
+            mjx_model=None,
+        )
+        if D > 1:  # If multiple drones, arrange them in a grid
+            grid = grid_2d(D)
+            states = data.states.replace(pos=data.states.pos.at[..., :2].set(grid))
+            data = data.replace(states=states)
+        data = self.sync_sim2mjx(data, self.mjx_model)
+
+        return data, data.replace()  # TODO: Only save the data of one world
+
+    def build(self, mjx: bool = True, data: bool = True, step: bool = True):
+        """Build the simulation pipeline.
+
+        This method is used to (re)build the simulation pipeline after changing the MuJoCo
+        model specification or any of the default functions that are used in the compiled step
+        function.
+
+        Warning:
+            Depending on what you build, you reset the simulation state. For example, rebuilding the
+            simulation data will reset the drone states.
+
+        Args:
+            mjx: Flag to (re)build the MuJoCo model and data structures.
+            data: Flag to (re)build the simulation data.
+            step: Flag to (re)build the simulation step function.
+        """
+        # TODO: Write tests for all options
+        if mjx:
+            if self.viewer is not None:
+                self.viewer.close()
+                self.viewer = None
+            self.mj_model, self.mj_data, self.mjx_model, mjx_data = self.init_mjx_model(self.spec)
+        if data:
+            self.data, self.default_data = self.init_data(
+                self.data.controls.state_freq,
+                self.data.controls.attitude_freq,
+                self.data.controls.thrust_freq,
+                self.data.core.rng_key,
+                self.data.mjx_data if not mjx else mjx_data,
+            )
+        if step:
+            self.init_step_fn()
+
     def reset(self, mask: Array | None = None):
         """Reset the simulation to the initial state.
 
@@ -236,6 +280,7 @@ def seed(self, seed: int):
     def close(self):
         if self.viewer is not None:
             self.viewer.close()
+        self.viewer = None
 
     @property
     def time(self) -> Array:
@@ -299,6 +344,7 @@ def sync_sim2mjx(data: SimData, mjx_model: Model | None = None) -> SimData:
         qvel = rearrange(jnp.concat([vel, local_ang_vel], axis=-1), "w d qvel -> w (d qvel)")
         mjx_data = data.mjx_data
         mjx_model = data.mjx_model if mjx_model is None else mjx_model
+        assert mjx_model is not None, "MuJoCo model is not initialized"
         mjx_data = mjx_data.replace(qpos=qpos, qvel=qvel)
         mjx_data = mjx_kinematics(mjx_model, mjx_data)
         mjx_data = mjx_collision(mjx_model, mjx_data)

crazyflow/sim/symbolic.py

@@ -20,7 +20,7 @@
 from numpy.typing import NDArray
 
 from crazyflow.constants import ARM_LEN, GRAVITY, SIGN_MIX_MATRIX
-from crazyflow.control.control import KF, KM
+from crazyflow.control.control import KF, KM, Control
 from crazyflow.sim import Sim
 
 
@@ -143,9 +143,87 @@ 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.
+    """
+    # # Define states.
+    z = MX.sym("z")
+    z_dot = MX.sym("z_dot")
+    g = GRAVITY
+    # # 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.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]
+    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
+    params_acc = [20.907574256269616, 3.653687545690674]
+    params_roll_rate = [-130.3, -16.33, 119.3]
+    params_pitch_rate = [-99.94, -13.3, 84.73]
+    # 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.
+    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)
+    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_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.
     """
@@ -173,9 +251,6 @@ def symbolic(mass: float, J: NDArray, dt: float) -> SymbolicModel:
     f1, f2, f3, f4 = MX.sym("f1"), MX.sym("f2"), MX.sym("f3"), MX.sym("f4")
     U = cs.vertcat(f1, f2, f3, f4)
 
-    # From Ch. 2 of Luis, Carlos, and Jérôme Le Ny. "Design of a trajectory tracking
-    # controller for a nanoquadcopter." arXiv preprint arXiv:1608.05786 (2016).
-
     # 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.
@@ -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/disturbance.py

@@ -36,7 +36,7 @@ def main(plot: bool = False):
 
     # Second run
     sim.disturbance_fn = disturbance_fn
-    sim.build()
+    sim.build(mjx=False, data=False, step=True)
     pos_disturbed, rpy_disturbed = [], []
     sim.reset()
     for _ in range(3 * sim.control_freq):