JAX examples¶
Applying JAX transformations to the simulator: differentiating through the dynamics with jax.grad, distributing the simulation across devices, and caching compiled functions. See Sharding for the details on multiple devices.
Gradient descent through dynamics¶
Optimize an attitude command with gradient descent through the simulation.
Because the simulator is built entirely from JAX operations, jax.grad can differentiate through
it. Starting the drone in the air keeps it away from the floor, so the floor-clipping stage never
applies and gradients flow freely through the entire trajectory.
examples/jax/gradient.py
"""Optimize an attitude command with gradient descent through the simulation.
Because the simulator is built entirely from JAX operations, `jax.grad` can differentiate through
it. Starting the drone in the air keeps it away from the floor, so the floor-clipping stage never
applies 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 command clipping¶
Compare gradients through the rotor command clip with straight-through and no clipping.
Hard clips such as the clip_rotor_vel_cmd stage zero the gradients while the command 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.
Run with --plot to plot the results, which requires matplotlib from the tests environment.
examples/jax/gradient_clipping.py
"""Compare gradients through the rotor command clip with straight-through and no clipping.
Hard clips such as the `clip_rotor_vel_cmd` stage zero the gradients while the command 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.
Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
"""
import argparse
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_cmd_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
cmd = data.controls.rotor_vel
cmd = cmd + jax.lax.stop_gradient(jnp.clip(cmd, lower, upper) - cmd)
return data.replace(controls=data.controls.replace(rotor_vel=cmd))
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 command is saturated
results = {"clip (default)": rollout(sim, cmds)}
# Option 2: clip the command in the forward pass, but keep the gradients flowing in the
# backward pass (straight-through estimator)
clip_fn = partial(clip_rotor_vel_cmd_nonblocking, lower=lower, upper=upper)
replace_fn(sim.step_pipeline, clip_fn, "clip_rotor_vel_cmd")
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_cmd")
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 command clipping and its effect on gradients")
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)
Sharding across devices¶
Shard the simulation across devices.
The host backend is asked for four logical devices, and the sharded results are compared against running on a single device.
examples/jax/sharding.py
"""Shard the simulation across devices.
The host backend is asked for four logical devices, and the sharded results are compared against
running on a single device.
"""
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 = 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(pos, single_device_pos, atol=1e-6)
sim.close()
if __name__ == "__main__":
main()
Persistent compilation cache¶
Speed up the first simulation step with a persistent JAX compilation cache.
JAX traces and compiles jitted functions to XLA on their first call. This is expensive, but later
calls reuse the compiled version. Sim uses many jitted functions internally, and its first step
compiles the whole chain of dynamics and control functions, so creating a Sim and stepping it
for the first time is slow.
The compiled functions are lost when the Python session ends. enable_cache stores them on disk
instead, so later sessions load them rather than recompiling. Compilation then only happens once
per machine, or again after the cache directory is deleted.
Run the script twice to compare. The first run compiles and fills the cache, the second loads it and deletes it again, so the runs alternate between slow and fast.
examples/jax/cache.py
"""Speed up the first simulation step with a persistent JAX compilation cache.
JAX traces and compiles jitted functions to XLA on their first call. This is expensive, but later
calls reuse the compiled version. `Sim` uses many jitted functions internally, and its first `step`
compiles the whole chain of dynamics and control functions, so creating a `Sim` and stepping it
for the first time is slow.
The compiled functions are lost when the Python session ends. `enable_cache` stores them on disk
instead, so later sessions load them rather than recompiling. Compilation then only happens once
per machine, or again after the cache directory is deleted.
Run the script twice to compare. The first run compiles and fills the cache, the second loads it
and deletes it again, so the runs alternate between slow and fast.
"""
import shutil
import time
from pathlib import Path
from crazyflow.sim import Sim
from crazyflow.utils import enable_cache
def main():
cache_dir = Path("/tmp/jax_cache_test")
if use_cache := cache_dir.exists():
print("Cache directory exists. This run will be fast.")
print("\nTo run without cache, run this script again.")
else:
print("Cache directory does not exist. This run will be slow.")
print("\nTo run with cache, run this script again.")
enable_cache(cache_path=cache_dir)
t0 = time.perf_counter()
sim = Sim()
t1 = time.perf_counter()
sim.step()
t2 = time.perf_counter()
prefix = "Using cache: " if use_cache else "Not using cache: "
print(f"{prefix}\n Init: {t1 - t0:.3f}s\n Step: {t2 - t1:.3f}s")
if use_cache:
shutil.rmtree(cache_dir) # Clean up cache so that the next run is slow again
sim.close()
if __name__ == "__main__":
main()