diff --git a/examples/quadrotor/plant.py b/examples/quadrotor/plant.py index b38a79d..72985fa 100644 --- a/examples/quadrotor/plant.py +++ b/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() diff --git a/tests/examples/quadrotor/test_plant.py b/tests/examples/quadrotor/test_plant.py index 7f0a5fc..35b4019 100644 --- a/tests/examples/quadrotor/test_plant.py +++ b/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]), + )