From 08b1bfc3d872059533689d36f67a288a30cf8d3e Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Fri, 12 Jun 2026 18:14:40 +0200 Subject: [PATCH 1/8] Restructure examples --- crazyflow/randomize/__init__.py | 3 - crazyflow/randomize/randomize.py | 48 ---------- docs/api/index.md | 1 - docs/examples/index.md | 36 +++---- docs/gen_ref_pages.py | 1 - docs/user-guide/gymnasium-envs.md | 14 +-- docs/user-guide/oo-api.md | 47 ++++----- docs/user-guide/pipelines.md | 53 ++++++---- examples/{ => contacts}/contacts.py | 0 examples/{ => contacts}/crash.py | 0 examples/{ => control}/attitude.py | 0 examples/{ => control}/change_pos.py | 0 examples/{ => control}/force_torque.py | 0 examples/{ => control}/hover.py | 0 examples/{ => control}/spiral.py | 0 examples/{ => environments}/figure8.py | 0 examples/{ => environments}/gymnasium_env.py | 0 examples/{ => jax}/cache.py | 0 examples/{ => jax}/gradient.py | 0 .../{plugins.py => plugins/action_delay.py} | 0 examples/{ => plugins}/disturbance.py | 0 examples/plugins/estimation.py | 1 + examples/plugins/randomize.py | 63 ++++++++++++ examples/randomize.py | 43 --------- examples/{ => rendering}/cameras.py | 0 examples/{ => rendering}/led_deck.py | 0 examples/{ => rendering}/raycasting.py | 0 examples/{ => rendering}/render.py | 0 tests/integration/test_randomize.py | 50 ---------- tests/unit/test_randomize.py | 96 ------------------- 30 files changed, 143 insertions(+), 313 deletions(-) delete mode 100644 crazyflow/randomize/__init__.py delete mode 100644 crazyflow/randomize/randomize.py rename examples/{ => contacts}/contacts.py (100%) rename examples/{ => contacts}/crash.py (100%) rename examples/{ => control}/attitude.py (100%) rename examples/{ => control}/change_pos.py (100%) rename examples/{ => control}/force_torque.py (100%) rename examples/{ => control}/hover.py (100%) rename examples/{ => control}/spiral.py (100%) rename examples/{ => environments}/figure8.py (100%) rename examples/{ => environments}/gymnasium_env.py (100%) rename examples/{ => jax}/cache.py (100%) rename examples/{ => jax}/gradient.py (100%) rename examples/{plugins.py => plugins/action_delay.py} (100%) rename examples/{ => plugins}/disturbance.py (100%) create mode 100644 examples/plugins/estimation.py create mode 100644 examples/plugins/randomize.py delete mode 100644 examples/randomize.py rename examples/{ => rendering}/cameras.py (100%) rename examples/{ => rendering}/led_deck.py (100%) rename examples/{ => rendering}/raycasting.py (100%) rename examples/{ => rendering}/render.py (100%) delete mode 100644 tests/integration/test_randomize.py delete mode 100644 tests/unit/test_randomize.py diff --git a/crazyflow/randomize/__init__.py b/crazyflow/randomize/__init__.py deleted file mode 100644 index 963740f..0000000 --- a/crazyflow/randomize/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from crazyflow.randomize.randomize import randomize_inertia, randomize_mass - -__all__ = ["randomize_mass", "randomize_inertia"] diff --git a/crazyflow/randomize/randomize.py b/crazyflow/randomize/randomize.py deleted file mode 100644 index d40f1c6..0000000 --- a/crazyflow/randomize/randomize.py +++ /dev/null @@ -1,48 +0,0 @@ -import jax -import jax.numpy as jnp -from jax import Array - -from crazyflow.sim import Sim -from crazyflow.sim.data import SimData -from crazyflow.utils import leaf_replace - - -def randomize_mass(sim: Sim, mass: Array, mask: Array | None = None): - """Randomize mass from a new masses. - - Args: - sim: The simulation object. - mass: The new masses. The shape always needs to be (n_worlds, n_drones). - mask: Boolean array of shape (n_worlds, ) that indicates which worlds to reset. If None, - all worlds are reset. - """ - sim.data = _randomize_mass_params(sim.data, mass, mask) - - -def randomize_inertia(sim: Sim, J: Array, mask: Array | None = None): - """Randomize inertia tensor from a new inertia tensors. - - Args: - sim: The simulation object. - J: The inertia tensors of shape (n_worlds, n_drones, 3, 3). - mask: Boolean array of shape (n_worlds, ) that indicates which worlds to reset. If None, - all worlds are reset. - - Warning: - This only works for first_principles dynamics. - """ - if not J.shape == (sim.n_worlds, sim.n_drones, 3, 3): - raise ValueError(f"Inertia tensor must have shape (n_worlds, n_drones, 3, 3), is {J.shape}") - sim.data = _randomize_inertia_params(sim.data, J, mask) - - -@jax.jit -def _randomize_mass_params(data: SimData, mass: Array, mask: Array | None = None) -> SimData: - mass = jnp.atleast_3d(mass) - assert mass.shape[2] == 1, f"Expected shape (n_worlds, n_drones, 1), is {mass.shape}" - return data.replace(params=leaf_replace(data.params, mask, mass=mass)) - - -@jax.jit -def _randomize_inertia_params(data: SimData, J: Array, mask: Array | None = None) -> SimData: - return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J))) diff --git a/docs/api/index.md b/docs/api/index.md index 63b6068..e1c9443 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -16,7 +16,6 @@ This section is auto-generated from the Crazyflow source code using [mkdocstring | `crazyflow.control` | `Control` enum | | `crazyflow.control.mellinger` | Mellinger controller data and parameters | | `crazyflow.envs` | Gymnasium vectorized environments | -| `crazyflow.randomize` | Domain randomization helpers | | `crazyflow.utils` | Grid utilities and pytree helpers | Navigate the full generated reference using the sidebar. diff --git a/docs/examples/index.md b/docs/examples/index.md index daad37a..5f555ea 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -1,6 +1,6 @@ # Examples -These examples build on each other — each one introduces one new concept on top of the previous. Start from the top if you're new, or jump to whichever section covers what you need. +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. --- @@ -9,11 +9,11 @@ These examples build on each other — each one introduces one new concept on to 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. ```{ .python notest } ---8<-- "examples/hover.py" +--8<-- "examples/control/hover.py" ``` ```bash -python examples/hover.py +python examples/control/hover.py ``` --- @@ -23,7 +23,7 @@ python examples/hover.py 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. ```{ .python notest } ---8<-- "examples/attitude.py" +--8<-- "examples/control/attitude.py" ``` --- @@ -33,17 +33,21 @@ Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses 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. ```{ .python notest } ---8<-- "examples/gradient.py" +--8<-- "examples/jax/gradient.py" ``` --- ## Domain randomization -Varying physical parameters per world at reset. Each world gets a slightly different mass, so identical commands produce diverging trajectories. +Randomizing mass and inertia through the reset pipeline. An optional mask limits randomization to selected worlds. ```{ .python notest } ---8<-- "examples/randomize.py" +--8<-- "examples/plugins/randomize.py" +``` + +```bash +python examples/plugins/randomize.py ``` --- @@ -53,7 +57,7 @@ Varying physical parameters per world at reset. Each world gets a slightly diffe Inserting a random external force and torque into the step pipeline. The disturbance fires on every physics tick, so the drone fights wind-like perturbations. ```{ .python notest } ---8<-- "examples/disturbance.py" +--8<-- "examples/plugins/disturbance.py" ``` --- @@ -67,11 +71,11 @@ Offscreen rendering returns RGB and depth images on every frame. The FPV camera ```{ .python notest } ---8<-- "examples/cameras.py" +--8<-- "examples/rendering/cameras.py" ``` ```bash -python examples/cameras.py +python examples/rendering/cameras.py ``` --- @@ -85,11 +89,11 @@ python examples/cameras.py ```{ .python notest } ---8<-- "examples/led_deck.py" +--8<-- "examples/rendering/led_deck.py" ``` ```bash -python examples/led_deck.py +python examples/rendering/led_deck.py ``` --- @@ -108,7 +112,7 @@ The default collision geometry is a sphere around the drone frame. `use_box_coll ```{ .python notest } ---8<-- "examples/contacts.py" +--8<-- "examples/contacts/contacts.py" ``` --- @@ -118,11 +122,11 @@ The default collision geometry is a sphere around the drone frame. `use_box_coll `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. ```{ .python notest } ---8<-- "examples/raycasting.py" +--8<-- "examples/rendering/raycasting.py" ``` ```bash -python examples/raycasting.py +python examples/rendering/raycasting.py ``` --- @@ -132,5 +136,5 @@ python examples/raycasting.py Evaluating a random policy in the figure-8 environment. The env wraps `Sim` behind the standard Gymnasium `VectorEnv` interface. ```{ .python notest } ---8<-- "examples/figure8.py" +--8<-- "examples/environments/figure8.py" ``` diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py index 58de12f..73a4324 100644 --- a/docs/gen_ref_pages.py +++ b/docs/gen_ref_pages.py @@ -60,7 +60,6 @@ * [envs.reach_pos_env](crazyflow/envs/reach_pos_env.md) * [envs.reach_vel_env](crazyflow/envs/reach_vel_env.md) * [envs.norm_actions_wrapper](crazyflow/envs/norm_actions_wrapper.md) -* [randomize](crazyflow/randomize/index.md) * [utils](crazyflow/utils.md) """ diff --git a/docs/user-guide/gymnasium-envs.md b/docs/user-guide/gymnasium-envs.md index 9ee0a4f..d0e8a38 100644 --- a/docs/user-guide/gymnasium-envs.md +++ b/docs/user-guide/gymnasium-envs.md @@ -41,7 +41,7 @@ All environments accept these common arguments: | `drone_model` | `"cf2x_L250"` | Drone configuration | | `freq` | 500 | Physics frequency, Hz | | `device` | `"cpu"` | `"cpu"` or `"gpu"` | -| `reset_randomization` | `None` | Optional `SimData → SimData` function applied at reset | +| `reset_randomization` | `None` | Optional `(SimData, mask) → SimData` function applied at reset | ## Action normalization @@ -60,18 +60,20 @@ env.close() ## Reset randomization -Pass a `reset_randomization` callable to vary initial conditions between episodes. The function receives `SimData` and a JAX random key and must return updated `SimData`: +Pass a `reset_randomization` callable to vary initial conditions between episodes. The function receives `SimData` and a boolean mask selecting the environments being reset, and must return updated `SimData`: ```{ .python notest } import jax -import jax.numpy as jnp from crazyflow.envs import ReachPosEnv from crazyflow.sim.data import SimData +from crazyflow.utils import leaf_replace -def randomize(data: SimData, key: jax.Array) -> SimData: - key, subkey = jax.random.split(key) +def randomize(data: SimData, mask: jax.Array | None) -> SimData: + key, subkey = jax.random.split(data.core.rng_key) + data = data.replace(core=data.core.replace(rng_key=key)) # Make sure to update the rng_key noise = jax.random.normal(subkey, data.states.pos.shape) * 0.05 - return data.replace(states=data.states.replace(pos=data.states.pos + noise)) + states = leaf_replace(data.states, mask, pos=data.states.pos + noise) + return data.replace(states=states) env = ReachPosEnv(num_envs=64, reset_randomization=randomize) ``` diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index 76253e4..7b3ecf1 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -196,44 +196,31 @@ sim.close() # close the viewer ## Domain randomization -Physical parameters can be randomized per-world using the `randomize` helpers: +Define physical-parameter randomization as a reset pipeline stage. Each stage receives the restored `SimData` and an optional world mask, and returns the randomized data: -```python +```{ .python notest } import jax -import numpy as np -from crazyflow.sim import Sim -from crazyflow.randomize import randomize_inertia, randomize_mass - -sim = Sim(n_worlds=4, n_drones=1) -sim.reset() - -nominal_mass = sim.data.params.mass -noise = jax.random.normal(jax.random.key(0), nominal_mass.shape) * 2e-3 -randomize_mass(sim, nominal_mass + noise) - -nominal_J = sim.data.params.J -J_noise = jax.random.normal(jax.random.key(1), nominal_J.shape) * 1e-6 -randomize_inertia(sim, nominal_J + J_noise) -``` +from jax import Array -To randomize only a subset of worlds, pass a boolean mask: - -```python -import jax -import numpy as np from crazyflow.sim import Sim -from crazyflow.randomize import randomize_mass +from crazyflow.sim.data import SimData +from crazyflow.utils import leaf_replace -sim = Sim(n_worlds=4, n_drones=1) -sim.reset() +@jax.jit +def randomize_mass(data: SimData, mask: Array | None = None) -> SimData: + key, mass_key = jax.random.split(data.core.rng_key) + mass = data.params.mass + jax.random.normal(mass_key, data.params.mass.shape) * 2e-3 + params = leaf_replace(data.params, mask, mass=mass) + return data.replace(params=params, core=data.core.replace(rng_key=key)) -import jax.numpy as jnp -mask = jnp.array([True, True, False, False]) # only worlds 0 and 1 -nominal_mass = sim.data.params.mass -noise = jax.random.normal(jax.random.key(0), nominal_mass.shape) * 2e-3 -randomize_mass(sim, nominal_mass + noise, mask=mask) +sim = Sim(n_worlds=4, n_drones=1) +sim.reset_pipeline = (randomize_mass,) +sim.build_reset_fn() +sim.reset() # randomizes every world ``` +Passing a boolean mask to `sim.reset(mask=mask)` randomizes only the worlds being reset. See the [domain randomization example](../examples/index.md#domain-randomization) for mass and inertia randomization in a complete simulation. + ## Next steps - [Functional API](functional-api.md) — run simulation inside `jax.jit` and `jax.grad` diff --git a/docs/user-guide/pipelines.md b/docs/user-guide/pipelines.md index ac8ed50..e808d26 100644 --- a/docs/user-guide/pipelines.md +++ b/docs/user-guide/pipelines.md @@ -21,7 +21,7 @@ print(sim.step_pipeline) ## The reset pipeline -`sim.reset_pipeline` is empty by default. When `sim.reset()` is called, it first restores `SimData` to the default state, then runs every function in the reset pipeline in order. The reset function signature is `(data: SimData, default_data: SimData, mask: Array | None) -> SimData`. +`sim.reset_pipeline` is empty by default. When `sim.reset()` is called, it first restores `SimData` to the default state, then runs every function in the reset pipeline in order. Each reset stage has the signature `(data: SimData, mask: Array | None) -> SimData`. Populate `sim.reset_pipeline` to add episode-level randomization without modifying the default state. @@ -36,36 +36,51 @@ To see how to modify the step pipeline with a stochastic disturbance, see the [D ## Modifying the reset pipeline -Add a function to the reset pipeline to vary initial conditions between episodes. The function receives the freshly-restored `data`, the `default_data` it was restored from, and an optional `mask` of worlds that were reset. +The [domain randomization example](../examples/index.md#domain-randomization) defines two reset stages that randomize mass and inertia. Each function receives the freshly restored `data` and an optional `mask` of worlds that were reset: ```{ .python notest } 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 jax import Array +from crazyflow.utils import leaf_replace -def randomize_initial_pos(data: SimData, default_data: SimData, mask: Array | None) -> SimData: - key, subkey = jax.random.split(data.core.rng_key) - noise = jax.random.normal(subkey, data.states.pos.shape) * 0.1 # ±10 cm - return data.replace( - states=data.states.replace(pos=default_data.states.pos + noise), - core=data.core.replace(rng_key=key), + +@jax.jit +def randomize_mass(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 + mass = ( + data.params.mass + + jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3 ) + return data.replace(params=leaf_replace(data.params, mask, mass=mass)) -sim = Sim(n_worlds=16) -sim.reset_pipeline = (randomize_initial_pos,) -sim.build_reset_fn() # recompile -sim.reset() -# Each of the 16 worlds now starts at a slightly different position -``` -Multiple stages can be chained; the output of each function is passed as input to the next: +@jax.jit +def randomize_inertia(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)) # Make sure to update the rng_key + J = ( + data.params.J + + jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8 + ) + return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J))) -```{ .python notest } -sim.reset_pipeline = (randomize_initial_pos, randomize_mass_fn, log_reset_fn) +sim = Sim(n_worlds=3, n_drones=4, control=Control.state) +sim.reset_pipeline = (randomize_mass, randomize_inertia) sim.build_reset_fn() + +mask = np.array([True, False, False]) # Only randomize the first world +sim.reset(mask=mask) # The mask is optional; omit it to reset and randomize all worlds ``` +Reset stages run in tuple order, with each stage receiving the output of the previous one. The mask ensures that parameter updates apply only to worlds selected by `sim.reset()`. + ## Removing a stage Remove any stage by excluding it from the tuple. A common case is removing the floor clip when computing gradients through a trajectory that starts high above the ground: @@ -80,7 +95,7 @@ sim.build_step_fn() ## Writing a custom stage -A step pipeline function must have the signature `(SimData) -> SimData`. A reset pipeline function must have the signature `(SimData, SimData, Array | None) -> SimData`. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled. +A step pipeline function must have the signature `(SimData) -> SimData`. A reset pipeline function must have the signature `(SimData, Array | None) -> SimData`. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled. ```{ .python notest } from crazyflow.sim.data import SimData diff --git a/examples/contacts.py b/examples/contacts/contacts.py similarity index 100% rename from examples/contacts.py rename to examples/contacts/contacts.py diff --git a/examples/crash.py b/examples/contacts/crash.py similarity index 100% rename from examples/crash.py rename to examples/contacts/crash.py diff --git a/examples/attitude.py b/examples/control/attitude.py similarity index 100% rename from examples/attitude.py rename to examples/control/attitude.py diff --git a/examples/change_pos.py b/examples/control/change_pos.py similarity index 100% rename from examples/change_pos.py rename to examples/control/change_pos.py diff --git a/examples/force_torque.py b/examples/control/force_torque.py similarity index 100% rename from examples/force_torque.py rename to examples/control/force_torque.py diff --git a/examples/hover.py b/examples/control/hover.py similarity index 100% rename from examples/hover.py rename to examples/control/hover.py diff --git a/examples/spiral.py b/examples/control/spiral.py similarity index 100% rename from examples/spiral.py rename to examples/control/spiral.py diff --git a/examples/figure8.py b/examples/environments/figure8.py similarity index 100% rename from examples/figure8.py rename to examples/environments/figure8.py diff --git a/examples/gymnasium_env.py b/examples/environments/gymnasium_env.py similarity index 100% rename from examples/gymnasium_env.py rename to examples/environments/gymnasium_env.py diff --git a/examples/cache.py b/examples/jax/cache.py similarity index 100% rename from examples/cache.py rename to examples/jax/cache.py diff --git a/examples/gradient.py b/examples/jax/gradient.py similarity index 100% rename from examples/gradient.py rename to examples/jax/gradient.py diff --git a/examples/plugins.py b/examples/plugins/action_delay.py similarity index 100% rename from examples/plugins.py rename to examples/plugins/action_delay.py diff --git a/examples/disturbance.py b/examples/plugins/disturbance.py similarity index 100% rename from examples/disturbance.py rename to examples/plugins/disturbance.py diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py new file mode 100644 index 0000000..252374e --- /dev/null +++ b/examples/plugins/estimation.py @@ -0,0 +1 @@ +# TODO add full state estimation pipeline example diff --git a/examples/plugins/randomize.py b/examples/plugins/randomize.py new file mode 100644 index 0000000..b882925 --- /dev/null +++ b/examples/plugins/randomize.py @@ -0,0 +1,63 @@ +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.utils import grid_2d, leaf_replace + + +@jax.jit +def randomize_mass(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 + mass = ( + data.params.mass + + jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3 + ) + return data.replace(params=leaf_replace(data.params, mask, mass=mass)) + + +@jax.jit +def randomize_inertia(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)) # Make sure to update the rng_key + J = ( + data.params.J + + jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8 + ) + return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J))) + + +def main(): + sim = Sim(n_worlds=3, n_drones=4, control=Control.state) + sim.reset_pipeline = (randomize_mass, randomize_inertia) + 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 anything + sim.reset(mask=mask) # Only reset the first world, the other two will stay the same + + sim.close() + + +if __name__ == "__main__": + main() diff --git a/examples/randomize.py b/examples/randomize.py deleted file mode 100644 index 304290a..0000000 --- a/examples/randomize.py +++ /dev/null @@ -1,43 +0,0 @@ -import jax -import numpy as np - -from crazyflow.control import Control -from crazyflow.randomize import randomize_inertia, randomize_mass -from crazyflow.sim import Sim -from crazyflow.utils import grid_2d - - -def main(): - sim = Sim(n_worlds=3, n_drones=4, control=Control.state) - sim.reset() - - duration = 5.0 - fps = 60 - - # Randomize the inertia and mass of the drones - mask = np.array([True, False, False]) # Only randomize the first world - mass = sim.data.params.mass - mass_rng = mass + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 1)) * 5e-3 - J = sim.data.params.J - J_rng = J + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 3, 3)) * 1e-6 - - randomize_mass(sim, mass_rng, mask) - # Note: The mask is optional. We can also randomize all worlds at once - randomize_mass(sim, mass_rng) - randomize_inertia(sim, J_rng, mask) - - cmd = np.zeros((sim.n_worlds, sim.n_drones, 13)) - cmd[..., 2] = 0.4 - cmd[..., :2] = grid_2d(sim.n_drones) * 0.25 - - # Simulate for 5 seconds. Each drone should behave slightly differently due to the randomization - 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() diff --git a/examples/cameras.py b/examples/rendering/cameras.py similarity index 100% rename from examples/cameras.py rename to examples/rendering/cameras.py diff --git a/examples/led_deck.py b/examples/rendering/led_deck.py similarity index 100% rename from examples/led_deck.py rename to examples/rendering/led_deck.py diff --git a/examples/raycasting.py b/examples/rendering/raycasting.py similarity index 100% rename from examples/raycasting.py rename to examples/rendering/raycasting.py diff --git a/examples/render.py b/examples/rendering/render.py similarity index 100% rename from examples/render.py rename to examples/rendering/render.py diff --git a/tests/integration/test_randomize.py b/tests/integration/test_randomize.py deleted file mode 100644 index 8c41e16..0000000 --- a/tests/integration/test_randomize.py +++ /dev/null @@ -1,50 +0,0 @@ -import jax -import numpy as np -import pytest - -from crazyflow.control import Control -from crazyflow.randomize import randomize_inertia, randomize_mass -from crazyflow.sim import Physics, Sim - - -@pytest.mark.integration -def test_randomize_mass(): - sim = Sim(n_worlds=2, n_drones=4, control=Control.state, physics=Physics.first_principles) - - control = np.zeros((sim.n_worlds, sim.n_drones, 13)) - control[:, :, 2] = 0.5 - - sim.state_control(control) - sim.step(sim.freq // sim.control_freq) - pos = sim.data.states.pos - - sim.reset() - mass_disturbance = jax.random.normal(jax.random.key(0), sim.data.params.mass.shape) * 1e-4 - randomize_mass(sim, sim.data.params.mass + mass_disturbance) - - sim.state_control(control) - sim.step(sim.freq // sim.control_freq) - pos_random = sim.data.states.pos - assert not np.all(pos == pos_random), "Mass randomization has no effect on dynamics" - - -@pytest.mark.integration -def test_randomize_inertia(): - sim = Sim(n_worlds=2, n_drones=4, control=Control.state, physics=Physics.first_principles) - steps = 50 # around 0.1s at 500Hz, not too long to still be in transient state - - control = np.zeros((sim.n_worlds, sim.n_drones, 13)) - control[:, :, :2] = 0.2 # Sideways motion to force tilt for inertia to have an effect - - sim.state_control(control) - sim.step(steps) - pos = sim.data.states.pos - - sim.reset() - J = sim.data.params.J + jax.random.normal(jax.random.key(0), sim.data.params.J.shape) * 1e-5 - randomize_inertia(sim, J) - - sim.state_control(control) - sim.step(steps) - pos_random = sim.data.states.pos - assert not np.all(pos == pos_random), "Inertia randomization has no effect on dynamics" diff --git a/tests/unit/test_randomize.py b/tests/unit/test_randomize.py deleted file mode 100644 index febdcd7..0000000 --- a/tests/unit/test_randomize.py +++ /dev/null @@ -1,96 +0,0 @@ -import jax -import jax.numpy as jnp -import numpy as np -import pytest -from flax.serialization import to_state_dict - -from crazyflow.control import Control -from crazyflow.randomize import randomize_inertia, randomize_mass -from crazyflow.sim import Sim - - -@pytest.mark.unit -@pytest.mark.parametrize("n_worlds", [1, 3]) -def test_randomize_mass(n_worlds: int): - sim = Sim(n_worlds=n_worlds, n_drones=4, control=Control.state) - - add_on_mass = np.random.uniform(-0.005, 0.005, size=(sim.n_worlds, sim.n_drones)) - masses = np.ones((sim.n_worlds, sim.n_drones)) * 0.025 - randomized_masses = masses + add_on_mass - randomize_mass(sim, randomized_masses) - - for k, v in to_state_dict(sim.data.params).items(): - default = getattr(sim.default_data.params, k) - if not isinstance(v, jnp.ndarray) or not isinstance(default, jnp.ndarray): - continue - assert v.shape == default.shape, f"{k} shape mismatch" - assert v.device == default.device, f"{k} device mismatch" - if k == "mass": - randomized_masses = randomized_masses.reshape(sim.n_worlds, sim.n_drones, 1) - assert np.allclose(v, randomized_masses), f"{k} value mismatch, {v - randomized_masses}" - - -@pytest.mark.unit -def test_randomize_mass_masked(): - sim = Sim(n_worlds=3, n_drones=4, control=Control.state) - - add_on_mass = np.random.uniform(-0.005, 0.005, size=(sim.n_worlds, sim.n_drones)) - masses = np.ones((sim.n_worlds, sim.n_drones)) * 0.025 - randomized_masses = masses + add_on_mass - randomized_masses = randomized_masses.reshape(sim.n_worlds, sim.n_drones, 1) - randomized_masses = jnp.array(randomized_masses, device=sim.device) - - mask = jnp.array([True, False, True]) - randomize_mass(sim, randomized_masses, mask) - - for k, v in to_state_dict(sim.data.params).items(): - default = getattr(sim.default_data.params, k) - if not isinstance(v, jnp.ndarray) or not isinstance(default, jnp.ndarray): - continue - assert v.shape == default.shape, f"{k} shape mismatch" - assert v.device == default.device, f"{k} device mismatch" - if k == "mass": - # Check that masked worlds match randomized masses - assert np.all(v[0] == randomized_masses[0]), "First world should be randomized" - assert np.all(v[1] == default[1]), "Second world should not be randomized" - assert np.all(v[2] == randomized_masses[2]), "Third world should be randomized" - - -@pytest.mark.unit -def test_randomize_inertia(): - sim = Sim(n_worlds=2, n_drones=4, control=Control.state) - - J_residual = jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 3, 3)) * 1e-5 - J = sim.data.params.J + J_residual - randomize_inertia(sim, J) - - for k, v in to_state_dict(sim.data.params).items(): - default = getattr(sim.default_data.params, k) - if not isinstance(v, jnp.ndarray) or not isinstance(default, jnp.ndarray): - continue - assert v.shape == default.shape, f"{k} shape mismatch" - assert v.device == default.device, f"{k} device mismatch" - if k == "J": - assert jnp.all(v == J), f"{k} value mismatch" - - -@pytest.mark.unit -def test_randomize_inertia_masked(): - sim = Sim(n_worlds=3, n_drones=4, control=Control.state) - - J_residual = jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 3, 3)) * 1e-5 - J = sim.data.params.J + J_residual - mask = jnp.array([True, False, True]) - randomize_inertia(sim, J, mask) - - for k, v in to_state_dict(sim.data.params).items(): - default = getattr(sim.default_data.params, k) - if not isinstance(v, jnp.ndarray) or not isinstance(default, jnp.ndarray): - continue - assert v.shape == default.shape, f"{k} shape mismatch" - assert v.device == default.device, f"{k} device mismatch" - if k == "J": - # Check that masked worlds match randomized inertias - assert jnp.all(v[0] == J[0]), "First world should be randomized" - assert jnp.all(v[1] == default[1]), "Second world should not be randomized" - assert jnp.all(v[2] == J[2]), "Third world should be randomized" From ac719ab3e2ae2649360d0dd24edb4bac28459051 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Fri, 12 Jun 2026 20:22:17 +0200 Subject: [PATCH 2/8] Add MPPI example --- docs/examples/index.md | 10 + examples/control/sampling.py | 348 +++++++++++++++++++++++++++++++++++ 2 files changed, 358 insertions(+) create mode 100644 examples/control/sampling.py diff --git a/docs/examples/index.md b/docs/examples/index.md index 5f555ea..7cd4774 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -28,6 +28,16 @@ Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses --- +## 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 a reduced dynamics model, 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. + +```bash +python examples/control/sampling.py +``` + +--- + ## 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. diff --git a/examples/control/sampling.py b/examples/control/sampling.py new file mode 100644 index 0000000..9dbc3c8 --- /dev/null +++ b/examples/control/sampling.py @@ -0,0 +1,348 @@ +"""Track a Lissajous curve with a simple sampling-based MPC controller.""" + +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 drone_models.core import load_params +from drone_models.transform import motor_force2rotor_vel +from jax import Array +from jax.lax import scan + +from crazyflow.control import Control +from crazyflow.sim import Physics, 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_MODEL = "cf21B_500" +DURATION = 18.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_PERCENTAGE = 0.01 # Percentage of samples to use for the mean update +MAX_CMD_ANGLE = np.deg2rad(60.0) + +# Lissajous reference configuration +REF_CENTER = jnp.array([0.0, 0.0, 1.0], dtype=jnp.float32) +REF_SCALE = jnp.array([1.0, 0.75, 0.0], dtype=jnp.float32) +REF_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 lissajous_reference(t: Array | float) -> dict[str, Array]: + """Return position, velocity, and yaw references at time ``t``.""" + t = jnp.asarray(t) + omega = 2.0 * jnp.pi / REF_PERIOD + + pos = jnp.stack( + ( + REF_CENTER[0] + REF_SCALE[0] * jnp.sin(omega * t), + REF_CENTER[1] + REF_SCALE[1] * jnp.sin(2.0 * omega * t), + jnp.broadcast_to(REF_CENTER[2], t.shape), + ), + axis=-1, + ) + vel = jnp.stack( + ( + REF_SCALE[0] * omega * jnp.cos(omega * t), + 2.0 * REF_SCALE[1] * omega * jnp.cos(2.0 * omega * t), + jnp.zeros_like(t), + ), + axis=-1, + ) + return {"pos": pos, "vel": vel, "yaw": jnp.zeros_like(t)} + + +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 model stores 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 = lissajous_reference(times) + costs, positions = rollout_fn(obs, candidates.transpose(1, 0, 2), references) + + n_elites = max(1, int(N_SAMPLES * ELITE_PERCENTAGE)) + 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_model=DRONE_MODEL, + physics=Physics.first_principles, + control=Control.attitude, + ) + sim.max_visual_geom = 100_000 # To be able to show all rollouts + sim.reset() + start_pos = lissajous_reference(0.0)["pos"] + drone_params = load_params("first_principles", DRONE_MODEL) + hover_thrust_value = np.asarray(drone_params["mass"] * 9.81, dtype=np.float32) + 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_model=DRONE_MODEL, + physics=Physics.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 + 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 = np.asarray(4.0 * drone_params["thrust_max"], dtype=np.float32) + 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, + ), + device=controller_device, + ) + 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, REF_PERIOD, round(REF_PERIOD * CTRL_FREQ) + 1) + reference = np.asarray(lissajous_reference(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 model. + "collective_thrust": thrust_estimate, + } + action, key, mean_controls, best_positions, sampled_positions = control( + t, obs, key, mean_controls, controller_fn, controller_device + ) + thrust_estimate += ( + drone_params["thrust_dyn_coef"] * (action[3] - thrust_estimate) / CTRL_FREQ + ) + sim.attitude_control(action[None, None]) + sim.step(sim.freq // CTRL_FREQ) + position_history.append(np.asarray(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() + + sim.close() + rollout_simulator.close() + + +if __name__ == "__main__": + main() From 32ca7a840cafb120831195e8210ce1d76e0497cd Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Fri, 12 Jun 2026 20:23:05 +0200 Subject: [PATCH 3/8] Add state estimation example --- examples/plugins/estimation.py | 221 ++++++++++++++++++++++++++++++++- 1 file changed, 220 insertions(+), 1 deletion(-) diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index 252374e..a57d8d2 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -1 +1,220 @@ -# TODO add full state estimation pipeline example +"""State estimation example using UWB position measurements.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import jax +import jax.numpy as jnp +import numpy as np +from drone_models.transform import motor_force2rotor_vel + +from crazyflow import Sim +from crazyflow.sim.visualize import draw_line, draw_points + +if TYPE_CHECKING: + from crazyflow.sim.data import SimData + +UWB_BASE_STATIONS = jnp.array( + [ + [-2.5, -2.5, 0.0], + [-2.5, 2.5, 0.0], + [2.5, -2.5, 0.0], + [2.5, 2.5, 0.0], + [-2.5, -2.5, 3.0], + [-2.5, 2.5, 3.0], + [2.5, -2.5, 3.0], + [2.5, 2.5, 3.0], + ] +) +HOVER_POSITION = jnp.array([0.0, 0.0, 1.0]) +TRAJECTORY_DURATION = 20.0 +RANGE_STD = 0.03 +RANGE_BIAS_MAX = 0.08 +# The constant-velocity model has no acceleration input, so it needs enough process noise to turn. +PROCESS_ACCEL_STD = 30.0 +MIN_ESTIMATOR_RANGE_STD = 1e-6 + + +def trajectory(t: float) -> np.ndarray: + """Return a slow figure-eight state command.""" + omega = 2 * np.pi / TRAJECTORY_DURATION + cmd = np.zeros((1, 1, 13)) + cmd[..., 0] = np.sin(omega * t) + cmd[..., 1] = 0.75 * np.sin(2 * omega * t) + cmd[..., 2] = 1.0 + cmd[..., 3] = omega * np.cos(omega * t) + cmd[..., 4] = 1.5 * omega * np.cos(2 * omega * t) + cmd[..., 6] = -(omega**2) * np.sin(omega * t) + cmd[..., 7] = -3 * omega**2 * np.sin(2 * omega * t) + return cmd + + +def simulate_uwb(data: SimData) -> SimData: + """Generate one range measurement for each UWB base station.""" + key, noise_key = jax.random.split(data.core.rng_key) + ranges = jnp.linalg.norm(data.states.pos[..., None, :] - UWB_BASE_STATIONS, axis=-1) + ranges += data.plugins["uwb_bias"] + ranges += data.plugins["range_std"] * jax.random.normal(noise_key, ranges.shape) + plugins = data.plugins | {"uwb_ranges": jnp.maximum(ranges, 0.0)} + return data.replace(plugins=plugins, core=data.core.replace(rng_key=key)) + + +def estimate_state(data: SimData) -> SimData: + """Run a constant-velocity EKF update using the UWB ranges.""" + dt = 1.0 / data.core.freq + estimate = data.plugins["estimate"] + covariance = data.plugins["covariance"] + + transition = jnp.eye(6).at[:3, 3:].set(jnp.eye(3) * dt) + accel_map = jnp.concat((jnp.eye(3) * (0.5 * dt**2), jnp.eye(3) * dt), axis=0) + process_covariance = PROCESS_ACCEL_STD**2 * accel_map @ accel_map.T + + predicted = estimate @ transition.T + predicted_covariance = transition @ covariance @ transition.T + process_covariance + + difference = predicted[..., None, :3] - UWB_BASE_STATIONS + predicted_ranges = jnp.linalg.norm(difference, axis=-1) + range_directions = difference / jnp.maximum(predicted_ranges[..., None], 1e-6) + n_ranges = UWB_BASE_STATIONS.shape[0] + measurement_jacobian = jnp.zeros((*predicted.shape[:-1], n_ranges, 6)) + measurement_jacobian = measurement_jacobian.at[..., :, :3].set(range_directions) + range_std = jnp.maximum(data.plugins["range_std"], MIN_ESTIMATOR_RANGE_STD) + measurement_variance = range_std**2 + measurement_covariance = jnp.eye(n_ranges) * measurement_variance + + innovation_covariance = ( + measurement_jacobian @ predicted_covariance @ jnp.swapaxes(measurement_jacobian, -1, -2) + + measurement_covariance + ) + covariance_times_jacobian = predicted_covariance @ jnp.swapaxes(measurement_jacobian, -1, -2) + gain = jnp.swapaxes( + jnp.linalg.solve(innovation_covariance, jnp.swapaxes(covariance_times_jacobian, -1, -2)), + -1, + -2, + ) + innovation = data.plugins["uwb_ranges"] - predicted_ranges + estimate = predicted + jnp.einsum("...ij,...j->...i", gain, innovation) + + identity = jnp.eye(6) + residual_map = identity - gain @ measurement_jacobian + covariance = residual_map @ predicted_covariance @ jnp.swapaxes(residual_map, -1, -2) + covariance += gain @ measurement_covariance @ jnp.swapaxes(gain, -1, -2) + covariance = 0.5 * (covariance + jnp.swapaxes(covariance, -1, -2)) + + return data.replace(plugins=data.plugins | {"estimate": estimate, "covariance": covariance}) + + +def use_estimate_for_control(data: SimData) -> SimData: + """Save the physical state and expose the estimate to the controllers.""" + estimate = data.plugins["estimate"] + plugins = data.plugins | {"ground_truth_state": data.states} + states = data.states.replace(pos=estimate[..., :3], vel=estimate[..., 3:]) + return data.replace(states=states, plugins=plugins) + + +def restore_ground_truth(data: SimData) -> SimData: + """Restore the physical state before evaluating and integrating the dynamics.""" + return data.replace(states=data.plugins["ground_truth_state"]) + + +def main(noisy: bool = False, render: bool = True) -> None: + """Run one perfect or noisy UWB estimation example.""" + name = "biased and noisy UWB measurements" if noisy else "perfect UWB measurements" + print(f"Running with {name}") + + sim = Sim(control="state", integrator="rk4", rng_key=42) + sim.max_visual_geom = 1000 + + hover_force = sim.data.params.mass * -sim.data.params.gravity_vec[2] / 4 + motor_forces = jnp.broadcast_to(hover_force, sim.data.states.rotor_vel.shape) + hover_rotor_vel = motor_force2rotor_vel(motor_forces, sim.data.params.rpm2thrust) + + states = sim.data.states.replace( + pos=trajectory(0.0)[..., :3], vel=trajectory(0.0)[..., 3:6], rotor_vel=hover_rotor_vel + ) + sim.data = sim.data.replace(states=states) + + key, bias_key = jax.random.split(sim.data.core.rng_key) + bias_shape = (sim.n_worlds, sim.n_drones, len(UWB_BASE_STATIONS)) + bias = ( + RANGE_BIAS_MAX * jax.random.uniform(bias_key, bias_shape) + if noisy + else jnp.zeros(bias_shape) + ) + range_std = RANGE_STD if noisy else 0.0 + estimate = jnp.concat((sim.data.states.pos, sim.data.states.vel), axis=-1) + covariance = jnp.broadcast_to(jnp.eye(6) * 0.1, (sim.n_worlds, sim.n_drones, 6, 6)) + plugins = { + "uwb_bias": bias, + "uwb_ranges": jnp.zeros_like(bias), + "range_std": jnp.asarray(range_std), + "estimate": estimate, + "covariance": covariance, + "ground_truth_state": sim.data.states, + } + sim.data = sim.data.replace( + plugins=sim.data.plugins | plugins, core=sim.data.core.replace(rng_key=key) + ) + + controllers = sim.step_pipeline[:-3] + integration = sim.step_pipeline[-3:] + sim.step_pipeline = ( + simulate_uwb, + estimate_state, + use_estimate_for_control, + *controllers, + restore_ground_truth, + *integration, + ) + sim.build_default_data() + sim.build_step_fn() + + duration = TRAJECTORY_DURATION + reference_times = np.linspace(0.0, duration, 200) + reference = np.array([trajectory(t)[0, 0, :3] for t in reference_times]) + estimation_errors = [] + tracking_errors = [] + fps = 60 + + for i in range(int(duration * sim.control_freq)): + t = i / sim.control_freq + cmd = trajectory(t) + sim.state_control(cmd) + sim.step(sim.freq // sim.control_freq) + + truth = np.asarray(sim.data.states.pos[0, 0]) + estimate = np.asarray(sim.data.plugins["estimate"][0, 0, :3]) + estimation_errors.append(np.linalg.norm(estimate - truth)) + tracking_errors.append(np.linalg.norm(truth - cmd[0, 0, :3])) + + if render and ((i * fps) % sim.control_freq) < fps: + truth = np.asarray(sim.data.states.pos[0, 0]) + estimate = np.asarray(sim.data.plugins["estimate"][0, 0, :3]) + anchors = np.asarray(UWB_BASE_STATIONS) + + draw_line( + sim, reference, rgba=np.array([0.4, 0.4, 0.4, 0.5]), start_size=0.5, end_size=0.5 + ) + for anchor in anchors: + draw_line( + sim, + np.stack((truth, anchor)), + rgba=np.array([0.2, 0.5, 1.0, 0.25]), + start_size=0.3, + end_size=0.3, + ) + draw_points(sim, anchors, rgba=np.array([0.1, 0.4, 1.0, 1.0]), size=0.05) + draw_points(sim, cmd[0, 0, :3], rgba=np.array([0.1, 0.8, 0.2, 1.0]), size=0.04) + draw_points(sim, estimate, rgba=np.array([1.0, 0.2, 0.1, 1.0]), size=0.04) + sim.render() + + sim.close() + estimation_rms = np.sqrt(np.mean(np.asarray(estimation_errors) ** 2)) + tracking_rms = np.sqrt(np.mean(np.asarray(tracking_errors) ** 2)) + print(f"Tracking RMS: {tracking_rms:.3f} m, estimation RMS: {estimation_rms:.3f} m") + + +if __name__ == "__main__": + main(noisy=False) + main(noisy=True) From 8a1eb269e1e70c9c0375b65b94d78cebfc7dd154 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Sat, 13 Jun 2026 18:37:39 +0200 Subject: [PATCH 4/8] Simplify UWB example --- examples/plugins/estimation.py | 82 ++++++++++++++++++++-------------- 1 file changed, 48 insertions(+), 34 deletions(-) diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index a57d8d2..8fb2e49 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -28,25 +28,25 @@ ] ) HOVER_POSITION = jnp.array([0.0, 0.0, 1.0]) -TRAJECTORY_DURATION = 20.0 +TRAJ_DURATION = 20.0 +TRAJ_SIZE = np.array([1.0, 0.75, 0.0]) RANGE_STD = 0.03 RANGE_BIAS_MAX = 0.08 -# The constant-velocity model has no acceleration input, so it needs enough process noise to turn. -PROCESS_ACCEL_STD = 30.0 +# Standard deviation of the unknown acceleration driving the constant-velocity process model. +PROCESS_NOISE_ACCEL_STD = 30.0 MIN_ESTIMATOR_RANGE_STD = 1e-6 def trajectory(t: float) -> np.ndarray: """Return a slow figure-eight state command.""" - omega = 2 * np.pi / TRAJECTORY_DURATION + omega = 2 * np.pi / TRAJ_DURATION cmd = np.zeros((1, 1, 13)) - cmd[..., 0] = np.sin(omega * t) - cmd[..., 1] = 0.75 * np.sin(2 * omega * t) - cmd[..., 2] = 1.0 - cmd[..., 3] = omega * np.cos(omega * t) - cmd[..., 4] = 1.5 * omega * np.cos(2 * omega * t) - cmd[..., 6] = -(omega**2) * np.sin(omega * t) - cmd[..., 7] = -3 * omega**2 * np.sin(2 * omega * t) + pos = HOVER_POSITION + TRAJ_SIZE * np.array([np.sin(omega * t), np.sin(2 * omega * t), 0.0]) + vel = TRAJ_SIZE * omega * np.array([np.cos(omega * t), 2 * np.cos(2 * omega * t), 0.0]) + acc = TRAJ_SIZE * omega**2 * np.array([-np.sin(omega * t), -4 * np.sin(2 * omega * t), 0.0]) + cmd[..., 0:3] = pos + cmd[..., 3:6] = vel + cmd[..., 6:9] = acc return cmd @@ -63,46 +63,60 @@ def simulate_uwb(data: SimData) -> SimData: def estimate_state(data: SimData) -> SimData: """Run a constant-velocity EKF update using the UWB ranges.""" dt = 1.0 / data.core.freq - estimate = data.plugins["estimate"] - covariance = data.plugins["covariance"] + # Predict: x_prior = F x, P_prior = F P F^T + Q + state = data.plugins["estimate"] + covariance = data.plugins["covariance"] + # Constant velocity: p_next = p + v * dt and v_next = v. transition = jnp.eye(6).at[:3, 3:].set(jnp.eye(3) * dt) - accel_map = jnp.concat((jnp.eye(3) * (0.5 * dt**2), jnp.eye(3) * dt), axis=0) - process_covariance = PROCESS_ACCEL_STD**2 * accel_map @ accel_map.T - - predicted = estimate @ transition.T + # Model unknown acceleration a as process noise: + # x_next = F x + acceleration_to_state * a, where delta_p = 0.5*a*dt^2 and delta_v = a*dt. + acceleration_to_state = jnp.concat((jnp.eye(3) * (0.5 * dt**2), jnp.eye(3) * dt), axis=0) + process_covariance = ( + PROCESS_NOISE_ACCEL_STD**2 * acceleration_to_state @ acceleration_to_state.T + ) + predicted_state = state @ transition.T predicted_covariance = transition @ covariance @ transition.T + process_covariance - difference = predicted[..., None, :3] - UWB_BASE_STATIONS - predicted_ranges = jnp.linalg.norm(difference, axis=-1) - range_directions = difference / jnp.maximum(predicted_ranges[..., None], 1e-6) + # Linearize ranges: z_prior = h(x_prior), H = dh/dx at x_prior + anchor_offsets = predicted_state[..., None, :3] - UWB_BASE_STATIONS + predicted_ranges = jnp.linalg.norm(anchor_offsets, axis=-1) + range_directions = anchor_offsets / jnp.maximum(predicted_ranges[..., None], 1e-6) n_ranges = UWB_BASE_STATIONS.shape[0] - measurement_jacobian = jnp.zeros((*predicted.shape[:-1], n_ranges, 6)) + measurement_jacobian = jnp.zeros((*predicted_state.shape[:-1], n_ranges, 6)) measurement_jacobian = measurement_jacobian.at[..., :, :3].set(range_directions) range_std = jnp.maximum(data.plugins["range_std"], MIN_ESTIMATOR_RANGE_STD) - measurement_variance = range_std**2 - measurement_covariance = jnp.eye(n_ranges) * measurement_variance + measurement_covariance = jnp.eye(n_ranges) * range_std**2 + measurement_jacobian_transpose = jnp.swapaxes(measurement_jacobian, -1, -2) + # Innovation and gain: y = z - z_prior, S = H P_prior H^T + R, K = P_prior H^T S^-1 + innovation = data.plugins["uwb_ranges"] - predicted_ranges innovation_covariance = ( - measurement_jacobian @ predicted_covariance @ jnp.swapaxes(measurement_jacobian, -1, -2) + measurement_jacobian @ predicted_covariance @ measurement_jacobian_transpose + measurement_covariance ) - covariance_times_jacobian = predicted_covariance @ jnp.swapaxes(measurement_jacobian, -1, -2) - gain = jnp.swapaxes( + covariance_times_jacobian = predicted_covariance @ measurement_jacobian_transpose + # K = P_prior H^T S^-1. Solve S K^T = (P_prior H^T)^T instead of forming S^-1; + # a linear solve is more numerically stable and directly computes the same result. + kalman_gain = jnp.swapaxes( jnp.linalg.solve(innovation_covariance, jnp.swapaxes(covariance_times_jacobian, -1, -2)), -1, -2, ) - innovation = data.plugins["uwb_ranges"] - predicted_ranges - estimate = predicted + jnp.einsum("...ij,...j->...i", gain, innovation) + # Correct: x = x_prior + K y + corrected_state = predicted_state + (kalman_gain @ innovation[..., None])[..., 0] + + # Joseph form: P = (I - K H) P_prior (I - K H)^T + K R K^T identity = jnp.eye(6) - residual_map = identity - gain @ measurement_jacobian - covariance = residual_map @ predicted_covariance @ jnp.swapaxes(residual_map, -1, -2) - covariance += gain @ measurement_covariance @ jnp.swapaxes(gain, -1, -2) - covariance = 0.5 * (covariance + jnp.swapaxes(covariance, -1, -2)) + residual_map = identity - kalman_gain @ measurement_jacobian + corrected_covariance = residual_map @ predicted_covariance @ jnp.swapaxes( + residual_map, -1, -2 + ) + kalman_gain @ measurement_covariance @ jnp.swapaxes(kalman_gain, -1, -2) + corrected_covariance = 0.5 * (corrected_covariance + jnp.swapaxes(corrected_covariance, -1, -2)) - return data.replace(plugins=data.plugins | {"estimate": estimate, "covariance": covariance}) + plugins = data.plugins | {"estimate": corrected_state, "covariance": corrected_covariance} + return data.replace(plugins=plugins) def use_estimate_for_control(data: SimData) -> SimData: @@ -170,7 +184,7 @@ def main(noisy: bool = False, render: bool = True) -> None: sim.build_default_data() sim.build_step_fn() - duration = TRAJECTORY_DURATION + duration = TRAJ_DURATION reference_times = np.linspace(0.0, duration, 200) reference = np.array([trajectory(t)[0, 0, :3] for t in reference_times]) estimation_errors = [] From 7ea6e621dc7e48f7c42cac01388450ef0d7405fd Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Mon, 15 Jun 2026 11:30:13 +0200 Subject: [PATCH 5/8] Add config example --- examples/rendering/cam_config.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 examples/rendering/cam_config.py diff --git a/examples/rendering/cam_config.py b/examples/rendering/cam_config.py new file mode 100644 index 0000000..f70d64d --- /dev/null +++ b/examples/rendering/cam_config.py @@ -0,0 +1,30 @@ +"""Simple example on how to change the camera configuration for rendering.""" + +import numpy as np + +from crazyflow.sim import Sim + +cam_config = {"distance": 0.8, "elevation": -45.0, "azimuth": -135.0, "lookat": [0.0, 0.0, 0.0]} + + +def main(cam_config: dict | None = None): + sim = Sim(control="state") + sim.reset() + + duration = 5.0 + fps = 60 + + cmd = np.zeros((sim.n_worlds, sim.n_drones, 13)) + cmd[..., :3] = 0.2 + + 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(cam_config=cam_config) + sim.close() + + +if __name__ == "__main__": + main(cam_config=None) + main(cam_config=cam_config) From 8ca44f0c3a07dac79155a00752680e6056b90ee0 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Mon, 15 Jun 2026 11:30:37 +0200 Subject: [PATCH 6/8] Improve clarity --- examples/plugins/estimation.py | 63 ++++++++++++++++------------------ examples/rendering/cameras.py | 21 +++++++++--- 2 files changed, 46 insertions(+), 38 deletions(-) diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index 8fb2e49..5c2c078 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -27,23 +27,19 @@ [2.5, 2.5, 3.0], ] ) -HOVER_POSITION = jnp.array([0.0, 0.0, 1.0]) -TRAJ_DURATION = 20.0 -TRAJ_SIZE = np.array([1.0, 0.75, 0.0]) -RANGE_STD = 0.03 -RANGE_BIAS_MAX = 0.08 -# Standard deviation of the unknown acceleration driving the constant-velocity process model. -PROCESS_NOISE_ACCEL_STD = 30.0 -MIN_ESTIMATOR_RANGE_STD = 1e-6 +RANGE_STD = 0.03 # measurement std dev [m] +RANGE_BIAS_MAX = 0.08 # measurement bias, uniformly distributed in [0, RANGE_BIAS_MAX] [m] -def trajectory(t: float) -> np.ndarray: +def trajectory(t: float, t_total: float = 20.0) -> np.ndarray: """Return a slow figure-eight state command.""" - omega = 2 * np.pi / TRAJ_DURATION + center = np.array([0.0, 0.0, 1.0]) + size = np.array([1.0, 0.75, 0.0]) + omega = 2 * np.pi / t_total cmd = np.zeros((1, 1, 13)) - pos = HOVER_POSITION + TRAJ_SIZE * np.array([np.sin(omega * t), np.sin(2 * omega * t), 0.0]) - vel = TRAJ_SIZE * omega * np.array([np.cos(omega * t), 2 * np.cos(2 * omega * t), 0.0]) - acc = TRAJ_SIZE * omega**2 * np.array([-np.sin(omega * t), -4 * np.sin(2 * omega * t), 0.0]) + pos = center + size * np.array([np.sin(omega * t), np.sin(2 * omega * t), 0.0]) + vel = size * omega * np.array([np.cos(omega * t), 2 * np.cos(2 * omega * t), 0.0]) + acc = size * omega**2 * np.array([-np.sin(omega * t), -4 * np.sin(2 * omega * t), 0.0]) cmd[..., 0:3] = pos cmd[..., 3:6] = vel cmd[..., 6:9] = acc @@ -72,9 +68,8 @@ def estimate_state(data: SimData) -> SimData: # Model unknown acceleration a as process noise: # x_next = F x + acceleration_to_state * a, where delta_p = 0.5*a*dt^2 and delta_v = a*dt. acceleration_to_state = jnp.concat((jnp.eye(3) * (0.5 * dt**2), jnp.eye(3) * dt), axis=0) - process_covariance = ( - PROCESS_NOISE_ACCEL_STD**2 * acceleration_to_state @ acceleration_to_state.T - ) + process_covariance = 30.0**2 * acceleration_to_state @ acceleration_to_state.T + # Standard deviation of 30 of the unknown acceleration driving the constant-velocity process. predicted_state = state @ transition.T predicted_covariance = transition @ covariance @ transition.T + process_covariance @@ -85,8 +80,7 @@ def estimate_state(data: SimData) -> SimData: n_ranges = UWB_BASE_STATIONS.shape[0] measurement_jacobian = jnp.zeros((*predicted_state.shape[:-1], n_ranges, 6)) measurement_jacobian = measurement_jacobian.at[..., :, :3].set(range_directions) - range_std = jnp.maximum(data.plugins["range_std"], MIN_ESTIMATOR_RANGE_STD) - measurement_covariance = jnp.eye(n_ranges) * range_std**2 + measurement_covariance = jnp.eye(n_ranges) * data.plugins["range_std"] ** 2 measurement_jacobian_transpose = jnp.swapaxes(measurement_jacobian, -1, -2) # Innovation and gain: y = z - z_prior, S = H P_prior H^T + R, K = P_prior H^T S^-1 @@ -138,25 +132,23 @@ def main(noisy: bool = False, render: bool = True) -> None: print(f"Running with {name}") sim = Sim(control="state", integrator="rk4", rng_key=42) - sim.max_visual_geom = 1000 - hover_force = sim.data.params.mass * -sim.data.params.gravity_vec[2] / 4 motor_forces = jnp.broadcast_to(hover_force, sim.data.states.rotor_vel.shape) hover_rotor_vel = motor_force2rotor_vel(motor_forces, sim.data.params.rpm2thrust) - - states = sim.data.states.replace( - pos=trajectory(0.0)[..., :3], vel=trajectory(0.0)[..., 3:6], rotor_vel=hover_rotor_vel + sim.data = sim.data.replace( + states=sim.data.states.replace( + pos=trajectory(0.0)[..., :3], vel=trajectory(0.0)[..., 3:6], rotor_vel=hover_rotor_vel + ) ) - sim.data = sim.data.replace(states=states) key, bias_key = jax.random.split(sim.data.core.rng_key) bias_shape = (sim.n_worlds, sim.n_drones, len(UWB_BASE_STATIONS)) bias = ( - RANGE_BIAS_MAX * jax.random.uniform(bias_key, bias_shape) + jax.random.uniform(bias_key, bias_shape, minval=0.0, maxval=RANGE_BIAS_MAX) if noisy else jnp.zeros(bias_shape) ) - range_std = RANGE_STD if noisy else 0.0 + range_std = RANGE_STD if noisy else 1e-6 # Nonzero for numerical stability estimate = jnp.concat((sim.data.states.pos, sim.data.states.vel), axis=-1) covariance = jnp.broadcast_to(jnp.eye(6) * 0.1, (sim.n_worlds, sim.n_drones, 6, 6)) plugins = { @@ -184,16 +176,16 @@ def main(noisy: bool = False, render: bool = True) -> None: sim.build_default_data() sim.build_step_fn() - duration = TRAJ_DURATION + duration = 20.0 reference_times = np.linspace(0.0, duration, 200) - reference = np.array([trajectory(t)[0, 0, :3] for t in reference_times]) + reference = np.array([trajectory(t, t_total=duration)[0, 0, :3] for t in reference_times]) estimation_errors = [] tracking_errors = [] fps = 60 for i in range(int(duration * sim.control_freq)): t = i / sim.control_freq - cmd = trajectory(t) + cmd = trajectory(t, t_total=duration) sim.state_control(cmd) sim.step(sim.freq // sim.control_freq) @@ -205,12 +197,11 @@ def main(noisy: bool = False, render: bool = True) -> None: if render and ((i * fps) % sim.control_freq) < fps: truth = np.asarray(sim.data.states.pos[0, 0]) estimate = np.asarray(sim.data.plugins["estimate"][0, 0, :3]) - anchors = np.asarray(UWB_BASE_STATIONS) draw_line( sim, reference, rgba=np.array([0.4, 0.4, 0.4, 0.5]), start_size=0.5, end_size=0.5 ) - for anchor in anchors: + for anchor in UWB_BASE_STATIONS: draw_line( sim, np.stack((truth, anchor)), @@ -218,10 +209,16 @@ def main(noisy: bool = False, render: bool = True) -> None: start_size=0.3, end_size=0.3, ) - draw_points(sim, anchors, rgba=np.array([0.1, 0.4, 1.0, 1.0]), size=0.05) + draw_points(sim, UWB_BASE_STATIONS, rgba=np.array([0.1, 0.4, 1.0, 1.0]), size=0.05) draw_points(sim, cmd[0, 0, :3], rgba=np.array([0.1, 0.8, 0.2, 1.0]), size=0.04) draw_points(sim, estimate, rgba=np.array([1.0, 0.2, 0.1, 1.0]), size=0.04) - sim.render() + cam_config = { + "distance": 3.0, + "elevation": -45.0, + "azimuth": 90.0, + "lookat": [0.0, 0.0, 1.0], + } + sim.render(cam_config=cam_config) sim.close() estimation_rms = np.sqrt(np.mean(np.asarray(estimation_errors) ** 2)) diff --git a/examples/rendering/cameras.py b/examples/rendering/cameras.py index 547b820..ef68274 100644 --- a/examples/rendering/cameras.py +++ b/examples/rendering/cameras.py @@ -1,3 +1,5 @@ +"""Example showing how to change the used camera and how to extract the pixel information.""" + import time import matplotlib.pyplot as plt @@ -71,8 +73,8 @@ def main(show_plot: bool = False, save_plot: bool = False): 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 = 8 - fps = 25 + duration = 5 + fps = 50 timings = [] # Set up matplotlib rendering @@ -95,6 +97,12 @@ def update_frame(_): # noqa: ANN202 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 drone_models. 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" ) @@ -108,16 +116,19 @@ def update_frame(_): # noqa: ANN202 im2.set_clim(np.nanmin(depth), np.nanmax(depth)) return im1, im2 - anim = animation.FuncAnimation(fig, update_frame, frames=int(duration * fps), blit=True) + anim = animation.FuncAnimation( + fig, update_frame, frames=int(duration * fps), interval=1000 / fps, blit=True, repeat=False + ) if show_plot: - plt.show() # this is slow + 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}ms, eqivalent to {1 / t_mean}fps") + 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__": From 0b908bf17304c60f9f034ed0dd76a6dc2303c196 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Mon, 15 Jun 2026 11:52:46 +0200 Subject: [PATCH 7/8] Add cam_config --- examples/control/sampling.py | 9 ++++++++- examples/plugins/estimation.py | 15 ++++++++------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/examples/control/sampling.py b/examples/control/sampling.py index 9dbc3c8..b0c2da2 100644 --- a/examples/control/sampling.py +++ b/examples/control/sampling.py @@ -338,7 +338,14 @@ def main() -> None: 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() + 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() diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index 5c2c078..71ec859 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -212,13 +212,14 @@ def main(noisy: bool = False, render: bool = True) -> None: draw_points(sim, UWB_BASE_STATIONS, rgba=np.array([0.1, 0.4, 1.0, 1.0]), size=0.05) draw_points(sim, cmd[0, 0, :3], rgba=np.array([0.1, 0.8, 0.2, 1.0]), size=0.04) draw_points(sim, estimate, rgba=np.array([1.0, 0.2, 0.1, 1.0]), size=0.04) - cam_config = { - "distance": 3.0, - "elevation": -45.0, - "azimuth": 90.0, - "lookat": [0.0, 0.0, 1.0], - } - sim.render(cam_config=cam_config) + sim.render( + cam_config={ + "distance": 3.0, + "elevation": -45.0, + "azimuth": 90.0, + "lookat": [0.0, 0.0, 1.0], + } + ) sim.close() estimation_rms = np.sqrt(np.mean(np.asarray(estimation_errors) ** 2)) From b94234ef79a1879848e13d6aaa9bba98a1645f26 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Mon, 15 Jun 2026 17:09:12 +0200 Subject: [PATCH 8/8] Implemented comments --- docs/user-guide/oo-api.md | 2 +- examples/control/sampling.py | 34 ++++++++++++++-------------------- examples/plugins/estimation.py | 5 ++++- examples/plugins/randomize.py | 12 ++++-------- 4 files changed, 23 insertions(+), 30 deletions(-) diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index 7b3ecf1..b0cfc6d 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -198,7 +198,7 @@ sim.close() # close the viewer Define physical-parameter randomization as a reset pipeline stage. Each stage receives the restored `SimData` and an optional world mask, and returns the randomized data: -```{ .python notest } +```python import jax from jax import Array diff --git a/examples/control/sampling.py b/examples/control/sampling.py index b0c2da2..5e5e539 100644 --- a/examples/control/sampling.py +++ b/examples/control/sampling.py @@ -1,4 +1,7 @@ -"""Track a Lissajous curve with a simple sampling-based MPC controller.""" +"""Track a Lissajous curve with a simple sampling-based MPC controller. + +WARNING: This is an advanced example meant for advanced Crazyflow users. +""" import os from collections import deque @@ -39,7 +42,7 @@ 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_PERCENTAGE = 0.01 # Percentage of samples to use for the mean update +ELITE_FRACTION = 0.01 # Fraction of samples in (0,1] to use for the mean update MAX_CMD_ANGLE = np.deg2rad(60.0) # Lissajous reference configuration @@ -65,23 +68,14 @@ def lissajous_reference(t: Array | float) -> dict[str, Array]: """Return position, velocity, and yaw references at time ``t``.""" t = jnp.asarray(t) omega = 2.0 * jnp.pi / REF_PERIOD - - pos = jnp.stack( - ( - REF_CENTER[0] + REF_SCALE[0] * jnp.sin(omega * t), - REF_CENTER[1] + REF_SCALE[1] * jnp.sin(2.0 * omega * t), - jnp.broadcast_to(REF_CENTER[2], t.shape), - ), - axis=-1, - ) - vel = jnp.stack( - ( - REF_SCALE[0] * omega * jnp.cos(omega * t), - 2.0 * REF_SCALE[1] * omega * jnp.cos(2.0 * omega * t), - jnp.zeros_like(t), - ), - axis=-1, - ) + pos_x = REF_CENTER[0] + REF_SCALE[0] * jnp.sin(omega * t) + pos_y = REF_CENTER[1] + REF_SCALE[1] * jnp.sin(2.0 * omega * t) + pos_z = jnp.broadcast_to(REF_CENTER[2], t.shape) + pos = jnp.stack((pos_x, pos_y, pos_z), axis=-1) + vel_x = REF_SCALE[0] * omega * jnp.cos(omega * t) + vel_y = 2.0 * REF_SCALE[1] * omega * jnp.cos(2.0 * omega * t) + vel_z = jnp.zeros_like(t) + vel = jnp.stack((vel_x, vel_y, vel_z), axis=-1) return {"pos": pos, "vel": vel, "yaw": jnp.zeros_like(t)} @@ -182,7 +176,7 @@ def update_controller( references = lissajous_reference(times) costs, positions = rollout_fn(obs, candidates.transpose(1, 0, 2), references) - n_elites = max(1, int(N_SAMPLES * ELITE_PERCENTAGE)) + 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) diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index 71ec859..7f83b81 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -1,4 +1,7 @@ -"""State estimation example using UWB position measurements.""" +"""State estimation example using UWB position measurements. + +WARNING: This is an advanced example meant for advanced Crazyflow users. +""" from __future__ import annotations diff --git a/examples/plugins/randomize.py b/examples/plugins/randomize.py index b882925..710df8d 100644 --- a/examples/plugins/randomize.py +++ b/examples/plugins/randomize.py @@ -13,10 +13,8 @@ def randomize_mass(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 - mass = ( - data.params.mass - + jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3 - ) + dist = jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3 + mass = data.params.mass + dist return data.replace(params=leaf_replace(data.params, mask, mass=mass)) @@ -24,10 +22,8 @@ def randomize_mass(data: SimData, mask: Array | None = None) -> SimData: def randomize_inertia(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)) # Make sure to update the rng_key - J = ( - data.params.J - + jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8 - ) + dist = jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8 + J = data.params.J + dist return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J)))