Plugin examples¶
Extending the step and reset pipelines with custom stages for domain randomization, disturbances, action delays, aerodynamic effects, and additional outputs. See Pipelines for how stages are composed.
Domain randomization¶
Randomize the drone parameters in the reset pipeline.
Mass, inertia and drag are randomized per drone, and thrust and torque curves, rotor dynamics, arm lengths and propeller inertias per motor. Each element of the default parameters is scaled by an independent uniform factor in [1 - x, 1 + x]. Using the default parameters as the base value ensures that repeated resets do not compound. An optional mask limits randomization to selected worlds.
examples/plugins/randomize.py
"""Randomize the drone parameters in the reset pipeline.
Mass, inertia and drag are randomized per drone, and thrust and torque curves, rotor dynamics, arm
lengths and propeller inertias per motor. Each element of the default parameters is scaled by an
independent uniform factor in [1 - x, 1 + x]. Using the default parameters as the base value
ensures that repeated resets do not compound. An optional mask limits randomization to selected
worlds.
"""
import os
os.environ["SCIPY_ARRAY_API"] = "1"
import jax
import jax.numpy as jnp
import numpy as np
from jax import Array
from scipy.spatial.transform import Rotation as R
from crazyflow.control import Control
from crazyflow.sim import Sim
from crazyflow.sim.data import SimData
from crazyflow.sim.pipeline import append_fn
from crazyflow.utils import grid_2d, leaf_replace
@jax.jit
def randomize_mass(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
key, mass_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key)) # Make sure to update the rng_key
# The default mass (1,) is shared by all drones. Scaling it to (n_worlds, n_drones, 1) gives
# every drone its own mass.
shape = (data.core.n_worlds, data.core.n_drones, 1)
amount = 0.1 # 10% variation
scale = jax.random.uniform(mass_key, shape, minval=1 - amount, maxval=1 + amount)
mass = default_data.params.mass * scale
return data.replace(params=leaf_replace(data.params, mask, mass=mass))
@jax.jit
def randomize_inertia(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
key, inertia_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
shape = (data.core.n_worlds, data.core.n_drones, 3, 3) # Randomize J across worlds
amount = 0.1
scale = jax.random.uniform(inertia_key, shape, minval=1 - amount, maxval=1 + amount)
J = default_data.params.J * scale
return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J)))
@jax.jit
def randomize_thrust_curve(
data: SimData, default_data: SimData, mask: Array | None = None
) -> SimData:
key, thrust_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
# The default thrust curve coefficients are shared by all drones and motors with shape (1, 3).
# Multiplying them with a (n_worlds, n_drones, 4, 3) scale gives every motor its own curve.
shape = (data.core.n_worlds, data.core.n_drones, 4, 3)
amount = 0.05
scale = jax.random.uniform(thrust_key, shape, minval=1 - amount, maxval=1 + amount)
rpm2thrust = default_data.params.rpm2thrust * scale
return data.replace(params=leaf_replace(data.params, mask, rpm2thrust=rpm2thrust))
@jax.jit
def randomize_torque_curve(
data: SimData, default_data: SimData, mask: Array | None = None
) -> SimData:
key, torque_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
shape = (data.core.n_worlds, data.core.n_drones, 4, 3) # 3 per motor, so (N, M, 4, 3)
amount = 0.1
scale = jax.random.uniform(torque_key, shape, minval=1 - amount, maxval=1 + amount)
rpm2torque = default_data.params.rpm2torque * scale
return data.replace(params=leaf_replace(data.params, mask, rpm2torque=rpm2torque))
@jax.jit
def randomize_rotor_dynamics(
data: SimData, default_data: SimData, mask: Array | None = None
) -> SimData:
key, rotor_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
shape = (data.core.n_worlds, data.core.n_drones, 4, 4)
amount = 0.05
scale = jax.random.uniform(rotor_key, shape, minval=1 - amount, maxval=1 + amount)
rotor_dyn_coef = default_data.params.rotor_dyn_coef * scale
return data.replace(params=leaf_replace(data.params, mask, rotor_dyn_coef=rotor_dyn_coef))
@jax.jit
def randomize_prop_inertia(
data: SimData, default_data: SimData, mask: Array | None = None
) -> SimData:
key, prop_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
# Default propeller inertia is shared by all drones and motors with shape (1,). Using
# (n_worlds, n_drones, 4) gives every propeller its own inertia.
shape = (data.core.n_worlds, data.core.n_drones, 4)
amount = 0.2
scale = jax.random.uniform(prop_key, shape, minval=1 - amount, maxval=1 + amount)
prop_inertia = default_data.params.prop_inertia * scale
return data.replace(params=leaf_replace(data.params, mask, prop_inertia=prop_inertia))
@jax.jit
def randomize_arm_length(
data: SimData, default_data: SimData, mask: Array | None = None
) -> SimData:
key, arm_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
shape = (data.core.n_worlds, data.core.n_drones, 4) # Same as for propeller inertia
amount = 0.01
scale = jax.random.uniform(arm_key, shape, minval=1 - amount, maxval=1 + amount)
L = default_data.params.L * scale
return data.replace(params=leaf_replace(data.params, mask, L=L))
@jax.jit
def randomize_drag(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
key, drag_key = jax.random.split(data.core.rng_key)
data = data.replace(core=data.core.replace(rng_key=key))
shape = (data.core.n_worlds, data.core.n_drones, 3, 3)
amount = 0.3
scale = jax.random.uniform(drag_key, shape, minval=1 - amount, maxval=1 + amount)
drag_matrix = default_data.params.drag_matrix * scale
return data.replace(params=leaf_replace(data.params, mask, drag_matrix=drag_matrix))
def main():
sim = Sim(n_worlds=3, n_drones=4, control=Control.state)
append_fn(sim.reset_pipeline, randomize_mass)
append_fn(sim.reset_pipeline, randomize_inertia)
append_fn(sim.reset_pipeline, randomize_thrust_curve)
append_fn(sim.reset_pipeline, randomize_torque_curve)
append_fn(sim.reset_pipeline, randomize_rotor_dynamics)
append_fn(sim.reset_pipeline, randomize_prop_inertia)
append_fn(sim.reset_pipeline, randomize_arm_length)
append_fn(sim.reset_pipeline, randomize_drag)
sim.build_reset_fn()
mask = np.array([True, False, False]) # Only randomize the first world
duration = 5.0
fps = 60
for _ in range(3):
cmd = np.zeros((sim.n_worlds, sim.n_drones, 16))
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
cmd[..., 2] = 0.4
cmd[..., :2] = grid_2d(sim.n_drones) * 0.25
# After the first reset, each drone should behave slightly differently
for i in range(int(duration * sim.control_freq)):
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
# Note: The mask is optional. We can also randomize all worlds at once by not passing it
sim.reset(mask=mask) # Only reset the first world, the other two will stay the same
sim.close()
if __name__ == "__main__":
main()
Disturbance injection¶
Inject a random external force and torque into the step pipeline.
The disturbance is applied on every dynamics tick, so the drone fights wind-like perturbations.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/plugins/disturbance.py
"""Inject a random external force and torque into the step pipeline.
The disturbance is applied on every dynamics tick, so the drone fights wind-like perturbations.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
import argparse
import os
import jax
import numpy as np
from numpy.typing import NDArray
from crazyflow.sim import Sim
from crazyflow.sim.data import SimData
from crazyflow.sim.pipeline import insert_fn_before
os.environ["SCIPY_ARRAY_API"] = "1"
from scipy.spatial.transform import Rotation as R
def disturbance_fn(data: SimData) -> SimData:
key, subkey = jax.random.split(data.core.rng_key)
states = data.states
disturbance_force = jax.random.normal(subkey, states.force.shape) * 0.2 # N, world frame
states = states.replace(force=disturbance_force)
key, subkey = jax.random.split(key)
disturbance_torque = jax.random.normal(subkey, states.torque.shape) * 0.0002 # Nm, world frame
states = states.replace(torque=disturbance_torque)
return data.replace(states=states, core=data.core.replace(rng_key=key))
def main(plot: bool = False):
sim = Sim(control="state")
control = np.zeros((sim.n_worlds, sim.n_drones, 16))
control[..., 9:13] = R.from_euler("z", 0.0).as_quat()
control[..., :3] = 0.2
# First run
pos, quat = [], []
sim.reset()
for _ in range(3 * sim.control_freq):
sim.state_control(control)
sim.step(sim.freq // sim.control_freq)
pos.append(sim.data.states.pos[0, 0])
quat.append(sim.data.states.quat[0, 0])
sim.render()
# Second run
# We insert the disturbance function into the step pipeline before the integration step. You can
# inspect the step pipeline stages with
# print(sim.step_pipeline)
insert_fn_before(sim.step_pipeline, "integration", disturbance_fn)
sim.build_step_fn()
pos_disturbed, quat_disturbed = [], []
sim.reset()
for _ in range(3 * sim.control_freq):
sim.state_control(control)
sim.step(sim.freq // sim.control_freq)
pos_disturbed.append(sim.data.states.pos[0, 0])
quat_disturbed.append(sim.data.states.quat[0, 0])
sim.render()
sim.close()
if plot:
plot_results(pos, pos_disturbed, quat, quat_disturbed)
def plot_results(
pos: list[NDArray],
pos_disturbed: list[NDArray],
quat: list[NDArray],
quat_disturbed: list[NDArray],
):
# Only import if plotting is desired to avoid a dependency on matplotlib
import matplotlib.pyplot as plt # noqa: F401
pos, pos_disturbed = np.array(pos), np.array(pos_disturbed)
rpy = R.from_quat(quat).as_euler("xyz")
rpy_disturbed = R.from_quat(quat_disturbed).as_euler("xyz")
fig, ax = plt.subplots(3, 2, sharex="all", figsize=(10, 6))
t = np.linspace(0, 3, len(pos))
# XYZ position
ax[0, 0].plot(t, pos[:, 0], label="x undisturbed", color="r")
ax[0, 0].plot(t, pos_disturbed[:, 0], label="x disturbed", color="r", linestyle="--")
ax[1, 0].plot(t, pos[:, 1], label="y undisturbed", color="g")
ax[1, 0].plot(t, pos_disturbed[:, 1], label="y perturbed", color="g", linestyle="--")
ax[2, 0].plot(t, pos[:, 2], label="z undisturbed", color="b")
ax[2, 0].plot(t, pos_disturbed[:, 2], label="z disturbed", color="b", linestyle="--")
# RPY angles
ax[0, 1].plot(t, rpy[:, 0], label="roll undisturbed", color="r")
ax[0, 1].plot(t, rpy_disturbed[:, 0], label="roll disturbed", color="r", linestyle="--")
ax[1, 1].plot(t, rpy[:, 1], label="pitch undisturbed", color="g")
ax[1, 1].plot(t, rpy_disturbed[:, 1], label="pitch disturbed", color="g", linestyle="--")
ax[2, 1].plot(t, rpy[:, 2], label="yaw undisturbed", color="b")
ax[2, 1].plot(t, rpy_disturbed[:, 2], label="yaw disturbed", color="b", linestyle="--")
fig.suptitle("Dynamics with disturbance")
ax[2, 0].set_xlabel("Time (s)")
ax[2, 1].set_xlabel("Time (s)")
for _ax in ax.flatten():
_ax.legend()
plt.tight_layout()
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot)
Action delay¶
Delay the attitude commands by 0.03 s with a plugin.
The plugin queues the staged attitude commands in the plugins dict and is prepended to the step pipeline. To demonstrate the effect, we fly the same trajectory with and without action delay and plot the difference between the two.
examples/plugins/action_delay.py
"""Delay the attitude commands by 0.03 s with a plugin.
The plugin queues the staged attitude commands in the plugins dict and is prepended to the step
pipeline. To demonstrate the effect, we fly the same trajectory with and without action delay and
plot the difference between the two.
"""
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 scipy.spatial.transform import Rotation as R
from crazyflow import Sim
from crazyflow.sim.pipeline import prepend_fn
from crazyflow.sim.visualize import draw_points
if TYPE_CHECKING:
from crazyflow.sim.data import SimData
def trajectory(t: float) -> np.ndarray:
cmd = np.zeros((1, 1, 16))
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
cmd[..., :3] = [np.cos(t) - 1, np.sin(t), 0.2 * t]
return cmd
def action_delay(data: SimData) -> SimData:
"""Delay state control actions."""
queued_actions = data.plugins["queued_actions"]
next_action = queued_actions[0]
queued_actions = jnp.roll(queued_actions, shift=-1, axis=0)
queued_actions = queued_actions.at[-1].set(data.controls.attitude.staged_cmd)
data = data.replace(
controls=data.controls.replace(
attitude=data.controls.attitude.replace(staged_cmd=next_action)
),
plugins=data.plugins | {"queued_actions": queued_actions},
)
return data
def main():
sim = Sim(control="state")
states = []
steps = 500
for i in range(steps):
cmd = trajectory(i / sim.control_freq)
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)
sim.render(camera="track_cam:0")
states.append(sim.data.states.pos[0, 0])
# Delay settings
delay: float = 0.03 # seconds
delay_steps = int(delay * sim.data.controls.attitude.freq)
sim.reset()
# Now we add our action delay into the simulation. We first insert the data we need into the
# plugins dict, then add our plugin function into the step pipeline, and finally rebuild the sim
# default data and step function to make sure our plugin is included and data persists across
# resets.
custom_data = {"queued_actions": jnp.zeros((delay_steps, 1, 1, 4))}
sim.data = sim.data.replace(plugins=sim.data.plugins | custom_data)
prepend_fn(sim.step_pipeline, action_delay, name="action_delay")
sim.build_default_data()
sim.build_step_fn()
# Run the simulation again, this time with an action delay. The states should be different
delayed_states = []
for i in range(steps):
cmd = trajectory(i / sim.control_freq)
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)
draw_points(sim, states[i], size=0.02)
sim.render()
delayed_states.append(sim.data.states.pos[0, 0])
position_differences = jnp.linalg.norm(jnp.array(states) - jnp.array(delayed_states), axis=-1)
print(f"Mean position difference: {jnp.mean(position_differences) * 100:.2f}cm")
sim.close()
if __name__ == "__main__":
main()
Downwash¶
Model the downwash between identical Crazyflies with a far-field external-wrench plugin.
The plugin uses the far-field jet from [1] and the thrust-decay model of [2]. It samples the flow at each target rotor and center of mass, computes rotor thrust loss and drag, and writes the resulting world-frame external force and torque.
The demo holds one cf21B_500 drone in hover while a second makes three passes underneath at
different speeds and heights.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
[1] Bauersfeld et al. https://arxiv.org/abs/2403.13321 [2] Su et al. https://arxiv.org/abs/2207.09645
examples/plugins/downwash.py
"""Model the downwash between identical Crazyflies with a far-field external-wrench plugin.
The plugin uses the far-field jet from [1] and the thrust-decay model of [2]. It samples the flow
at each target rotor and center of mass, computes rotor thrust loss and drag, and writes the
resulting world-frame external force and torque.
The demo holds one `cf21B_500` drone in hover while a second makes three passes underneath at
different speeds and heights.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
[1] Bauersfeld et al. https://arxiv.org/abs/2403.13321
[2] Su et al. https://arxiv.org/abs/2207.09645
"""
from __future__ import annotations
import argparse
from typing import TYPE_CHECKING
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 jax import Array
from crazyflow.sim.data import SimData
# Physical parameters for the cf21B_500
AIR_DENSITY = 1.225 # kg/m^3
PROPELLER_RADIUS = 27.5e-3 # m
MOTOR_DISTANCE = 0.1 # m, distance between opposite motors
# This must be fitted for the propeller/downwash setup.
THRUST_DECAY_COEFFICIENT = 0.07 # s/m
# Far-field fit in Eq. (9) of [1]
BD = 10.11
S = 0.07668
S0 = -5.817
def downwash_speed(s: Array, r: Array, u_hover: Array) -> Array:
"""Return downwash speed at axial and radial distances in metres.
Positive s points downstream along the source drone's negative body z-axis.
Inputs broadcast to the sampling grid; the returned speeds are in m/s.
"""
# Normalization according to [1] Eq. (8).
s_normalized = s / MOTOR_DISTANCE
r_normalized = r / MOTOR_DISTANCE
# Keep the fit finite upstream, where its contribution is masked below.
axial_distance = jnp.maximum(s_normalized - S0, 1e-6)
half_width = S * axial_distance # [1] Eq. (6)
centerline_speed = u_hover * BD / axial_distance # [1] Eq. (5)
radial_ratio = r_normalized / half_width # [1] Eq. (4)
# [1] Eq. (3).
speed = centerline_speed / (1.0 + (jnp.sqrt(2.0) - 1.0) * radial_ratio**2) ** 2
return jnp.where(s_normalized > 0.1, speed, 0.0)
def downwash_fn(data: SimData) -> SimData:
"""Apply downwash-induced thrust loss and drag as a world-frame external wrench.
The source flow originates at each drone centre, while the field is sampled
at every target rotor and CoM in the source's body frame.
"""
body_to_world = R.from_quat(data.states.quat)
mixing_matrix = data.params.mixing_matrix
drag_matrix = data.params.drag_matrix
offsets = data.params.L * jnp.stack(
[-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], axis=-1
)
rotor_offsets_world = R.from_quat(data.states.quat[..., None, :]).apply(offsets)
rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
sample_positions = jnp.concatenate([rotor_positions, data.states.pos[..., None, :]], axis=2)
# Axis 1 indexes the source, axis 2 the target, and axis 3 its rotors then CoM.
source_to_target = data.states.pos[:, :, None, None, :] - sample_positions[:, None, :, :, :]
# Broadcast each source rotation across all target drones and sampling points.
source_to_target_body = R.from_quat(data.states.quat[..., None, None, :]).apply(
source_to_target, inverse=True
)
s = source_to_target_body[..., 2]
r = jnp.linalg.vector_norm(source_to_target_body[..., :2], axis=-1)
mass = data.params.mass[0]
gravity = -data.params.gravity_vec[2]
n_propellers = mixing_matrix.shape[-1]
u_hover = jnp.sqrt(
mass * gravity / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * n_propellers)
) # [1] Eq. (1)
# Shape: (world, source, target, sample), with rotors followed by the CoM.
sample_speed = downwash_speed(s, r, u_hover)
z_axes = body_to_world.as_matrix()[..., 2]
# The final sample is the CoM; each source's wind follows its negative z-axis.
wind_com_world = jnp.sum(-sample_speed[..., -1, None] * z_axes[:, :, None, :], axis=1)
wind_com_body = body_to_world.apply(wind_com_world, inverse=True)
# Project only the rotor samples onto the target axis for thrust loss.
cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
# Sum all sources at each target rotor: (world, target, rotor).
rotor_inflow = jnp.sum(sample_speed[..., :-1] * cos_theta[..., None], axis=1)
# [2] Eq. (5): each motor loses a fraction b_v * U_D of its current thrust
loss_fraction = THRUST_DECAY_COEFFICIENT * rotor_inflow
rotor_vel = data.states.rotor_vel
k0, k1, k2 = (
data.params.rpm2thrust[..., 0],
data.params.rpm2thrust[..., 1],
data.params.rpm2thrust[..., 2],
)
motor_thrust = k0 + k1 * rotor_vel + k2 * rotor_vel**2
thrust_delta = -loss_fraction * motor_thrust
# Map the per-motor force changes to a body-frame wrench, as in [2] Eq. (7).
total_thrust_delta = jnp.sum(thrust_delta, axis=-1)
zeros = jnp.zeros_like(total_thrust_delta)
force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
# Compute drag induced through downwash
drag_body = (-drag_matrix @ wind_com_body[..., None])[..., 0]
force_body += drag_body
# Compute the torque generated by downwash
lever = jnp.array([1.0, 1.0, 0.0])
torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[..., 0] * lever
states = data.states.replace(
force=body_to_world.apply(force_body), torque=body_to_world.apply(torque_body)
)
return data.replace(states=states)
def plot_hover_velocity_field(source_positions: np.ndarray, data: SimData) -> None:
"""Plot the far-field downwash-speed magnitude in the y=0 plane."""
import matplotlib.pyplot as plt
x = np.linspace(-0.6, 0.6, 300)
z = np.linspace(0.0, 1.15, 300)
X, Z = np.meshgrid(x, z)
# Every grid point lies in the y=0 plane.
points = np.stack((X, np.zeros_like(X), Z), axis=-1)
u_downwash = np.zeros_like(X)
gravity = -data.params.gravity_vec[2]
n_propellers = data.params.mixing_matrix.shape[-1]
mass = data.params.mass[0]
u_hover = np.sqrt(
mass * gravity / (2.0 * AIR_DENSITY * np.pi * PROPELLER_RADIUS**2 * n_propellers)
)
for source_pos in source_positions:
source_to_point = source_pos - points
s = source_to_point[..., 2]
r = np.linalg.vector_norm(source_to_point[..., :2], axis=-1)
u_downwash += downwash_speed(jnp.asarray(s), jnp.asarray(r), u_hover)
fig, ax = plt.subplots(figsize=(6, 4.5), layout="constrained")
image = ax.pcolormesh(X, Z, u_downwash, shading="auto", cmap="viridis")
ax.scatter(source_positions[:, 0], source_positions[:, 2], color="red", label="source drone")
ax.set_xlabel("x (m)", fontsize=14)
ax.set_ylabel("z (m)", fontsize=14)
ax.tick_params(axis="both", labelsize=12)
ax.set_title("Hovering-drone downwash speed", fontsize=14)
ax.legend()
colorbar = fig.colorbar(image, ax=ax, pad=0.02)
colorbar.set_label("downward airspeed $U_D$ (m/s)", fontsize=14)
colorbar.ax.tick_params(labelsize=12)
plt.show()
def main(plot: bool = False) -> None:
"""Hover drone 0 while drone 1 makes three downwash passes at different height and velocity."""
sim = Sim(n_drones=2, drone="cf21B_500", control="state")
insert_fn_before(sim.step_pipeline, "integration", downwash_fn)
sim.build_step_fn()
upper_pos = np.array([0.0, 0.0, 1.2])
first_run_height = 0.5
second_run_height = 0.95
lower_start = np.array([-0.5, 0.0, first_run_height])
sim.data = sim.data.replace(
states=sim.data.states.replace(pos=jnp.array([[upper_pos, lower_start]]))
)
sim.build_default_data()
command = np.zeros((1, 2, 16))
command[..., 9:13] = R.from_euler("z", 0).as_quat()
command[0, 0, :3] = upper_pos
waypoints = np.concatenate(
(
# First pass slow
np.linspace(
lower_start, [0.3, 0.0, first_run_height], 3 * sim.control_freq, endpoint=False
),
# Pause to stabilize
np.tile([0.3, 0.0, first_run_height], (2 * sim.control_freq, 1)),
# Second pass fast
np.linspace([0.3, 0.0, first_run_height], lower_start, int(0.5 * sim.control_freq)),
# Pause to stabilize
np.tile(lower_start, (2 * sim.control_freq, 1)),
# Increase altitude
np.linspace(
lower_start, [-0.5, 0.0, second_run_height], sim.control_freq, endpoint=False
),
# Third pass closer to hovering drone
np.linspace(
[-0.5, 0.0, second_run_height], [0.5, 0.0, second_run_height], 3 * sim.control_freq
),
)
)
z_positions = []
downwash_force_z = []
downwash_pitch_torque = []
for position in waypoints:
command[0, 1, :3] = position
sim.state_control(command)
sim.step(sim.freq // sim.control_freq)
z_positions.append(sim.data.states.pos[0, :, 2])
downwash_force_z.append(sim.data.states.force[0, 1, 2])
downwash_pitch_torque.append(sim.data.states.torque[0, 1, 1])
sim.render()
sim.close()
if plot:
plot_results(
z_positions,
downwash_force_z,
downwash_pitch_torque,
upper_pos,
sim.data,
sim.control_freq,
)
def plot_results(
z_positions: list[Array],
downwash_force_z: list[Array],
downwash_pitch_torque: list[Array],
upper_pos: np.ndarray,
data: SimData,
control_freq: int,
) -> None:
"""Plot the passes of the lower drone and the downwash field of the hovering drone."""
import matplotlib.pyplot as plt
t = np.arange(len(z_positions)) / control_freq
z_positions = np.asarray(z_positions)
fig, axes = plt.subplots(3, 1, sharex=True)
axes[0].plot(t, z_positions[:, 0], label="upper drone")
axes[0].plot(t, z_positions[:, 1], label="lower drone")
axes[0].set_ylabel("z position (m)")
axes[0].legend()
axes[1].plot(t, downwash_force_z, label="lower drone")
axes[1].set_ylabel("downwash force z (N)")
axes[1].legend()
axes[2].plot(t, downwash_pitch_torque, label="lower drone")
axes[2].set_xlabel("time (s)")
axes[2].set_ylabel("downwash pitch torque y (Nm)")
axes[2].legend()
plot_hover_velocity_field(upper_pos[None], data)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot)
Ground effect¶
Model the ground effect of a cf21B_500 drone as an external force.
The plugin implements Eq. (15) of Shi et al., arXiv:1811.08027. It computes the nominal thrust from the rotor speeds, applies a height-dependent gain, and adds the extra thrust along the drone's body z-axis as a world-frame external force.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/plugins/ground_effect.py
"""Model the ground effect of a `cf21B_500` drone as an external force.
The plugin implements Eq. (15) of Shi et al., arXiv:1811.08027. It computes the nominal thrust
from the rotor speeds, applies a height-dependent gain, and adds the extra thrust along the drone's
body z-axis as a world-frame external force.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
from __future__ import annotations
import argparse
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[list[float], list[float]]:
"""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 hover_heights, hover_thrusts
def main(plot: bool = False, 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:
plot_results(hover_heights, hover_thrusts)
def plot_results(hover_heights: list[float], hover_thrusts: list[float]):
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__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot, render=True)
State derivatives¶
Add the state derivatives to the step pipeline.
We compare two methods for the state derivatives: evaluating the dynamics and finite differences.
The dynamics need the current state and input, so we evaluate them before integration. At the end
of the step, they are the derivative at the previous state. Finite differences are taken after
integration and give the derivative from the previous to the current state. For Euler integration,
both methods are identical up to numerical noise. Other integrators differ, as shown here for RK4.
Both plugins store a SimStateDeriv in the plugins dict. The plot compares them while the drone
flies a figure-eight.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/plugins/derivatives.py
"""Add the state derivatives to the step pipeline.
We compare two methods for the state derivatives: evaluating the dynamics and finite differences.
The dynamics need the current state and input, so we evaluate them before integration. At the end
of the step, they are the derivative at the previous state. Finite differences are taken after
integration and give the derivative from the previous to the current state. For Euler integration,
both methods are identical up to numerical noise. Other integrators differ, as shown here for RK4.
Both plugins store a `SimStateDeriv` in the plugins dict. The plot compares them while the drone
flies a figure-eight.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
from __future__ import annotations
import argparse
import os
from typing import TYPE_CHECKING
os.environ["SCIPY_ARRAY_API"] = "1"
import jax
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.data import SimStateDeriv
from crazyflow.sim.dynamics import first_principles_dynamics
from crazyflow.sim.pipeline import append_fn, insert_fn_before
if TYPE_CHECKING:
from crazyflow.sim.data import SimData
DURATION = 10.0 # s, one loop of the figure-eight
def dynamics_deriv(data: SimData) -> SimData:
"""Evaluate the dynamics."""
return data.replace(plugins=data.plugins | {"states_deriv": first_principles_dynamics(data)})
def finite_diff_deriv(data: SimData) -> SimData:
"""Differentiate the states over the last step."""
prev, states, freq = data.plugins["prev_states"], data.states, data.core.freq
rot = R.from_quat(prev.quat).inv() * R.from_quat(states.quat)
deriv = SimStateDeriv(
vel=(states.pos - prev.pos) * freq,
ang_vel=rot.as_rotvec() * freq,
acc=(states.vel - prev.vel) * freq,
ang_acc=(states.ang_vel - prev.ang_vel) * freq,
rotor_acc=(states.rotor_vel - prev.rotor_vel) * freq,
)
return data.replace(plugins=data.plugins | {"fd_states_deriv": deriv, "prev_states": states})
def trajectory(t: float) -> np.ndarray:
"""Return a figure-eight state command."""
omega = 2 * np.pi / DURATION
cmd = np.zeros((1, 1, 16))
cmd[..., :3] = [2 * np.sin(omega * t), np.sin(2 * omega * t), 1.0]
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
return cmd
def main(plot: bool = False):
results = {}
for integrator in ("euler", "rk4"):
sim = Sim(dynamics="first_principles", control="state", integrator=integrator)
pos = jnp.asarray(trajectory(0.0)[..., :3], device=sim.device)
sim.data = sim.data.replace(states=sim.data.states.replace(pos=pos))
plugins = {
"states_deriv": SimStateDeriv.create(sim.n_worlds, sim.n_drones, sim.device),
"fd_states_deriv": SimStateDeriv.create(sim.n_worlds, sim.n_drones, sim.device),
"prev_states": sim.data.states,
}
sim.data = sim.data.replace(plugins=sim.data.plugins | plugins)
insert_fn_before(sim.step_pipeline, "integration", dynamics_deriv)
append_fn(sim.step_pipeline, finite_diff_deriv)
sim.build_default_data()
sim.build_step_fn()
log = {"states_deriv": [], "fd_states_deriv": []}
for i in range(int(DURATION * sim.control_freq)):
sim.state_control(trajectory(i / sim.control_freq))
sim.step(sim.freq // sim.control_freq)
for key in log:
log[key].append(jax.tree.map(lambda x: x[0, 0], sim.data.plugins[key]))
sim.close()
dynamics, fd = (jax.tree.map(lambda *x: np.stack(x), *log[key]) for key in log)
results[integrator] = dynamics, fd
for name in ("vel", "ang_vel", "acc", "ang_acc", "rotor_acc"):
x, x_fd = getattr(dynamics, name), getattr(fd, name)
diff = np.abs(x - x_fd).max() / np.abs(x).max()
print(f"{integrator} {name}: max relative difference {diff:.1e}")
if plot:
plot_results(results, sim.control_freq)
def plot_results(results: dict[str, tuple[SimStateDeriv, SimStateDeriv]], freq: int):
import matplotlib.pyplot as plt
fig, axes = plt.subplots(2, 2, sharex=True, figsize=(12, 7))
quantities = (
("acc", "Linear acceleration", "m/s$^2$"),
("ang_acc", "Angular acceleration", "rad/s$^2$"),
)
for row, (integrator, (dynamics, fd)) in enumerate(results.items()):
for col, (name, title, unit) in enumerate(quantities):
x, x_fd = getattr(dynamics, name), getattr(fd, name)
t = np.arange(len(x)) / freq
for i, axis in enumerate("xyz"):
axes[row, col].plot(t, x_fd[:, i] - x[:, i], f"C{i}", lw=0.8, label=axis)
ylabel = f"Finite differences - dynamics [{unit}]"
axes[row, col].set(title=f"{title} ({integrator})", ylabel=ylabel)
for ax in axes[-1]:
ax.set(xlabel="Time [s]")
axes[0, 0].legend()
for ax in axes.flat:
ax.grid()
fig.tight_layout()
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot)
IMU¶
Simulate an IMU as a plugin.
The accelerometer measures the specific force and the gyroscope the body rates, both with noise and drifting biases, while the drone flies two loops of a figure-eight. The plot compares the readings with the true state, skipping the first loop to remove initial transients.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/plugins/imu.py
"""Simulate an IMU as a plugin.
The accelerometer measures the specific force and the gyroscope the body rates, both with noise and
drifting biases, while the drone flies two loops of a figure-eight. The plot compares the readings
with the true state, skipping the first loop to remove initial transients.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
from __future__ import annotations
import argparse
import os
from functools import partial
from typing import TYPE_CHECKING, Callable
os.environ["SCIPY_ARRAY_API"] = "1"
import jax
import jax.numpy as jnp
import numpy as np
from flax.struct import dataclass, field
from scipy.spatial.transform import Rotation as R
from crazyflow.sim import Sim
from crazyflow.sim.pipeline import prepend_fn
from crazyflow.sim.sim import select_dynamics_fn
from crazyflow.utils import CORE_NDIM_KEY
if TYPE_CHECKING:
from jax import Array, Device
from crazyflow.sim.data import SimData, SimStateDeriv
ACCEL_STD = 0.03 # m/s^2
GYRO_STD = 0.0025 # rad/s
ACCEL_BIAS_WALK_STD = 1e-3 # m/s^2/sqrt(s)
GYRO_BIAS_WALK_STD = 0.002 # rad/s/sqrt(s), inflated such that the bias walk is visible
GYRO_BIAS_MAX = 0.01 # rad/s, the initial bias is uniform in [-max, max]
PERIOD = 20.0 # s, one loop of the figure-eight
DURATION = 2 * PERIOD # s, the first loop is dropped to remove initial transients
@dataclass
class IMUData:
"""IMU measurements, biases and noise parameters."""
accel: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3)
"""Latest accelerometer measurement as specific force in the body frame."""
gyro: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3)
"""Latest gyroscope measurement as angular velocity in the body frame."""
accel_bias: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3)
"""Accelerometer bias."""
gyro_bias: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 3)
"""Gyroscope bias."""
accel_std: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Standard deviation of the accelerometer noise."""
gyro_std: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Standard deviation of the gyroscope noise."""
accel_bias_walk_std: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Standard deviation of the accelerometer bias random walk."""
gyro_bias_walk_std: Array = field(metadata={CORE_NDIM_KEY: 1}) # (N, M, 1)
"""Standard deviation of the gyroscope bias random walk."""
@staticmethod
def create(
n_worlds: int,
n_drones: int,
accel_std: float,
gyro_std: float,
accel_bias_walk_std: float,
gyro_bias_walk_std: float,
device: Device,
) -> IMUData:
"""Create IMU data with zero measurements and biases."""
zeros = jnp.zeros((n_worlds, n_drones, 3), device=device)
std_shape = (n_worlds, n_drones, 1)
return IMUData(
accel=zeros,
gyro=zeros,
accel_bias=zeros,
gyro_bias=zeros,
accel_std=jnp.full(std_shape, accel_std, device=device),
gyro_std=jnp.full(std_shape, gyro_std, device=device),
accel_bias_walk_std=jnp.full(std_shape, accel_bias_walk_std, device=device),
gyro_bias_walk_std=jnp.full(std_shape, gyro_bias_walk_std, device=device),
)
def simulate_imu(data: SimData, deriv_fn: Callable[[SimData], SimStateDeriv]) -> SimData:
"""Sample the IMU at the current state."""
imu: IMUData = data.plugins["imu"]
key, accel_key, gyro_key, accel_walk_key, gyro_walk_key = jax.random.split(data.core.rng_key, 5)
sqrt_dt = (1 / data.core.freq) ** 0.5
accel_walk = jax.random.normal(accel_walk_key, imu.accel_bias.shape) * imu.accel_bias_walk_std
accel_bias = imu.accel_bias + accel_walk * sqrt_dt
gyro_walk = jax.random.normal(gyro_walk_key, imu.gyro_bias.shape) * imu.gyro_bias_walk_std
gyro_bias = imu.gyro_bias + gyro_walk * sqrt_dt
# Accelerometers measure the specific force in the body frame, f_b = R_wb^T (a_w - g_w)
acc = deriv_fn(data).acc
accel = R.from_quat(data.states.quat).apply(acc - data.params.gravity_vec, inverse=True)
accel = accel + accel_bias + jax.random.normal(accel_key, acc.shape) * imu.accel_std
gyro = data.states.ang_vel + gyro_bias + jax.random.normal(gyro_key, acc.shape) * imu.gyro_std
imu = imu.replace(accel=accel, gyro=gyro, accel_bias=accel_bias, gyro_bias=gyro_bias)
return data.replace(plugins=data.plugins | {"imu": imu}, core=data.core.replace(rng_key=key))
def add_imu(
sim: Sim,
accel_std: float = ACCEL_STD,
gyro_std: float = GYRO_STD,
accel_bias_walk_std: float = ACCEL_BIAS_WALK_STD,
gyro_bias_walk_std: float = GYRO_BIAS_WALK_STD,
):
"""Add the IMU data to the plugins and sample the IMU at the start of each step.
The caller rebuilds the default data and the step function.
"""
imu = IMUData.create(
sim.n_worlds,
sim.n_drones,
accel_std,
gyro_std,
accel_bias_walk_std,
gyro_bias_walk_std,
sim.device,
)
sim.data = sim.data.replace(plugins=sim.data.plugins | {"imu": imu})
simulate_fn = partial(simulate_imu, deriv_fn=select_dynamics_fn(sim.dynamics))
prepend_fn(sim.step_pipeline, simulate_fn, name="simulate_imu")
def trajectory(t: float) -> np.ndarray:
"""Return a figure-eight position command."""
omega = 2 * np.pi / PERIOD
cmd = np.zeros((1, 1, 16))
cmd[..., :3] = [np.sin(omega * t), 0.75 * np.sin(2 * omega * t), 1.0]
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
return cmd
def main(plot: bool = False):
sim = Sim(control="state")
states = sim.data.states.replace(pos=trajectory(0.0)[..., :3])
add_imu(sim)
imu: IMUData = sim.data.plugins["imu"]
key, bias_key = jax.random.split(sim.data.core.rng_key)
bias = jax.random.uniform(
bias_key, imu.gyro_bias.shape, minval=-GYRO_BIAS_MAX, maxval=GYRO_BIAS_MAX
)
sim.data = sim.data.replace(
states=states,
plugins=sim.data.plugins | {"imu": imu.replace(gyro_bias=bias)},
core=sim.data.core.replace(rng_key=key),
)
sim.build_default_data()
sim.build_step_fn()
log = {key: [] for key in ("vel", "quat", "ang_vel", "accel", "gyro")}
for i in range(int(DURATION * sim.control_freq)):
sim.state_control(trajectory(i / sim.control_freq))
sim.step(sim.freq // sim.control_freq)
log["vel"].append(sim.data.states.vel[0, 0])
log["quat"].append(sim.data.states.quat[0, 0])
log["ang_vel"].append(sim.data.states.ang_vel[0, 0])
log["accel"].append(sim.data.plugins["imu"].accel[0, 0])
log["gyro"].append(sim.data.plugins["imu"].gyro[0, 0])
sim.close()
log = {key: np.array(value) for key, value in log.items()}
# Compare the accelerations without gravity, a_b = R_wb^T a_w, so that the noise is more visible
rot = R.from_quat(log["quat"])
acc = np.gradient(log["vel"], 1 / sim.control_freq, axis=0)
log["acc"] = rot.apply(acc, inverse=True)
log["accel"] = log["accel"] + rot.apply(sim.data.params.gravity_vec, inverse=True)
# Drop the first loop
log = {key: value[int(PERIOD * sim.control_freq) :] for key, value in log.items()}
accel_error = np.sqrt(np.mean(np.sum((log["accel"] - log["acc"]) ** 2, axis=-1)))
gyro_error = np.sqrt(np.mean(np.sum((log["gyro"] - log["ang_vel"]) ** 2, axis=-1)))
print(
f"Accelerometer error RMS {accel_error:.3f} m/s^2, "
f"gyroscope error RMS {gyro_error:.4f} rad/s"
)
if plot:
plot_results(log)
def plot_results(log: dict[str, np.ndarray]):
import matplotlib.pyplot as plt
fig, (ax_accel, ax_gyro) = plt.subplots(1, 2, figsize=(11, 4.5))
t = np.linspace(0, PERIOD, len(log["gyro"]))
sensors = (
(ax_accel, log["acc"], log["accel"], "Accelerometer without gravity", "[m/s$^2$]"),
(ax_gyro, log["ang_vel"], log["gyro"], "Gyroscope", "[rad/s]"),
)
for ax, true, reading, title, unit in sensors:
readings = [
ax.plot(t, reading[:, i], lw=0.6, label=f"{x} reading")[0] for i, x in enumerate("xyz")
]
true_lines = ax.plot(t, true, "k", lw=0.8, label="true")
ax.legend(handles=[true_lines[0], *readings])
ax.set(title=title, xlabel="Time [s]", ylabel=unit)
fig.tight_layout()
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot)