diff --git a/examples/control/dynamics.py b/examples/control/dynamics.py new file mode 100644 index 00000000..98ace94f --- /dev/null +++ b/examples/control/dynamics.py @@ -0,0 +1,120 @@ +import jax.numpy as jnp +import numpy as np + +from crazyflow.control import Control +from crazyflow.control.transform import motor_force2rotor_vel +from crazyflow.sim import Dynamics, Sim + +DURATION = 5.0 +FPS = 60 +TRAJECTORY_CENTER = np.array([0.0, 0.0, 1.0]) +TRAJECTORY_SIZE = np.array([1.0, 0.0, 0.5]) + + +def figure_eight(t: float) -> np.ndarray: + """Return the position, velocity, and acceleration reference at time ``t``.""" + omega = 2.0 * np.pi / DURATION + phase = omega * t + cmd = np.zeros((1, 1, 13)) + cmd[..., 0:3] = TRAJECTORY_CENTER + TRAJECTORY_SIZE * np.array( + [np.sin(phase), 0.0, np.sin(2.0 * phase)] + ) + cmd[..., 3:6] = ( + TRAJECTORY_SIZE * omega * np.array([np.cos(phase), 0.0, 2.0 * np.cos(2.0 * phase)]) + ) + cmd[..., 6:9] = ( + -TRAJECTORY_SIZE * omega**2 * np.array([np.sin(phase), 0.0, 4.0 * np.sin(2.0 * phase)]) + ) + return cmd + + +def reset_sim(sim: Sim) -> None: + """Reset a simulation to the initial reference state and hover rotor speed.""" + sim.reset() + initial_state = figure_eight(0.0) + hover_thrust = sim.data.params.mass * -sim.data.params.gravity_vec[-1] + + if sim.dynamics == Dynamics.first_principles: + motor_forces = jnp.broadcast_to(hover_thrust / 4.0, sim.data.states.rotor_vel.shape) + rotor_vel = motor_force2rotor_vel(motor_forces, sim.data.params.rpm2thrust) + elif sim.dynamics in (Dynamics.so_rpy_rotor, Dynamics.so_rpy_rotor_drag): + # These models store collective thrust in the shared rotor velocity state. + hover_state = (hover_thrust - sim.data.params.acc_coef) / sim.data.params.cmd_f_coef + rotor_vel = jnp.broadcast_to(hover_state, sim.data.states.rotor_vel.shape) + else: + # The rotor state is unused by so_rpy, but keep its reset state nonzero too. + rotor_vel = jnp.broadcast_to(hover_thrust, sim.data.states.rotor_vel.shape) + + states = sim.data.states.replace( + pos=sim.data.states.pos.at[...].set(initial_state[..., 0:3]), + vel=sim.data.states.vel.at[...].set(initial_state[..., 3:6]), + rotor_vel=sim.data.states.rotor_vel.at[...].set(rotor_vel), + ) + sim.data = sim.data.replace(states=states) + + +def main(plot: bool = False, render: bool = False) -> None: + trajectories: dict[Dynamics, np.ndarray] = {} + rmses: dict[Dynamics, float] = {} + reference = np.empty((0, 3)) + + for dynamics in Dynamics: + sim = Sim(dynamics=dynamics, control=Control.state) + reset_sim(sim) + + n_steps = int(DURATION * sim.control_freq) + times = np.arange(n_steps + 1) / sim.control_freq + commands = np.stack([figure_eight(t) for t in times]) + reference = commands[:, 0, 0, 0:3] + trajectory = [np.asarray(sim.data.states.pos[0, 0])] + + for i, cmd in enumerate(commands[1:]): + sim.state_control(cmd) + sim.step(sim.freq // sim.control_freq) + trajectory.append(np.asarray(sim.data.states.pos[0, 0])) + if ((i * FPS) % sim.control_freq) < FPS and render: + sim.render() + + sim.close() + trajectory = np.asarray(trajectory) + trajectories[dynamics] = trajectory + rmses[dynamics] = np.sqrt(np.mean(np.sum((trajectory - reference) ** 2, axis=-1))) + print(f"{dynamics.value}: position RMSE = {rmses[dynamics]:.3f} m") + + if plot: + plot_results(reference, trajectories, rmses) + + +def plot_results( + reference: np.ndarray, trajectories: dict[Dynamics, np.ndarray], rmses: dict[Dynamics, float] +) -> None: + """Plot the x-z trajectories and lateral positions against the reference.""" + # Only import if plotting is desired to avoid a dependency on matplotlib. + import matplotlib.pyplot as plt + + fig, (ax_xz, ax_y) = plt.subplots(1, 2, figsize=(12, 5), layout="constrained") + ax_xz.plot(reference[:, 0], reference[:, 2], "k--", linewidth=2, label="reference") + times = np.linspace(0.0, DURATION, len(reference)) + ax_y.plot(times, reference[:, 1], "k--", linewidth=2, label="reference") + for dynamics, trajectory in trajectories.items(): + label = f"{dynamics.value} (RMSE={rmses[dynamics] * 1_000:.0f}mm)" + ax_xz.plot(trajectory[:, 0], trajectory[:, 2], label=label) + ax_y.plot(times, trajectory[:, 1], label=label) + + ax_xz.set_title("Figure-eight trajectory") + ax_xz.set_xlabel("x (m)") + ax_xz.set_ylabel("z (m)") + ax_xz.set_aspect("equal", adjustable="datalim") + ax_xz.legend() + ax_y.set_title("Out-of-plane position") + ax_y.set_xlabel("time (s)") + ax_y.set_ylabel("y (m)") + for ax in (ax_xz, ax_y): + ax.set_box_aspect(1) + ax.grid() + fig.suptitle("Figure-eight tracking across dynamics models") + plt.show() + + +if __name__ == "__main__": + main(plot=True, render=False)