Add the dynamics for QuadrotorPlant. (#75)

This system uses roll-pitch-yaw as the state.

examples/quadrotor/plant.py

@@ -5,7 +5,7 @@
 import jax
 import jax.numpy as jnp
 import pydrake.symbolic as sym
-from pydrake.examples import QuadrotorPlant
+import pydrake.examples
 import pydrake.systems.framework
 import pydrake.geometry
 import pydrake.multibody.plant
@@ -41,6 +41,124 @@ def set_val(i, j, val):
     return R
 
 
+def calc_matrix_relating_rpydt_to_angular_velocity_in_child(
+    rpy: np.ndarray, T
+) -> np.ndarray:
+    """
+    This is the same as CalcMatrixRelatingRpyDtToAngularVelocityInChild in Drake,
+    but supports symbolic computation.
+    """
+    sp = np.sin(rpy[1])
+    cp = np.cos(rpy[1])
+    one_over_cp = 1.0 / cp
+    sr = np.sin(rpy[0])
+    cr = np.cos(rpy[0])
+    cr_over_cp = cr * one_over_cp
+    sr_over_cp = sr * one_over_cp
+    M = np.array(
+        [
+            [T(1), sr_over_cp * sp, cr_over_cp * sp],
+            [T(0), cr, -sr],
+            [T(0), sr_over_cp, cr_over_cp],
+        ]
+    )
+    return M
+
+
+class QuadrotorPlant(pydrake.systems.framework.LeafSystem):
+    """
+    The state is (pos, rpy, pos_dot, omega_WB_B)
+    """
+
+    m: float
+    I: np.ndarray
+    l: float
+    g: float
+    kF: float
+    kM: float
+
+    def __init__(self):
+        super().__init__()
+        self.DeclareVectorInputPort("thrust", 4)
+        state_index = self.DeclareContinuousState(12)
+        self.DeclareStateOutputPort("x", state_index)
+        drake_plant = pydrake.examples.QuadrotorPlant()
+        self.m = drake_plant.m()
+        self.I = drake_plant.inertia()
+        self.g = drake_plant.g()
+        self.l = drake_plant.length()
+        self.kF = drake_plant.force_constant()
+        self.kM = drake_plant.moment_constant()
+        self.I_inv = np.linalg.inv(self.I)
+
+    def DoCalcTimeDerivatives(
+        self,
+        context: pydrake.systems.framework.Context,
+        derivatives: pydrake.systems.framework.ContinuousState,
+    ):
+        x = context.get_continuous_state_vector().CopyToVector()
+        u = self.EvalVectorInput(context, 0).CopyToVector()
+        xdot: np.ndarray = self.dynamics(x, u)
+        derivatives.SetFromVector(xdot)
+
+    def dynamics(self, x: np.ndarray, u: np.ndarray) -> np.ndarray:
+        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 = np.stack((Mx, My, Mz))
+        Fgravity_N = np.array([0, 0, -self.m * self.g])
+
+        rpy = pydrake.math.RollPitchYaw(x[3:6])
+        w_NB_B = x[-3:]
+        rpy_dot = rpy.CalcRpyDtFromAngularVelocityInChild(w_NB_B)
+        R_NB = pydrake.math.RotationMatrix(rpy)
+
+        xyzDDt = (Fgravity_N + R_NB @ Faero_B) / self.m
+
+        wIw = np.cross(w_NB_B, self.I @ w_NB_B)
+        alpha_NB_B = self.I_inv @ (tau_B - wIw)
+
+        xDt = np.concatenate((x[6:9], rpy_dot, xyzDDt, alpha_NB_B))
+
+        return xDt
+
+    def affine_dynamics(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        T = sym.Expression if x.dtype == object else float
+        f = np.zeros((12,), dtype=x.dtype)
+        g = np.zeros((12, 4), dtype=x.dtype)
+        for i in range(12):
+            f[i] = T(0)
+            for j in range(4):
+                g[i, j] = T(0)
+
+        f[:3] = np.array([T(x[i]) for i in range(6, 9)])
+        f[3:6] = (
+            calc_matrix_relating_rpydt_to_angular_velocity_in_child(x[3:6], T=T)
+            @ x[-3:]
+        )
+        rpy = pydrake.math.RollPitchYaw_[T](x[3], x[4], x[5])
+        R = pydrake.math.RotationMatrix_[T](rpy)
+        f[8] = T(-self.g)
+        g[6:9, :] = np.outer(R.matrix()[:, 2], self.kF * np.ones((4,))) / self.m
+        w_NB_B = x[-3:]
+        f[9:] = -self.I_inv @ np.cross(w_NB_B, self.I @ w_NB_B)
+        # tau_B = u_to_M * u
+        u_to_M = np.array(
+            [
+                [T(0), T(self.l * self.kF), T(0), T(-self.l * self.kF)],
+                [T(-self.l * self.kF), T(0), T(self.l * self.kF), T(0)],
+                [T(self.kM), T(-self.kM), T(self.kM), T(-self.kM)],
+            ]
+        )
+        g[9:, :] = self.I_inv @ u_to_M
+        return f, g
+
+
 class QuadrotorPolyPlant(pydrake.systems.framework.LeafSystem):
     """
     The state is x=(quat - [1, 0, 0, 0], pos, pos_dot, omega_WB_B)
@@ -60,7 +178,7 @@ def __init__(self):
         self.DeclareVectorInputPort("thrust", 4)
         state_index = self.DeclareContinuousState(13)
         self.DeclareStateOutputPort("x", state_index)
-        drake_plant = QuadrotorPlant()
+        drake_plant = pydrake.examples.QuadrotorPlant()
         self.m = drake_plant.m()
         self.I = drake_plant.inertia()
         self.g = drake_plant.g()

tests/examples/quadrotor/test_plant.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest  # noqa
+
 import pydrake.symbolic as sym
 import pydrake.math
 
@@ -105,3 +106,61 @@ def check(x_val, u_val):
         u_val = np.array([1, 2, 3, 4.0])
         A, B = dut.linearize_dynamics(x_val, u_val)
         check(x_val, u_val)
+
+
+class TestQuadrotorPlant:
+    def test_dynamics(self):
+        dut = mut.QuadrotorPlant()
+        quadrotor_poly = mut.QuadrotorPolyPlant()
+
+        pos = np.array([0.5, 1.2, -0.4])
+        quat = np.array([0.5, 1.2, 0.3, -0.8])
+        quat = quat / np.linalg.norm(quat)
+        R = pydrake.math.RotationMatrix(
+            pydrake.common.eigen_geometry.Quaternion(quat[0], quat[1], quat[2], quat[3])
+        )
+        rpy = pydrake.math.RollPitchYaw(R)
+        pos_dot = np.array([0.9, -1.1, 0.3])
+        omega_WB_B = np.array([0.5, 2.2, 0.9])
+
+        x = np.concatenate((pos, rpy.vector(), pos_dot, omega_WB_B))
+        u = np.array([0.5, 1.2, 0.9, 2.1])
+        xdot = dut.dynamics(x, u)
+
+        x_poly = np.concatenate(
+            (
+                np.array([quat[0] - 1, quat[1], quat[2], quat[3]]),
+                pos,
+                pos_dot,
+                omega_WB_B,
+            )
+        )
+        x_poly_dot = quadrotor_poly.dynamics(x_poly, u)
+        np.testing.assert_allclose(xdot[:3], x_poly_dot[4:7])
+        np.testing.assert_allclose(xdot[6:], x_poly_dot[7:])
+        np.testing.assert_allclose(
+            xdot[3:6], rpy.CalcRpyDtFromAngularVelocityInChild(omega_WB_B)
+        )
+
+    def test_affine_dynamics(self):
+        dut = mut.QuadrotorPlant()
+
+        x_sym = sym.MakeVectorContinuousVariable(12, "x")
+        f_sym, g_sym = dut.affine_dynamics(x_sym)
+
+        def check(x, u):
+            xdot = dut.dynamics(x, u)
+            f, g = dut.affine_dynamics(x)
+            np.testing.assert_allclose(f + g @ u, xdot, atol=1e-6)
+            env = {x_sym[i]: x[i] for i in range(12)}
+            f_val = np.array([f_sym[i].Evaluate(env) for i in range(12)])
+            g_val = np.array(
+                [[g_sym[i, j].Evaluate(env) for j in range(4)] for i in range(12)]
+            )
+            np.testing.assert_allclose(f_val, f, atol=1e-6)
+            np.testing.assert_allclose(g_val, g, atol=1e-6)
+
+        check(
+            np.array([0.2, 1.2, 0.4, -1.5, 0.4, 0.3, 1.5, 2.3, 4.1, 0.5, -1.7, 1.5]),
+            np.array([0.5, 2.1, 0.4, 1.3]),
+        )