Examples¶
These runnable examples cover control, JAX transformations, pipeline extensions, rendering, contacts, and Gymnasium environments. Start with hover if you're new, or jump to the section that matches your use case.
Hover¶
A single drone commanded to hold a fixed height using state control. This is the minimal end-to-end loop: create a Sim, reset it, apply a state command, and step forward.
import numpy as np
from crazyflow.control import Control
from crazyflow.sim import Dynamics, Sim
def main():
sim = Sim(
n_worlds=1,
n_drones=1,
dynamics=Dynamics.first_principles,
control=Control.state,
freq=500,
attitude_freq=500,
state_freq=100,
device="cpu",
)
sim.reset()
duration = 5.0
fps = 60
# State cmd is [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]
cmd = np.zeros((sim.n_worlds, sim.n_drones, 13))
cmd[..., :3] = 0.1
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()
sim.close()
if __name__ == "__main__":
main()
Attitude control¶
Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses the Mellinger position loop and is typical for RL agents that output attitude targets.
from functools import partial
import numpy as np
from crazyflow.control import Control, parametrize
from crazyflow.control.mellinger import state2attitude
from crazyflow.sim import Sim
def control(t: float, pos_start: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Compute the attitude command to track a circle with a slow climb."""
cmd = np.zeros(13)
cmd[:3] = pos_start + np.array([np.cos(t) - 1, np.sin(t), 0.2 * t])
cmd[3:6] = np.array([-np.sin(t), np.cos(t), 0.2])
cmd[6:9] = np.array([-np.cos(t), -np.sin(t), 0.0])
cmd[9] = t # Yaw
return cmd
def main():
sim = Sim(control=Control.attitude)
sim.reset()
duration = 6.5
fps = 60
# We use the Mellinger position controller to generate attitude commands. This could be any
# controller that outputs [roll, pitch, yaw, thrust], e.g. a learned policy.
position_ctrl = partial(parametrize(state2attitude, sim.drone), ctrl_freq=sim.control_freq)
pos_err_i = np.zeros(3)
cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) # [roll, pitch, yaw, thrust]
pos_start = np.asarray(sim.data.states.pos[0, 0])
for i in range(int(duration * sim.control_freq)):
pos, quat = np.asarray(sim.data.states.pos[0, 0]), np.asarray(sim.data.states.quat[0, 0])
vel = np.asarray(sim.data.states.vel[0, 0])
ref = control(i / sim.control_freq, pos_start)
cmd[0, 0, :], pos_err_i = position_ctrl(pos, quat, vel, ref, pos_err_i)
sim.attitude_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
sim.close()
if __name__ == "__main__":
main()
Body rate control¶
Commanding body-frame angular rates and collective thrust. The firmware controller has no dedicated body rate mode and levels the drone with its attitude terms, so the example sets the kR and ki_m gains of the body rate controller to zero.
import os
os.environ["SCIPY_ARRAY_API"] = "1"
from functools import partial
import jax.numpy as jnp
import numpy as np
from scipy.spatial.transform import Rotation as R
from crazyflow.control import Control, parametrize
from crazyflow.control.mellinger import state2attitude
from crazyflow.sim import Sim
kp_att = 8.0 # Proportional gain from the attitude error to body rates
def trajectory(t: float, pos_start: np.ndarray) -> np.ndarray:
"""Compute the full state command of a circle with a slow climb."""
cmd = np.zeros(13)
cmd[:3] = pos_start + np.array([np.cos(t) - 1, np.sin(t), 0.2 * t])
cmd[3:6] = np.array([-np.sin(t), np.cos(t), 0.2])
cmd[6:9] = np.array([-np.cos(t), -np.sin(t), 0.0])
cmd[9] = t # Yaw
return cmd
def control(quat: np.ndarray, rpyt: np.ndarray) -> np.ndarray:
"""Convert an attitude command into body rates with a proportional attitude loop."""
rot_err = (R.from_quat(quat).inv() * R.from_euler("xyz", rpyt[:3])).as_rotvec()
return np.concatenate([kp_att * rot_err, rpyt[3:]])
def main():
sim = Sim(control=Control.body_rate, body_rate_freq=250)
# The firmware has no dedicated body rate mode. Its attitude terms level the drone at the
# current yaw and would counteract the commanded rates. Disable them to track body rates.
body_rate = sim.data.controls.body_rate
params = body_rate.params | {"kR": jnp.zeros(3), "ki_m": jnp.zeros(3)}
controls = sim.data.controls.replace(body_rate=body_rate.replace(params=params))
sim.data = sim.data.replace(controls=controls)
sim.build_default_data()
sim.reset()
duration = 6.5
fps = 60
# We use the Mellinger position controller to generate attitude commands, which we then convert
# to body rates. This could be any controller that outputs [w_x, w_y, w_z, thrust].
position_ctrl = partial(parametrize(state2attitude, sim.drone), ctrl_freq=sim.control_freq)
pos_err_i = np.zeros(3)
cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) # [roll_rate, pitch_rate, yaw_rate, thrust]
pos_start = np.asarray(sim.data.states.pos[0, 0])
for i in range(int(duration * sim.control_freq)):
pos, quat = np.asarray(sim.data.states.pos[0, 0]), np.asarray(sim.data.states.quat[0, 0])
vel = np.asarray(sim.data.states.vel[0, 0])
ref = trajectory(i / sim.control_freq, pos_start)
rpyt, pos_err_i = position_ctrl(pos, quat, vel, ref, pos_err_i)
cmd[0, 0, :] = control(quat, rpyt)
sim.body_rate_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
sim.close()
if __name__ == "__main__":
main()
Sampling-based MPC¶
A sampling-based model predictive controller tracks a Lissajous curve while avoiding a grid of obstacles. It rolls out thousands of candidate control sequences in parallel using identified dynamics, then applies the first action from a cost-weighted update of the best samples. The controller automatically uses a GPU when one is available and lowers the sample count on CPU.
Gradient descent through dynamics¶
Because the simulator is built entirely from JAX operations, jax.grad can differentiate through it. Starting the drone above the target height keeps it away from the floor, so the floor-clipping stage never fires and gradients flow freely through the entire trajectory.
import time
import jax
import jax.numpy as jnp
from numpy.typing import NDArray
from crazyflow.control import Control
from crazyflow.sim import Dynamics, Sim
from crazyflow.sim.data import SimData
def main():
sim = Sim(control=Control.attitude, dynamics=Dynamics.first_principles, attitude_freq=50)
# Remove clipping floor function which kills gradients
sim_step = sim.build_step_fn()
# If the drone starts on the floor, the gradient gets killed by the floor clipping function. We
# thus start in the air to avoid zero gradients. Alternatively, we could also remove the floor
# clipping function
sim.data = sim.data.replace(
states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.5))
)
def step(cmd: NDArray, data: SimData) -> jax.Array:
data = data.replace(
controls=data.controls.replace(attitude=data.controls.attitude.replace(staged_cmd=cmd))
)
data = sim_step(data, 10)
return (data.states.pos[0, 0, 2] - 1.0) ** 2 # Quadratic cost to reach 1m height
step_grad = jax.jit(jax.grad(step))
cmd = jnp.zeros((1, 1, 4), dtype=jnp.float32)
cmd = cmd.at[..., 3].set(sim.data.params.mass[0] * 9.81 * 1.05)
# Trigger jax's jit to compile the gradient function. This is not necessary, but it ensures that
# the timings are not affected by the compilation time.
step_grad(cmd, sim.data).block_until_ready()
# JAX compiles again if static properties change. Not sure why this is happening here, but this
# is a simple way to enforce all recompilations before measuring performance.
step_grad(cmd - 0.1 * step_grad(cmd, sim.data), sim.data).block_until_ready()
print(f"Initial command: {cmd}")
t0 = time.perf_counter()
for _ in range(10):
grad = step_grad(cmd, sim.data)
cmd = cmd - 0.1 * grad
t1 = time.perf_counter()
print(f"Loss: {step(cmd, sim.data)}\nGradient: {grad}")
print(f"Time taken: {t1 - t0:.2e}s ({(t1 - t0) / 10:.2e}s per step)")
# The final command should increase the z position (3rd array element) as well as the z velocity
# (6th array element) to minimize the cost function.
print(f"Final command: {cmd}")
if __name__ == "__main__":
main()
Gradients and state clipping¶
Hard state clips such as the clip_rotor_vel stage zero the gradients while the state is saturated. This example ramps a motor command beyond the rotor limits and compares the rotor state and the gradient of the vertical acceleration w.r.t. the command for three options: the default clip, a straight-through clip (clipped forward pass, unclipped gradients), and no clip. Replacing the default clip with the straight-through variant can help gradient-based methods such as trajectory optimization or policy learning, which would otherwise receive zero gradients whenever the motors saturate.
from functools import partial
import jax
import jax.numpy as jnp
import numpy as np
from numpy.typing import NDArray
import crazyflow.sim.functional as F
from crazyflow.control import Control
from crazyflow.sim import Sim
from crazyflow.sim.data import SimData
from crazyflow.sim.pipeline import remove_fn, replace_fn
from crazyflow.sim.sim import rotor_vel_limits
def clip_rotor_vel_nonblocking(data: SimData, lower: float, upper: float) -> SimData:
# Straight-through estimator: x + stop_gradient(clip(x) - x) evaluates to clip(x) in the
# forward pass, while its derivative w.r.t. x is 1 in the backward pass
rotor_vel = data.states.rotor_vel
rotor_vel = rotor_vel + jax.lax.stop_gradient(jnp.clip(rotor_vel, lower, upper) - rotor_vel)
return data.replace(states=data.states.replace(rotor_vel=rotor_vel))
def rollout(sim: Sim, cmds: NDArray) -> tuple[NDArray, NDArray]:
step_fn = sim.build_step_fn()
# The command only enters the acceleration through the rotor state of the *next* step, so we
# measure the acceleration with a one-step lookahead while the outer loop advances by a single
# step per command
def acc_z(cmd: jax.Array, data: SimData) -> tuple[jax.Array, SimData]:
data = F.rotor_vel_control(data, jnp.full((1, 1, 4), cmd))
data = step_fn(data, 1)
lookahead = step_fn(data, 1)
acc = (lookahead.states.vel[0, 0, 2] - data.states.vel[0, 0, 2]) * sim.freq
return acc, data
grad_fn = jax.jit(jax.value_and_grad(acc_z, has_aux=True))
data, rotor_vel, grads = sim.data, [], []
for cmd in cmds:
(_, data), grad = grad_fn(jnp.float32(cmd), data)
rotor_vel.append(data.states.rotor_vel[0, 0, 0])
grads.append(grad)
return np.array(rotor_vel), np.array(grads)
def main(plot: bool = False):
sim = Sim(control=Control.rotor_vel)
lower, upper = rotor_vel_limits(sim.dynamics, sim.drone)
# Start in the air so that the drone never reaches the floor, where the floor clipping would
# zero the velocity and kill the gradients (see gradient.py)
sim.data = sim.data.replace(
states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(2.0))
)
# Motor command (RPM): ramp up beyond the upper limit, hold, ramp back down to zero
ramp = float(upper) + 10_000
n = 250 # 0.5 s per segment at 500 Hz
cmds = np.concatenate([np.linspace(0, ramp, n), np.full(n, ramp), np.linspace(ramp, 0, n)])
# Option 1: keep the clipping as is. The rotor state respects the limits, but the gradient is
# zero while the state is saturated
results = {"clip (default)": rollout(sim, cmds)}
# Option 2: clip the state in the forward pass, but keep the gradients flowing in the backward
# pass (straight-through estimator)
clip_fn = partial(clip_rotor_vel_nonblocking, lower=lower, upper=upper)
replace_fn(sim.step_pipeline, clip_fn, "clip_rotor_vel")
results["nonblocking clip"] = rollout(sim, cmds)
# Option 3: remove the clipping. Gradients always flow, but the state can leave the limits
remove_fn(sim.step_pipeline, "clip_rotor_vel")
results["no clip"] = rollout(sim, cmds)
sim.close()
if plot:
plot_results(cmds, results, float(lower), float(upper), sim.freq)
def plot_results(
cmds: NDArray,
results: dict[str, tuple[NDArray, NDArray]],
lower: float,
upper: float,
freq: int,
):
# Only import if plotting is desired to avoid a dependency on matplotlib
import matplotlib.pyplot as plt
t = np.arange(len(cmds)) / freq
fig, (ax_state, ax_grad) = plt.subplots(1, 2, sharex="all", figsize=(12, 5))
ax_state.plot(t, cmds, label="command", color="gray", linestyle=":")
# The forward pass of the nonblocking clip is identical to the default clip, so we plot it
# dashed to keep both curves visible
styles = {"nonblocking clip": {"linestyle": "--"}}
for name, (rotor_vel, grads) in results.items():
ax_state.plot(t, rotor_vel, label=name, **styles.get(name, {}))
ax_grad.plot(t, grads, label=name, **styles.get(name, {}))
for limit in (lower, upper):
ax_state.axhline(limit, color="black", linestyle="--", linewidth=0.8)
ax_state.set_title("Rotor state")
ax_state.set_ylabel("rotor_vel (RPM)")
ax_grad.set_title("Gradient d acc_z / d cmd")
ax_grad.set_ylabel("(m/s$^2$) / RPM")
for ax in (ax_state, ax_grad):
ax.set_xlabel("Time (s)")
ax.legend()
ax.grid(True)
fig.suptitle("Rotor state clipping and its effect on gradients")
plt.tight_layout()
plt.show()
if __name__ == "__main__":
main(plot=True)
Sharding across devices¶
We can also distribute the simulation across devices. The host backend is asked for four logical devices. We then compare this against running on a single device. See Sharding for the details.
"""Example on how to shard the simulation across devices."""
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
import jax
import numpy as np
from crazyflow.sim import Sim
from crazyflow.sim.sharding import world_mesh
def main():
devices = jax.devices("cpu")
sim = Sim(n_worlds=4 * len(devices), device="cpu") # Worlds must be divisible by n_devices
sim.step(10)
single_device_pos = np.asarray(sim.data.states.pos)
sim.reset()
sim.shard(world_mesh(devices))
sim.step(10)
pos, gravity = sim.data.states.pos, sim.data.params.gravity_vec
print(f"Distributing {sim.n_worlds} worlds over {len(devices)} devices")
# Arrays that carry a world axis are partitioned. Shared params are replicated on every device
print(" Placement")
print(f" {'position':<{10}} {str(pos.shape):<{10}} {pos.sharding.spec}")
print(f" {'gravity':<{10}} {str(gravity.shape):<{10}} {gravity.sharding.spec}")
print(" Shards of states.pos")
for shard in pos.addressable_shards:
print(f" {shard.device} {shard.data.shape[0]} worlds")
assert np.allclose(np.asarray(pos), single_device_pos, atol=1e-6)
sim.close()
if __name__ == "__main__":
main()
Domain randomization¶
Randomizing mass, inertia and drag per drone, and thrust and torque curves, rotor dynamics, arm lengths and propeller inertias per motor through the reset pipeline, each by a uniform factor around its default value. An optional mask limits randomization to selected worlds.
"""Example showing how to randomize parameters of the simulation.
All shown parameters can be randomized per world and per drone. All randomizations scale each
element of the default parameters 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.
"""
import jax
import jax.numpy as jnp
import numpy as np
from jax import Array
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, 13))
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¶
Inserting a random external force and torque into the step pipeline. The disturbance fires on every dynamics tick, so the drone fights wind-like perturbations.
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, 13))
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__":
main(plot=True) # Default is False to disable plotting during testing
Cameras and RGBD¶
Offscreen rendering returns RGB-D images on every frame. The FPV camera (fpv_cam) is attached to the drone and moves with it.
"""Example showing how to change the used camera and how to extract the pixel information."""
import time
import matplotlib.pyplot as plt
import mujoco
import numpy as np
from matplotlib import animation
from crazyflow.control import Control
from crazyflow.dynamics import Dynamics
from crazyflow.sim import Sim
from crazyflow.sim.integration import Integrator
def control(t: float, t_tot: float) -> np.ndarray:
phi = 2 * np.pi * t / t_tot + np.pi
circle = np.array([np.cos(phi), np.sin(phi)])
cmd = np.zeros((1, 1, 13))
cmd[..., :2] = circle # xy
cmd[..., 2] = 0.1 + 0.5 * t / t_tot # z
cmd[..., -4] = 1.9 * np.pi * t / t_tot # yaw
return cmd
def add_smiley(sim: Sim):
# Add 3d object to sim
# create box spec from an XML string
box_xml = """
<mujoco model="box_model">
<worldbody>
<body name="cube" pos="0 0 0">
<geom type="box" size="0.05 0.05 0.05" rgba="0.8 0.4 0.2 1"/>
</body>
</worldbody>
</mujoco>
"""
box_spec = mujoco.MjSpec.from_string(box_xml)
frame = sim.spec.worldbody.add_frame()
boxes = [
# eyes
((0.0, -0.15, 0.6), (1, 0, 0, 0)),
((0.0, 0.15, 0.6), (1, 0, 0, 0)),
# mouth
((0.0, -0.2, 0.4), (1, 0, 0, 0)),
((0.0, 0.2, 0.4), (1, 0, 0, 0)),
((0.0, -0.1, 0.3), (1, 0, 0, 0)),
((0.0, 0.0, 0.3), (1, 0, 0, 0)),
((0.0, 0.1, 0.3), (1, 0, 0, 0)),
]
for i, x in enumerate(boxes):
box_body = box_spec.body("cube")
box = frame.attach_body(box_body, "", f":{i}")
box.pos = x[0]
box.quat = x[1]
sim.build_mjx()
sim.build_reset_fn()
def main(show_plot: bool = False, save_plot: bool = False):
"""Example showing the rendering feature and saving a gif via FuncAnimation."""
# Setup sim
sim = Sim(
n_drones=1,
control=Control.state,
integrator=Integrator.rk4,
dynamics=Dynamics.first_principles,
drone="cf2x_T350",
)
add_smiley(sim)
sim.reset()
pos = sim.data.states.pos.at[...].set([-1, 0, 0])
states = sim.data.states.replace(pos=pos)
sim.data = sim.data.replace(states=states)
duration = 5
fps = 50
timings = []
# Set up matplotlib rendering
resolution = (160, 120)
rgb = np.zeros((resolution[1], resolution[0], 3))
d = np.zeros((resolution[1], resolution[0]))
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
im1 = ax1.imshow(rgb)
ax1.set_title("RGB")
ax1.axis("off")
im2 = ax2.imshow(d, cmap="viridis")
ax2.set_title("Depth")
ax2.axis("off")
fig.tight_layout()
# Animation setup
def update_frame(_): # noqa: ANN202
t = sim.data.core.steps[0, 0] / sim.freq
sim.state_control(control(t, duration))
sim.step(sim.freq // fps)
t1 = time.perf_counter()
# mode: Either "human" for the regular window, "rgb_array" for an RGB array,
# "depth_array" for a depth array, or "rgbd_tuple" for both at the same time.
# camera: The name or id of the camera. The names are specified in the corresponding
# xml file in crazyflow/drones. For example, "fpv_cam:0" is the first-person view
# camera of the first drone, "track_cam:0" is the tracking camera of the first
# drone. Id -1 is the global camera.
rgbd = sim.render(
width=resolution[0], height=resolution[1], mode="rgbd_tuple", camera="fpv_cam:0"
)
t2 = time.perf_counter()
timings.append(t2 - t1)
if rgbd is None:
return im1, im2
rgb, depth = rgbd
im1.set_data(rgb)
im2.set_data(depth)
im2.set_clim(np.nanmin(depth), np.nanmax(depth))
return im1, im2
anim = animation.FuncAnimation(
fig, update_frame, frames=int(duration * fps), interval=1000 / fps, blit=True, repeat=False
)
if show_plot:
plt.show()
if save_plot:
anim.save("cameras.gif", writer="pillow", fps=fps)
sim.close()
t_mean = np.mean(timings)
print(f"Average render time {t_mean * 1000:.2f}ms, eqivalent to {1 / t_mean:.2f}fps")
print("For more optimized depth rendering, check out the raycasting.py example.")
if __name__ == "__main__":
main(show_plot=True, save_plot=False)
LED deck and materials¶
change_material updates the RGBA colour and emission of any named material on any subset of drones at runtime.
import tempfile
from pathlib import Path
import numpy as np
from crazyflow.control import Control
from crazyflow.sim import Sim
from crazyflow.sim.visualize import change_material
scene_dark_xml = """
<mujoco model="Drone scene">
<option integrator="RK4" density="1.225" viscosity="1.8e-5" timestep="0.001"/>
<compiler inertiafromgeom="false" meshdir="assets" autolimits="true"/>
<statistic center="0 0 2" extent="2.5"/>
<visual>
<rgba haze="0.15 0.25 0.35 0" fog="1 1 1 0"/>
<map fogstart="0" fogend="0"/>
<global azimuth="-20" elevation="-20" ellipsoidinertia="true"/>
</visual>
<asset>
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
markrgb="0.8 0.8 0.8" width="512" height="512"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="2 2" reflectance="0.2"/>
</asset>
<worldbody>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
</worldbody>
</mujoco>
""" # noqa: E501
def main():
"""Spawn 25 drones in one world and activate led decks."""
try:
# Use a named temporary xml file
with tempfile.NamedTemporaryFile(suffix=".xml", delete=False) as tmp:
tmp.write(scene_dark_xml.encode())
tmp.flush()
tmp_path = Path(tmp.name)
sim = Sim(n_drones=25, drone="cf21B_500", control=Control.state, xml_path=tmp_path)
fps = 60
cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))
cmd[..., 3] = sim.data.params.mass[0] * 9.81
rgbas = np.random.default_rng(0).uniform(0, 1, (sim.n_drones, 4))
rgbas[..., 3] = 1.0
init_pos = np.array(sim.data.states.pos[0, :, :])
cmd = np.zeros((sim.n_worlds, sim.n_drones, 13))
cmd[:, :, :3] = init_pos
cmd[:, :, 2] += 1.5
for i in range(int(10 * sim.control_freq)):
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
even_ids = np.arange(0, sim.n_drones, 2)
odd_ids = np.arange(1, sim.n_drones, 2)
emission = np.sin(i / sim.control_freq * np.pi)
change_material(
sim,
mat_name="led_top",
drone_ids=even_ids,
rgba=rgbas[even_ids, :],
emission=emission,
)
change_material(
sim,
mat_name="led_bot",
drone_ids=odd_ids,
rgba=rgbas[odd_ids, :],
emission=emission,
)
sim.render()
sim.close()
finally:
# clean up the temporary file
if tmp_path is not None and tmp_path.exists():
try:
tmp_path.unlink()
except Exception:
pass
if __name__ == "__main__":
main()
Contact queries¶
The default collision geometry is a sphere around the drone frame. use_box_collision replaces it with a tighter oriented box, useful for narrow-gap flight and accurate contact debugging.
import numpy as np
from crazyflow.sim import Dynamics, Sim
from crazyflow.sim.sim import use_box_collision
def main():
"""Spawn multiple drones in multiple worlds and check for contacts."""
n_worlds, n_drones = 2, 3
sim = Sim(n_worlds=n_worlds, n_drones=n_drones, dynamics=Dynamics.so_rpy, device="cpu")
use_box_collision(sim, enable=True) # Enable box collision for all drones
fps = 60
cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))
cmd[..., 3] = sim.data.params.mass[0] * 9.81 * 1.04
for i in range(int(2 * sim.control_freq)):
sim.attitude_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
print(f"Contacts: {sim.contacts().any()}")
sim.close()
if __name__ == "__main__":
main()
Raycasting and depth sensing¶
render_depth fires rays from a camera and returns per-pixel distances. This is faster than full RGB rendering and useful for obstacle sensing or depth-based controllers.
import jax.numpy as jnp
import matplotlib.pyplot as plt
from crazyflow.sim import Sim
from crazyflow.sim.sensors.depth import build_render_depth_fn, render_depth
def main(plot: bool = False):
sim = Sim()
sim.data = sim.data.replace(
states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.2))
)
# The easiest way to get depth images is to use the render_depth function
dist = render_depth(sim, camera=0, resolution=(100, 100), include_drone=False)
dist = dist.at[dist > 1.5].set(jnp.nan) # Cap max distance for better visualization
if plot:
plt.imshow(dist[0], cmap="viridis")
plt.colorbar(label="Distance (m)")
plt.title("Raycast Distance from Camera")
plt.show()
# We can also build a depth renderer function for better performance if we need maximum speed or
# more fine-grained control. Here we only render the drone collision geometry to avoid expensive
# raycasting against the high-poly visual mesh of the drone.
render_depth_fn = build_render_depth_fn(
sim.mjx_model, camera=0, resolution=(200, 200), geomgroup=(1, 1, 0, 1, 1, 1, 1, 1)
)
dist_fn = render_depth_fn(sim)
dist_fn = dist_fn.at[dist_fn > 1.5].set(jnp.nan) # Cap max distance for better visualization
if plot:
plt.imshow(dist_fn[0], cmap="viridis")
plt.colorbar(label="Distance (m)")
plt.title("Raycast Distance from Camera (Compiled)")
plt.show()
if __name__ == "__main__":
main(plot=True)
Gymnasium environment¶
Evaluating a random policy in the figure-8 environment. The env wraps Sim behind the standard Gymnasium VectorEnv interface.
import gymnasium
import jax.numpy as jnp
import numpy as np
from gymnasium.wrappers.vector import JaxToNumpy # , JaxToTorch
from crazyflow.envs import NormalizeActions # noqa: F401
from crazyflow.utils import enable_cache
def main():
enable_cache()
# Create environment that contains a figure eight trajectory. You can parametrize the
# observation space, i.e., which part of the trajectory is contained in the observation. Please
# refer to the documentation of the environment for more information.
envs = gymnasium.make_vec(
"DroneFigureEightTrajectory-v0",
num_envs=20,
freq=50,
n_samples=10,
samples_dt=0.1,
trajectory_time=10.0,
)
# NormalizeActions wrapper to clip the actions to [-1, 1] and rescale them for use with common
# DRL libraries.
envs = NormalizeActions(envs)
envs = JaxToNumpy(envs)
# dummy action for going up (in attitude control)
action = np.zeros((20, 4), dtype=np.float32)
action[..., 3] = 0.3
obs, info = envs.reset()
# Step through the environment
for _ in range(1_000):
# Prevent alignment warnings. Related issue: https://github.com/jax-ml/jax/issues/29810
# TODO: Remove once https://github.com/jax-ml/jax/pull/29963 is merged.
action = np.asarray(jnp.asarray(action))
observation, reward, terminated, truncated, info = envs.step(action)
envs.render()
envs.close()
if __name__ == "__main__":
main()