Control examples¶
Flying drones with the different control modes, from state commands down to force and torque, comparing the dynamics models, and model-based control. See Control Modes for the controller behind each mode.
Hover¶
Fly a single drone to a fixed position and hold it with state control.
This is the minimal end-to-end loop: create a Sim, reset it, apply a state command, and step
forward.
examples/control/hover.py
"""Fly a single drone to a fixed position and hold it with state control.
This is the minimal end-to-end loop: create a `Sim`, reset it, apply a state command, and step
forward.
"""
import os
os.environ["SCIPY_ARRAY_API"] = "1"
import numpy as np
from scipy.spatial.transform import Rotation as R
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, qx, qy, qz, qw, wx, wy, wz]
cmd = np.zeros((sim.n_worlds, sim.n_drones, 16))
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
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()
Changing the state¶
Overwrite the state of a drone before flying it to a setpoint.
JAX arrays are immutable, so sim.data cannot be changed in place. Instead, replace creates a
new sim.data with the position and rotor velocities of the drone set to new values.
examples/control/change_pos.py
"""Overwrite the state of a drone before flying it to a setpoint.
JAX arrays are immutable, so `sim.data` cannot be changed in place. Instead, `replace` creates a
new `sim.data` with the position and rotor velocities of the drone set to new values.
"""
import os
os.environ["SCIPY_ARRAY_API"] = "1"
import numpy as np
from scipy.spatial.transform import Rotation as R
from crazyflow.sim import Sim
def main():
sim = Sim(control="state")
sim.reset()
# Replace the position of the first drone in the first world. JAX arrays are immutable, which is
# why we cannot change the sim.data object in-place. Instead, we need to create a new sim.data
# object with the desired changes by calling sim.data.replace(). The same logic applies to the
# sim.data.states object contained within sim.data, and the sim.data.states.pos array. For more
# information on changing JAX arrays, see:
# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates
sim.data = sim.data.replace(
states=sim.data.states.replace(
pos=sim.data.states.pos.at[0, 0].set(np.array([0.5, 0.5, 0.2])),
rotor_vel=sim.data.states.rotor_vel.at[0, 0].set(np.ones(4) * 20000),
)
)
control = np.zeros((sim.n_worlds, sim.n_drones, 16))
control[..., 9:13] = R.from_euler("z", 0.0).as_quat()
control[..., :3] = np.array([[0.0, 0.0, 0.3]])
for _ in range(3 * sim.control_freq):
sim.state_control(control)
sim.step(sim.freq // sim.control_freq)
sim.render()
sim.close()
if __name__ == "__main__":
main()
Spirals¶
Example where 4 drones fly upwards in a spiral motion.
examples/control/spiral.py
"""Example where 4 drones fly upwards in a spiral motion."""
import os
os.environ["SCIPY_ARRAY_API"] = "1"
import numpy as np
from scipy.spatial.transform import Rotation as R
from crazyflow.control import Control
from crazyflow.sim import Sim
def trajectory(start_xy: np.ndarray, t: float) -> np.ndarray:
circle = np.array([np.cos(t) - 1, np.sin(t)])
cmd = np.zeros((*start_xy.shape[:-1], 16))
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
cmd[..., :2] = start_xy + circle
cmd[..., 2] = 0.2 * t
return cmd
def main():
sim = Sim(n_drones=4, control=Control.state, integrator="rk4", dynamics="first_principles")
sim.reset()
duration = 5.0
fps = 60
start_xy = sim.data.states.pos[..., :2]
for i in range(int(duration * sim.control_freq)):
sim.state_control(trajectory(start_xy, i / sim.control_freq))
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
sim.close()
if __name__ == "__main__":
main()
Attitude control¶
Command roll, pitch, yaw, and collective thrust directly.
This level bypasses the Mellinger position loop and is typical for RL agents that output attitude targets.
examples/control/attitude.py
"""Command roll, pitch, yaw, and collective thrust directly.
This level bypasses the Mellinger position loop and is typical for RL agents that output attitude
targets.
"""
import os
os.environ["SCIPY_ARRAY_API"] = "1"
from functools import partial
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
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(16)
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:13] = R.from_euler("z", t).as_quat()
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 = 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)
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¶
Command body-frame angular rates and collective thrust.
The body rate controller keeps the attitude terms of the Mellinger controller, which level the
drone at its current yaw and counteract the commanded rates. The example sets their kR and ki_m
gains to zero to freely track the body rates.
examples/control/body_rate.py
"""Command body-frame angular rates and collective thrust.
The body rate controller keeps the attitude terms of the Mellinger controller, which level the
drone at its current yaw and counteract the commanded rates. The example sets their `kR` and `ki_m`
gains to zero to freely track the body rates.
"""
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(16)
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:13] = R.from_euler("z", t).as_quat()
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 attitude terms of the body rate controller 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)) # [wx, wy, wz, thrust]
pos_start = 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()
Force and torque control¶
Command collective thrust and body torques directly.
The thrust slightly exceeds the weight of the drone, so it slowly accelerates upwards.
examples/control/force_torque.py
"""Command collective thrust and body torques directly.
The thrust slightly exceeds the weight of the drone, so it slowly accelerates upwards.
"""
import numpy as np
from crazyflow.control import Control
from crazyflow.sim import Sim
def main():
sim = Sim(control=Control.force_torque)
sim.reset()
duration = 5.0
fps = 60
cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) # [fz, tx, ty, tz]
mass, gravity = sim.data.params.mass, -sim.data.params.gravity_vec[-1]
cmd[..., 0] = (mass + 1e-4) * gravity # Plus a small margin to accelerate slightly
for i in range(int(duration * sim.control_freq)):
sim.force_torque_control(cmd)
sim.step(sim.freq // sim.control_freq)
if ((i * fps) % sim.control_freq) < fps:
sim.render()
sim.close()
if __name__ == "__main__":
main()
Comparing integrators¶
Fly straight up with rotor velocity control and compare the integrators.
All rotors get the same constant command above hover. The plot shows the height difference of every integrator to the default one.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/control/integrators.py
"""Fly straight up with rotor velocity control and compare the integrators.
All rotors get the same constant command above hover. The plot shows the height difference of every
integrator to the default one.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
import argparse
import numpy as np
from crazyflow.control import Control
from crazyflow.control.transform import motor_force2rotor_vel
from crazyflow.sim import Sim
from crazyflow.sim.integration import Integrator
DURATION = 2.0
def main(plot: bool = False):
heights = {}
for integrator in Integrator:
sim = Sim(control=Control.rotor_vel, integrator=integrator)
sim.reset()
thrust = 1.2 * sim.data.params.mass * -sim.data.params.gravity_vec[-1]
rpm = motor_force2rotor_vel(thrust / 4, sim.data.params.rpm2thrust)
sim.rotor_vel_control(np.broadcast_to(rpm, (sim.n_worlds, sim.n_drones, 4)))
z = [sim.data.states.pos[0, 0, 2]]
for _ in range(int(DURATION * sim.freq)):
sim.step(1)
z.append(sim.data.states.pos[0, 0, 2])
sim.close()
heights[integrator] = np.array(z)
print(f"{integrator.value}: z = {z[-1]:.4f} m after {DURATION} s")
if plot:
plot_results(heights, sim.freq)
def plot_results(heights: dict[Integrator, np.ndarray], freq: int):
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
t = np.arange(len(heights[Integrator.default])) / freq
for integrator, z in heights.items():
if integrator != Integrator.default:
ax.plot(t, z - heights[Integrator.default], label=integrator.value)
ax.set(xlabel="Time [s]", ylabel="Δz [m]")
ax.set_title(f"Height difference to {Integrator.default.value}")
ax.legend()
ax.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)
Comparing dynamics models¶
Compare all dynamics models while tracking the same figure-eight with state control.
Each model starts from the same reference state, and the position RMSE to the reference is printed for every model. The plot shows the x-z trajectories and the lateral positions against the reference.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/control/dynamics.py
"""Compare all dynamics models while tracking the same figure-eight with state control.
Each model starts from the same reference state, and the position RMSE to the reference is
printed for every model. The plot shows the x-z trajectories and the lateral positions against the
reference.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
import argparse
import os
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.control import Control
from crazyflow.control.transform import motor_force2rotor_vel
from crazyflow.sim import Dynamics, Sim
DURATION = 5.0
FPS = 60
TRAJ_CENTER = np.array([0.0, 0.0, 1.0])
TRAJ_SIZE = np.array([1.0, 0.0, 0.5])
def trajectory(t: float) -> np.ndarray:
"""Return the position, velocity, and acceleration reference at time ``t``."""
omega = 2.0 * np.pi / DURATION
phi = omega * t
cmd = np.zeros((1, 1, 16))
cmd[..., 9:13] = R.from_euler("z", 0.0).as_quat()
cmd[..., 0:3] = TRAJ_CENTER + TRAJ_SIZE * np.array([np.sin(phi), 0.0, np.sin(2.0 * phi)])
cmd[..., 3:6] = TRAJ_SIZE * omega * np.array([np.cos(phi), 0.0, 2.0 * np.cos(2.0 * phi)])
cmd[..., 6:9] = -TRAJ_SIZE * omega**2 * np.array([np.sin(phi), 0.0, 4.0 * np.sin(2.0 * phi)])
return cmd
def reset_sim(sim: Sim) -> None:
"""Reset a simulation to the initial reference state and hover rotor speed."""
sim.reset()
initial_state = trajectory(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([trajectory(t) for t in times])
reference = commands[:, 0, 0, 0:3]
positions = [sim.data.states.pos[0, 0]]
for i, cmd in enumerate(commands[1:]):
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)
positions.append(sim.data.states.pos[0, 0])
if ((i * FPS) % sim.control_freq) < FPS and render:
sim.render()
sim.close()
positions = np.asarray(positions)
trajectories[dynamics] = positions
rmses[dynamics] = np.sqrt(np.mean(np.sum((positions - 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, positions in trajectories.items():
label = f"{dynamics.value} (RMSE={rmses[dynamics] * 1_000:.0f}mm)"
ax_xz.plot(positions[:, 0], positions[:, 2], label=label)
ax_y.plot(times, positions[:, 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__":
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Plot the results")
main(plot=parser.parse_args().plot)
Symbolic dynamics¶
Create symbolic CasADi dynamics for model-based control.
symbolic_dynamics returns the dynamics as CasADi expressions with the nominal parameters of the
simulated drone, so they can be used in model-based controllers such as MPC. A CasADi integrator
turns them into discrete-time dynamics at the simulation frequency.
examples/control/symbolic.py
"""Create symbolic CasADi dynamics for model-based control.
`symbolic_dynamics` returns the dynamics as CasADi expressions with the nominal parameters of the
simulated drone, so they can be used in model-based controllers such as MPC. A CasADi integrator
turns them into discrete-time dynamics at the simulation frequency.
"""
import casadi as cs
import numpy as np
from crazyflow.dynamics import parametrize
from crazyflow.dynamics.so_rpy import symbolic_dynamics
from crazyflow.sim import Sim
def main():
# We can create a symbolic dynamics directly from the simulation. Note that this will use the
# nominal parameters of the simulation and choose the control type based on the simulation.
sim = Sim(dynamics="so_rpy", freq=500)
X_dot, X, U, Y = parametrize(symbolic_dynamics, sim.drone)()
assert X_dot.shape == (13, 1)
assert X.shape == (13, 1)
assert U.shape == (4, 1) # Attitude control
assert Y.shape == (7, 1) # 3 for pos and 4 for quat
# To create a discrete-time dynamics that you can integrate, you can use the integrator function
# from CasADi.
fd = cs.integrator("fd", "cvodes", {"x": X, "p": U, "ode": X_dot}, 0, 1 / sim.freq)
x0 = np.ones((13, 1))
u = np.ones((4, 1))
res = fd(x0=x0, p=u)
assert res["xf"].shape == (13, 1)
if __name__ == "__main__":
main()
Sampling-based MPC¶
Track a Lissajous curve with a sampling-based MPC while avoiding a grid of obstacles.
The controller 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. It uses a GPU when one is available and lowers the sample count on CPU.
WARNING: This is an advanced example meant for advanced Crazyflow users.
examples/control/sampling.py
"""Track a Lissajous curve with a sampling-based MPC while avoiding a grid of obstacles.
The controller 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. It uses a
GPU when one is available and lowers the sample count on CPU.
WARNING: This is an advanced example meant for advanced Crazyflow users.
"""
import os
from collections import deque
from functools import partial
from typing import Callable
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import jax
import jax.numpy as jnp
import numpy as np
from jax import Array
from jax.lax import scan
from crazyflow.control import Control
from crazyflow.control.transform import motor_force2rotor_vel
from crazyflow.dynamics import load_params
from crazyflow.sim import Dynamics, Sim
from crazyflow.sim.data import SimData
from crazyflow.sim.visualize import draw_capsule, draw_line
try:
# If available, use the GPU for the controller to be able to run more samples
jax.devices("gpu")
DEVICE_CONTROLLER = "gpu"
except RuntimeError:
DEVICE_CONTROLLER = "cpu"
# Simulation configuration
DRONE = "cf21B_500"
DURATION = 10.0
FPS = 60
RENDER = True
# Controller configuration
CTRL_FREQ = 50
T = 1.0 # Prediction horizon in seconds
N = 25 # Prediction steps
N_SAMPLES = 500_000 if DEVICE_CONTROLLER == "gpu" else 2_000
NOISE_SIGMA = jnp.array([0.10, 0.10, 0.0, 0.08], dtype=jnp.float32)
ELITE_FRACTION = 0.01 # Fraction of samples in (0,1] to use for the mean update
MAX_CMD_ANGLE = np.deg2rad(60.0)
# Lissajous trajectory configuration
TRAJ_CENTER = jnp.array([0.0, 0.0, 1.0], dtype=jnp.float32)
TRAJ_SIZE = jnp.array([1.0, 0.75, 0.0], dtype=jnp.float32)
TRAJ_PERIOD = 6.0
# Visualization configuration
HISTORY_DURATION = 1.0
N_RENDER_SAMPLES = 32 # How many rollouts are visualized
# Pole-grid configuration
OBSTACLE_GRID = (6, 5)
OBSTACLE_SPACING = (0.5, 0.5)
OBSTACLE_CENTER = (0.0, 0.0)
OBSTACLE_RADIUS = 0.055
OBSTACLE_HEIGHT = 1.6
DRONE_RADIUS = 0.12
OBSTACLE_MARGIN = 0.02
def trajectory(t: Array | float) -> dict[str, Array]:
"""Return the position, velocity, and yaw reference at time ``t``."""
omega = 2.0 * jnp.pi / TRAJ_PERIOD
phi = omega * jnp.asarray(t)
zero = jnp.zeros_like(phi)
pos = TRAJ_CENTER + TRAJ_SIZE * jnp.stack([jnp.sin(phi), jnp.sin(2.0 * phi), zero], axis=-1)
vel = TRAJ_SIZE * omega * jnp.stack([jnp.cos(phi), 2.0 * jnp.cos(2.0 * phi), zero], axis=-1)
return {"pos": pos, "vel": vel, "yaw": zero}
def obstacle_grid() -> np.ndarray:
"""Return pole base positions for a centered rectangular grid."""
nx, ny = OBSTACLE_GRID
x = (np.arange(nx) - (nx - 1) / 2) * OBSTACLE_SPACING[0] + OBSTACLE_CENTER[0]
y = (np.arange(ny) - (ny - 1) / 2) * OBSTACLE_SPACING[1] + OBSTACLE_CENTER[1]
xx, yy = np.meshgrid(x, y, indexing="ij")
return np.column_stack((xx.ravel(), yy.ravel(), np.zeros(nx * ny))).astype(np.float32)
def step_sim(
data: SimData,
inputs: tuple[Array, dict[str, Array]],
step_fn: Callable[[SimData, int], SimData],
obstacles: Array,
hover_thrust: Array,
) -> tuple[SimData, tuple[Array, Array]]:
"""Apply one candidate input per world, advance all rollouts one step and compute the cost."""
command, reference = inputs
data = data.replace(
controls=data.controls.replace(
attitude=data.controls.attitude.replace(staged_cmd=command[:, None, :])
)
)
next_data = step_fn(data, 1)
pos = next_data.states.pos[:, 0]
vel = next_data.states.vel[:, 0]
cmd = next_data.controls.attitude.staged_cmd[:, 0]
# Tracking cost
pos_error = jnp.linalg.norm(pos - reference["pos"], axis=-1)
vel_error = jnp.linalg.norm(vel - reference["vel"], axis=-1)
track_cost = 50.0 * pos_error**2 + vel_error**2
# Control cost
tilt_cost = 5.0 * jnp.linalg.norm(cmd[:, :2], axis=-1) ** 2
thrust_cost = 5.0 * (cmd[:, 3] - hover_thrust) ** 2
yaw_cost = 100.0 * (cmd[:, 2] - reference["yaw"]) ** 2
input_cost = tilt_cost + thrust_cost + yaw_cost
# Obstacle cost
obstacle_distance = jnp.linalg.norm(pos[:, None, :2] - obstacles[None, :, :2], axis=-1)
obstacle_hits = obstacle_distance < OBSTACLE_RADIUS + DRONE_RADIUS + OBSTACLE_MARGIN
obstacle_cost = 1_000.0 * jnp.sum(obstacle_hits, axis=-1)
cost = track_cost + input_cost + obstacle_cost
return next_data, (cost, next_data.states.pos[:, 0])
def rollout_sim(
obs: dict[str, Array],
command: Array,
reference: dict[str, Array],
rollout_data: SimData,
step_fn: Callable[[SimData, int], SimData],
obstacles: Array,
hover_thrust: Array,
) -> tuple[Array, Array]:
"""Roll out control sequences from the current state (obs) and compute their costs."""
states = rollout_data.states.replace(
pos=rollout_data.states.pos.at[...].set(obs["pos"]),
quat=rollout_data.states.quat.at[...].set(obs["quat"]),
vel=rollout_data.states.vel.at[...].set(obs["vel"]),
ang_vel=rollout_data.states.ang_vel.at[...].set(obs["ang_vel"]),
# The reduced dynamics store collective thrust in its rotor_vel state.
rotor_vel=rollout_data.states.rotor_vel.at[...].set(obs["collective_thrust"]),
)
data = rollout_data.replace(states=states)
_, (costs, positions) = scan(
partial(step_sim, step_fn=step_fn, obstacles=obstacles, hover_thrust=hover_thrust),
data,
(command, reference),
)
return jnp.sum(costs, axis=0), positions
def update_controller(
t: Array,
obs: dict[str, Array],
key: Array,
mean_controls: Array,
rollout_fn: Callable[[dict[str, Array], Array, dict[str, Array]], tuple[Array, Array]],
hover_cmd: Array,
action_low: Array,
action_high: Array,
noise_sigma: Array,
) -> tuple[Array, Array, Array, Array, Array]:
"""Update the control trajectory from a cost-weighted mean of the elite samples."""
key, sample_key = jax.random.split(key)
noise = jax.random.normal(sample_key, (N_SAMPLES, N, 4)) * noise_sigma
candidates = jnp.clip(mean_controls[None] + noise, action_low, action_high)
candidates = candidates.at[0].set(mean_controls)
prediction_dt = T / N
times = t + (jnp.arange(N) + 1) * prediction_dt
references = trajectory(times)
costs, positions = rollout_fn(obs, candidates.transpose(1, 0, 2), references)
n_elites = max(1, int(N_SAMPLES * ELITE_FRACTION))
elite_indices = jnp.argsort(costs)[:n_elites]
updated_controls = jnp.mean(candidates[elite_indices], axis=0)
action = updated_controls[0]
prediction_times = jnp.arange(N) * prediction_dt
shifted_times = prediction_times + 1.0 / CTRL_FREQ
def shift_control(control: Array, hover: Array) -> Array:
return jnp.interp(shifted_times, prediction_times, control, right=hover)
mean_controls = jax.vmap(shift_control, in_axes=(1, 0), out_axes=1)(updated_controls, hover_cmd)
best_index = elite_indices[0]
best_positions = positions[:, best_index]
render_indices = elite_indices[jnp.linspace(0, n_elites - 1, N_RENDER_SAMPLES, dtype=jnp.int32)]
sampled_positions = positions[:, render_indices].transpose(1, 0, 2)
return action, key, mean_controls, best_positions, sampled_positions
def control(
t: float,
obs: dict[str, Array],
key: Array,
mean_controls: Array,
controller_fn: Callable,
controller_device: jax.Device,
) -> tuple[np.ndarray, Array, Array, np.ndarray, np.ndarray]:
"""Move the observation to the rollout device and compute one SMPC update."""
obs = jax.tree.map(lambda value: jax.device_put(value, controller_device), obs)
t_device = jax.device_put(jnp.asarray(t, dtype=jnp.float32), controller_device)
action, key, mean_controls, best_positions, sampled_positions = controller_fn(
t_device, obs, key, mean_controls
)
return (
np.asarray(action),
key,
mean_controls,
np.asarray(best_positions),
np.asarray(sampled_positions),
)
def main() -> None:
obstacles = obstacle_grid()
# Set up the main sim
sim = Sim(n_worlds=1, drone=DRONE, dynamics=Dynamics.first_principles, control=Control.attitude)
sim.max_visual_geom = 100_000 # To be able to show all rollouts
sim.reset()
start_pos = trajectory(0.0)["pos"]
drone_params = load_params(Dynamics.first_principles, DRONE)
hover_thrust_value = drone_params["mass"] * 9.81
hover_rotor_vel = motor_force2rotor_vel(
np.full(4, hover_thrust_value / 4.0, dtype=np.float32), drone_params["rpm2thrust"]
)
sim.data = sim.data.replace(
states=sim.data.states.replace(
pos=sim.data.states.pos.at[0, 0].set(start_pos),
rotor_vel=sim.data.states.rotor_vel.at[0, 0].set(hover_rotor_vel),
)
)
# Set up the controller
controller_device = jax.devices(DEVICE_CONTROLLER)[0]
rollout_freq = int(N / T)
rollout_simulator = Sim(
n_worlds=N_SAMPLES,
device=controller_device.platform,
drone=DRONE,
dynamics=Dynamics.so_rpy_rotor_drag,
control=Control.attitude,
freq=rollout_freq,
attitude_freq=rollout_freq,
)
rollout_simulator.reset()
thrust_estimate = hover_thrust_value # Initial thrust estimate
thrust_dyn_coef = float(rollout_simulator.data.params.thrust_dyn_coef[0])
hover_cmd = jax.device_put(
jnp.array([0.0, 0.0, 0.0, hover_thrust_value], dtype=jnp.float32), controller_device
)
action_low = jax.device_put(
jnp.array([-MAX_CMD_ANGLE, -MAX_CMD_ANGLE, 0.0, 0.0], dtype=jnp.float32), controller_device
)
max_thrust_value = 4.0 * drone_params["thrust_max"]
action_high = jax.device_put(
jnp.array([MAX_CMD_ANGLE, MAX_CMD_ANGLE, 0.0, max_thrust_value], dtype=jnp.float32),
controller_device,
)
noise_sigma = jax.device_put(NOISE_SIGMA, controller_device)
rollout_fn = partial(
rollout_sim,
rollout_data=rollout_simulator.data,
step_fn=rollout_simulator.build_step_fn(),
obstacles=jax.device_put(jnp.asarray(obstacles), controller_device),
hover_thrust=jax.device_put(jnp.asarray(hover_thrust_value), controller_device),
)
controller_fn = jax.jit(
partial(
update_controller,
rollout_fn=rollout_fn,
hover_cmd=hover_cmd,
action_low=action_low,
action_high=action_high,
noise_sigma=noise_sigma,
)
)
mean_controls = jax.device_put(jnp.broadcast_to(hover_cmd, (N, 4)), controller_device)
key = jax.device_put(jax.random.key(0), controller_device)
# Plotting helpers
position_history = deque(maxlen=max(2, round(HISTORY_DURATION * CTRL_FREQ)))
render_times = np.linspace(0.0, TRAJ_PERIOD, round(TRAJ_PERIOD * CTRL_FREQ) + 1)
reference = np.asarray(trajectory(render_times)["pos"])
for step in range(int(DURATION * CTRL_FREQ)):
t = step / CTRL_FREQ
obs = {
"pos": sim.data.states.pos[0, 0],
"quat": sim.data.states.quat[0, 0],
"vel": sim.data.states.vel[0, 0],
"ang_vel": sim.data.states.ang_vel[0, 0],
# Thrust is not observable and difficult to estimate, so use the thrust dynamics.
"collective_thrust": thrust_estimate,
}
action, key, mean_controls, best_positions, sampled_positions = control(
t, obs, key, mean_controls, controller_fn, controller_device
)
thrust_estimate += thrust_dyn_coef * (action[3] - thrust_estimate) / CTRL_FREQ
sim.attitude_control(action[None, None])
sim.step(sim.freq // CTRL_FREQ)
position_history.append(sim.data.states.pos[0, 0])
if RENDER and ((step * FPS) % CTRL_FREQ) < FPS:
col_ref = np.array([0.0, 0.1, 1.0, 1.0])
col_pred = np.array([0.2, 1.0, 0.2, 0.8])
col_rollouts = np.array([0.25, 0.85, 0.85, 0.18])
col_hist = np.array([1.0, 0.25, 0.0, 1.0])
col_pole = np.array([0.9, 0.9, 0.9, 1.0])
draw_line(sim, reference, rgba=col_ref, start_size=2.0, end_size=2.0)
for sampled_path in sampled_positions:
draw_line(sim, sampled_path, rgba=col_rollouts, start_size=0.15, end_size=0.15)
draw_line(sim, best_positions, rgba=col_pred)
if len(position_history) > 1:
draw_line(sim, np.array(position_history), rgba=col_hist, start_size=0.01)
for pole in obstacles:
pole_top = pole + np.array([0.0, 0.0, OBSTACLE_HEIGHT])
draw_capsule(sim, pole, pole_top, radius=OBSTACLE_RADIUS, rgba=col_pole)
sim.render(
cam_config={
"distance": 3.0,
"elevation": -45.0,
"azimuth": 90.0,
"lookat": [0.0, 0.0, 1.0],
}
)
sim.close()
rollout_simulator.close()
if __name__ == "__main__":
main()