diff --git a/examples/plugins/ground_effect.py b/examples/plugins/ground_effect.py new file mode 100644 index 00000000..e0c830d9 --- /dev/null +++ b/examples/plugins/ground_effect.py @@ -0,0 +1,114 @@ +"""External-force implementation of the ground effect in Eq. (15) of Shi et al arXiv:1811.08027v.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +os.environ["SCIPY_ARRAY_API"] = "1" + +import jax.numpy as jnp +import numpy as np +from jax.scipy.spatial.transform import Rotation as R + +from crazyflow.sim import Sim +from crazyflow.sim.pipeline import insert_fn_before + +if TYPE_CHECKING: + from crazyflow.sim.data import SimData + + +# Parameters for cf21B_500 +PROPELLER_DIAMETER = 55e-3 # m +MU = 2.0 +MIN_HEIGHT = 0.02 # m; Eq. (15) is not valid arbitrarily close to the floor +MAX_GAIN = 2.0 # avoid the model's singularity near the floor + +# Descent points +HOVER_HEIGHTS = np.linspace(0.50, 0.02, 15) +SETTLE_DURATION = 10.0 # s +SAMPLE_DURATION = 0.2 # s + + +def ground_effect_fn(data: SimData) -> SimData: + rpm = data.states.rotor_vel + c, b, a = ( + data.params.rpm2thrust[..., 0], + data.params.rpm2thrust[..., 1], + data.params.rpm2thrust[..., 2], + ) + nominal_thrust = jnp.sum(c + b * rpm + a * rpm**2, axis=-1) + + height = jnp.maximum(data.states.pos[..., 2], MIN_HEIGHT) + gain = 1.0 / (1.0 - MU * (PROPELLER_DIAMETER / (8.0 * height)) ** 2) + gain = jnp.minimum(gain, MAX_GAIN) + extra_thrust = nominal_thrust * (gain - 1.0) + + # The force follows the body z thrust axis; ``states.force`` needs world coordinates. + body_force = jnp.zeros_like(data.states.pos).at[..., 2].set(extra_thrust) + ground_force = R.from_quat(data.states.quat).apply(body_force) + return data.replace(states=data.states.replace(force=ground_force)) + + +def measure_hover_points( + sim: Sim, heights: np.ndarray, render: bool = False +) -> tuple[np.ndarray, np.ndarray]: + """Hold at each setpoint and return the mean measured height and thrust command.""" + command = np.zeros((sim.n_worlds, sim.n_drones, 16)) + command[..., 9:13] = R.from_euler("z", 0.0).as_quat() + + settle_steps = int(SETTLE_DURATION * sim.control_freq) + total_steps = settle_steps + int(SAMPLE_DURATION * sim.control_freq) + fps = 60 + + hover_heights, hover_thrusts = [], [] + + for height in heights: + command[..., 2] = height + height_samples = [] + thrust_samples = [] + for step in range(total_steps): + sim.state_control(command) + sim.step(sim.freq // sim.control_freq) + + if step >= settle_steps: + height_samples.append(float(sim.data.states.pos[0, 0, 2])) + # This is the collective force command passed to the motor mixer. + thrust_samples.append(float(sim.data.controls.force_torque.cmd[0, 0, 0])) + if render: + if ((step * fps) % sim.control_freq) < fps: + sim.render() + + hover_heights.append(np.mean(height_samples)) + hover_thrusts.append(np.mean(thrust_samples)) + + return np.asarray(hover_heights), np.asarray(hover_thrusts) + + +def main(plot: bool = True, render: bool = False) -> None: + + sim = Sim(n_drones=1, drone="cf21B_500", control="state") + + insert_fn_before(sim.step_pipeline, "integration", ground_effect_fn) + sim.build_step_fn() + + sim.data = sim.data.replace( + states=sim.data.states.replace(pos=jnp.array([[[0.0, 0.0, HOVER_HEIGHTS[0]]]])) + ) + + hover_heights, hover_thrusts = measure_hover_points(sim, HOVER_HEIGHTS, render=render) + + sim.close() + + if plot: + import matplotlib.pyplot as plt + + plt.scatter(hover_heights, hover_thrusts) + plt.xlabel("Measured hover height (m)") + plt.ylabel("Commanded collective thrust (N)") + plt.grid() + plt.show() + + +if __name__ == "__main__": + main(render=True)