Crazyflow

https://learnsyslab.github.io/crazyflow
Admin
Crazyflow is a research simulator for Crazyflie-style quadrotors, built on JAX, that runs millions...
Tokens:121,981
Snippets:714
Trust Score:9.7
Update:3 months ago
Show doc for...
Tokens:
Raw
### Run Camera Rendering Example

Source: https://learnsyslab.github.io/crazyflow/examples

This snippet shows the command to execute the Python script for the camera rendering example.

```bash
python examples/rendering/cameras.py

```

--------------------------------

### Install CrazyFlow from Source with Pixi

Source: https://learnsyslab.github.io/crazyflow/get-started/installation

Clone the repository and use Pixi to set up the development environment for installing CrazyFlow from source.

```bash
git clone https://github.com/learnsyslab/crazyflow.git
cd crazyflow
pixi shell

```

--------------------------------

### Install CrazyFlow from Source with Pixi for Testing

Source: https://learnsyslab.github.io/crazyflow/get-started/installation

Clone the repository and use Pixi to set up the development environment, including test dependencies, for installing CrazyFlow from source.

```bash
git clone https://github.com/learnsyslab/crazyflow.git
cd crazyflow
pixi shell -e tests

```

--------------------------------

### Install CrazyFlow with Pip

Source: https://learnsyslab.github.io/crazyflow/get-started/installation

Use this command for a standard installation of CrazyFlow.

```bash
pip install crazyflow

```

--------------------------------

### Verify CrazyFlow Installation

Source: https://learnsyslab.github.io/crazyflow/get-started/installation

Run this Python command to verify that CrazyFlow has been installed correctly and can initialize a Sim object.

```python
python -c "from crazyflow.sim import Sim; sim = Sim(); sim.reset(); print('OK')"

```

--------------------------------

### Minimal Simulation Example

Source: https://learnsyslab.github.io/crazyflow

Sets up a minimal simulation with one world and one drone, controlling its state to hover at a specified altitude. Retrieves the position of the drone after stepping the simulation.

```python
import numpy as np
from crazyflow.sim import Sim
from crazyflow.control import Control

sim = Sim(n_worlds=1, n_drones=1, control=Control.state)
sim.reset()

# State command: [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]
cmd = np.zeros((1, 1, 13), dtype=np.float32)
cmd[0, 0, 2] = 0.5  # hover at 0.5 m

sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)

pos = sim.data.states.pos[0, 0]  # shape (3,) — position of world 0, drone 0

```

--------------------------------

### Initialize Simulation with First Principles Dynamics and Rotor Velocity Control

Source: https://learnsyslab.github.io/crazyflow/user-guide/dynamics

Configure the simulator to use 'first_principles' dynamics along with 'rotor_vel' control mode. This setup is compatible with first-principles dynamics.

```python
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(
    dynamics=Dynamics.first_principles,
    control=Control.rotor_vel,
)
sim.reset()

```

--------------------------------

### Force-Torque Control Example

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Demonstrates direct force and torque input for testing dynamics or custom controllers. This requires the `Dynamics.first_principles` model.

```python
import numpy as np
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(n_worlds=1, n_drones=1, control=Control.force_torque, dynamics=Dynamics.first_principles)
sim.reset()

# [collective_force_N, torque_x_Nm, torque_y_Nm, torque_z_Nm]
cmd = np.zeros((1, 1, 4), dtype=np.float32)
cmd[0, 0, 0] = float(sim.data.params.mass[0, 0, 0]) * 9.81  # hover force

sim.force_torque_control(cmd)
sim.step(1)

```

--------------------------------

### Create Default Simulation Core Parameters

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/sim/data

Initializes a `SimCore` object with essential simulation parameters. This method handles the setup of simulation frequency, world/drone counts, and random number generation.

```python
@staticmethod
def create(
    freq: int,
    n_worlds: int,
    n_drones: int,
    drone_mocap_ids: Array,
    rng_key: int | Array,
    device: Device,
) -> SimCore:
    """Create a default set of core simulation parameters."""
    steps = jnp.zeros((n_worlds, 1), dtype=jnp.int32, device=device)
    if isinstance(rng_key, int):  # Only convert to an PRNG key if its not already one
        rng_key = jax.random.key(rng_key)
    rng_key = jax.device_put(rng_key, device)
    return SimCore(
        freq=freq,
        device=device,
        steps=steps,
        n_worlds=n_worlds,
        n_drones=n_drones,
        drone_mocap_ids=jnp.array(drone_mocap_ids, dtype=jnp.int32, device=device),
        rng_key=rng_key,
        mjx_synced=jnp.array(False, dtype=jnp.bool_, device=device),
    )
```

--------------------------------

### Install CrazyFlow with GPU Support

Source: https://learnsyslab.github.io/crazyflow/get-started/installation

Install CrazyFlow with the 'gpu' extra to enable GPU execution. Ensure your system meets the requirements for GPU support.

```bash
pip install "crazyflow[gpu]"

```

--------------------------------

### Gradient Descent Through Dynamics Simulation

Source: https://learnsyslab.github.io/crazyflow/examples

Demonstrates gradient descent through a JAX-based simulator. The example shows how to differentiate through the simulation dynamics to optimize control commands. It highlights the importance of initial conditions to avoid zero gradients from floor clipping.

```python
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, 0, 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)}
Gradient: {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()

```

--------------------------------

### Offscreen Rendering with FPV Camera

Source: https://learnsyslab.github.io/crazyflow/examples

This example shows how to set up offscreen rendering using the FPV camera, retrieve RGB and depth data, and visualize it with Matplotlib. It includes adding a custom 3D object to the simulation environment.

```python
"""Example showing how to change the used camera and how to extract the pixel information."""

import time

import matplotlib.pyplot as plt
import mujoco
import numpy as np
from matplotlib import animation

from crazyflow.control import Control
from crazyflow.dynamics import Dynamics
from crazyflow.sim import Sim
from crazyflow.sim.integration import Integrator


def control(t: float, t_tot: float) -> np.ndarray:
    phi = 2 * np.pi * t / t_tot + np.pi
    circle = np.array([np.cos(phi), np.sin(phi)])
    cmd = np.zeros((1, 1, 13))
    cmd[..., :2] = circle  # xy
    cmd[..., 2] = 0.1 + 0.5 * t / t_tot  # z
    cmd[..., -4] = 1.9 * np.pi * t / t_tot  # yaw

    return cmd


def add_smiley(sim: Sim):
    # Add 3d object to sim
    # create box spec from an XML string
    box_xml = """
    <mujoco model="box_model">
      <worldbody>
        <body name="cube" pos="0 0 0">
          <geom type="box" size="0.05 0.05 0.05" rgba="0.8 0.4 0.2 1"/>
        </body>
      </worldbody>
    </mujoco>
    """
    box_spec = mujoco.MjSpec.from_string(box_xml)
    frame = sim.spec.worldbody.add_frame()
    boxes = [
        # eyes
        ((0.0, -0.15, 0.6), (1, 0, 0, 0)),
        ((0.0, 0.15, 0.6), (1, 0, 0, 0)),
        # mouth
        ((0.0, -0.2, 0.4), (1, 0, 0, 0)),
        ((0.0, 0.2, 0.4), (1, 0, 0, 0)),
        ((0.0, -0.1, 0.3), (1, 0, 0, 0)),
        ((0.0, 0.0, 0.3), (1, 0, 0, 0)),
        ((0.0, 0.1, 0.3), (1, 0, 0, 0)),
    ]
    for i, x in enumerate(boxes):
        box_body = box_spec.body("cube")
        box = frame.attach_body(box_body, "", f":{i}")
        box.pos = x[0]
        box.quat = x[1]
    sim.build_mjx()
    sim.build_reset_fn()


def main(show_plot: bool = False, save_plot: bool = False):
    """Example showing the rendering feature and saving a gif via FuncAnimation."""
    # Setup sim
    sim = Sim(
        n_drones=1,
        control=Control.state,
        integrator=Integrator.rk4,
        dynamics=Dynamics.first_principles,
        drone="cf2x_T350",
    )
    add_smiley(sim)
    sim.reset()
    pos = sim.data.states.pos.at[...].set([-1, 0, 0])
    states = sim.data.states.replace(pos=pos)
    sim.data = sim.data.replace(states=states)
    duration = 5
    fps = 50
    timings = []

    # Set up matplotlib rendering
    resolution = (160, 120)
    rgb = np.zeros((resolution[1], resolution[0], 3))
    d = np.zeros((resolution[1], resolution[0]))
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    im1 = ax1.imshow(rgb)
    ax1.set_title("RGB")
    ax1.axis("off")
    im2 = ax2.imshow(d, cmap="viridis")
    ax2.set_title("Depth")
    ax2.axis("off")
    fig.tight_layout()

    # Animation setup
    def update_frame(_):  # noqa: ANN202
        t = sim.data.core.steps[0, 0] / sim.freq
        sim.state_control(control(t, duration))
        sim.step(sim.freq // fps)

        t1 = time.perf_counter()
        # mode: Either "human" for the regular window, "rgb_array" for an RGB array,
        #       "depth_array" for a depth array, or "rgbd_tuple" for both at the same time.
        # camera: The name or id of the camera. The names are specified in the corresponding
        #         xml file in crazyflow/drones. For example, "fpv_cam:0" is the first-person view
        #         camera of the first drone, "track_cam:0" is the tracking camera of the first
        #         drone. Id -1 is the global camera.
        rgbd = sim.render(
            width=resolution[0], height=resolution[1], mode="rgbd_tuple", camera="fpv_cam:0"
        )
        t2 = time.perf_counter()
        timings.append(t2 - t1)
        if rgbd is None:
            return im1, im2
        rgb, depth = rgbd
        im1.set_data(rgb)
        im2.set_data(depth)
        im2.set_clim(np.nanmin(depth), np.nanmax(depth))
        return im1, im2

    anim = animation.FuncAnimation(
        fig, update_frame, frames=int(duration * fps), interval=1000 / fps, blit=True, repeat=False
    )
    if show_plot:
        plt.show()
    if save_plot:
        anim.save("cameras.gif", writer="pillow", fps=fps)

    sim.close()

    t_mean = np.mean(timings)
    print(f"Average render time {t_mean * 1000:.2f}ms, eqivalent to {1 / t_mean:.2f}fps")
    print("For more optimized depth rendering, check out the raycasting.py example.")


if __name__ == "__main__":
    main(show_plot=True, save_plot=False)

```

--------------------------------

### Run Sampling-based MPC Example

Source: https://learnsyslab.github.io/crazyflow/examples

Executes a sampling-based model predictive controller to track a Lissajous curve while avoiding obstacles. The controller utilizes identified dynamics and GPU acceleration if available.

```bash
python examples/control/sampling.py

```

--------------------------------

### Create Default SimState

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/sim/data

Initializes a default SimState with zeroed positions, velocities, and orientations, and sets the identity quaternion. Useful for starting a new simulation.

```python
import jax.numpy as jnp
from typing import Literal
from flax.struct import dataclass

Device = Literal["cpu", "gpu", "tpu"]

@dataclass
class SimState:
    pos: jnp.ndarray
    quat: jnp.ndarray
    vel: jnp.ndarray
    ang_vel: jnp.ndarray
    force: jnp.ndarray
    torque: jnp.ndarray
    rotor_vel: jnp.ndarray

    @staticmethod
    def create(n_worlds: int, n_drones: int, device: Device) -> SimState:
        """Create a default set of states for the simulation."""
        zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
        q_identity = jnp.zeros((n_worlds, n_drones, 4), device=device)
        q_identity = q_identity.at[..., -1].set(1.0)
        rotor_vel = jnp.zeros((n_worlds, n_drones, 4), device=device)
        return SimState(
            pos=zeros_3d,
            quat=q_identity,
            vel=zeros_3d,
            ang_vel=zeros_3d,
            force=zeros_3d,
            torque=zeros_3d,
            rotor_vel=rotor_vel,
        )

```

--------------------------------

### Initialize Rotor Velocity Control Simulation

Source: https://learnsyslab.github.io/crazyflow/user-guide/control

Sets up a simulation for direct motor commands using the first principles dynamics model. The simulation is reset after setup.

```python
import numpy as np
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(control=Control.rotor_vel, dynamics=Dynamics.first_principles)
sim.reset()

cmd = np.full((1, 1, 4), 15_000.0, dtype=np.float32)

sim.rotor_vel_control(cmd)
sim.step(1)

```

--------------------------------

### Render Depth with Raycasting

Source: https://learnsyslab.github.io/crazyflow/examples

This example demonstrates rendering depth images using raycasting from a camera. It shows how to cap maximum distance for visualization and how to build a compiled depth renderer function for performance.

```python
import jax.numpy as jnp
import matplotlib.pyplot as plt

from crazyflow.sim import Sim
from crazyflow.sim.sensors import build_render_depth_fn, render_depth


def main(plot: bool = False):
    sim = Sim()
    sim.data = sim.data.replace(
        states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.2))
    )
    # The easiest way to get depth images is to use the render_depth function
    dist = render_depth(sim, camera=0, resolution=(100, 100), include_drone=False)
    dist = dist.at[dist > 1.5].set(jnp.nan)  # Cap max distance for better visualization
    if plot:
        plt.imshow(dist[0], cmap="viridis")
        plt.colorbar(label="Distance (m)")
        plt.title("Raycast Distance from Camera")
        plt.show()
    # We can also build a depth renderer function for better performance if we need maximum speed or
    # more fine-grained control. Here we only render the drone collision geometry to avoid expensive
    # raycasting against the high-poly visual mesh of the drone.
    render_depth_fn = build_render_depth_fn(
        sim.mjx_model, camera=0, resolution=(200, 200), geomgroup=(1, 1, 0, 1, 1, 1, 1, 1)
    )
    dist_fn = render_depth_fn(sim)
    dist_fn = dist_fn.at[dist_fn > 1.5].set(jnp.nan)  # Cap max distance for better visualization
    if plot:
        plt.imshow(dist_fn[0], cmap="viridis")
        plt.colorbar(label="Distance (m)")
        plt.title("Raycast Distance from Camera (Compiled)")
        plt.show()


if __name__ == "__main__":
    main(plot=True)

```

--------------------------------

### State Control Example

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Demonstrates state control by setting a desired hover position and then stepping the simulation. This method uses an internal Mellinger controller to convert state commands to attitude commands.

```python
import numpy as np
from crazyflow.sim import Sim
from crazyflow.control import Control

sim = Sim(n_worlds=1, n_drones=1, control=Control.state)
sim.reset()

# [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]
cmd = np.zeros((1, 1, 13), dtype=np.float32)
cmd[0, 0, 2] = 0.5  # hover at 0.5 m

sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)

```

--------------------------------

### Using First Principles Dynamics

Source: https://learnsyslab.github.io/crazyflow/user-guide/dynamics/dynamics-functions

This snippet demonstrates how to initialize and use the first principles dynamics model. It requires specifying the drone model and provides an example of calling the dynamics function with state variables and motor RPM commands.

```python
import numpy as np
from crazyflow.dynamics import parametrize
from crazyflow.dynamics.first_principles import dynamics

dynamics = parametrize(dynamics, drone="cf2x_L250")

pos, vel, ang_vel = np.zeros((3,)), np.zeros((3,)), np.zeros((3,))
quat = np.array([0.0, 0.0, 0.0, 1.0])
cmd = np.full((4,), 15_000.0)
rotor_vel = np.full((4,), 12_000.0)

pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = dynamics(
    pos, quat, vel, ang_vel,
    cmd,  # shape (4,) — motor RPMs
    rotor_vel,  # shape (4,) — current motor RPMs; pass None to skip rotor dynamics
)

```

--------------------------------

### Drone Hovering with State Control

Source: https://learnsyslab.github.io/crazyflow/examples

This example demonstrates a minimal end-to-end loop for commanding a single drone to hold a fixed height using state control. It covers creating a simulation, resetting it, applying a state command, and stepping forward.

```python
import numpy as np

from crazyflow.control import Control
from crazyflow.sim import Dynamics, Sim


def main():
    sim = Sim(
        n_worlds=1,
        n_drones=1,
        dynamics=Dynamics.first_principles,
        control=Control.state,
        freq=500,
        attitude_freq=500,
        state_freq=100,
        device="cpu",
    )

    sim.reset()
    duration = 5.0
    fps = 60

    # State cmd is [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]
    cmd = np.zeros((sim.n_worlds, sim.n_drones, 13))
    cmd[..., :3] = 0.1

    for i in range(int(duration * sim.control_freq)):
        sim.state_control(cmd)
        sim.step(sim.freq // sim.control_freq)
        if ((i * fps) % sim.control_freq) < fps:
            sim.render()
    sim.close()


if __name__ == "__main__":
    main()

```

```bash
python examples/control/hover.py

```

--------------------------------

### Create a Simulator Instance

Source: https://learnsyslab.github.io/crazyflow/get-started/quick-start

Initializes a simulator with specified parameters and resets it to a default state. Use this to begin a new simulation.

```python
from crazyflow.sim import Sim

sim = Sim(n_worlds=1, n_drones=1, freq=500)
sim.reset()

```

--------------------------------

### SimState.create

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/sim/data

Initializes a default SimState object with zeroed values for all state attributes, suitable for starting a new simulation.

```APIDOC
## SimState.create

### Description
Creates a default set of states for the simulation, initializing all state variables to their zero or identity values.

### Method
`create(n_worlds: int, n_drones: int, device: Device) -> SimState`

### Parameters
- **n_worlds** (int) - The number of simulation worlds.
- **n_drones** (int) - The number of drones per world.
- **device** (Device) - The computation device (e.g., CPU, GPU).

### Returns
- **SimState** - An instance of SimState with initial zeroed states.
```

--------------------------------

### contacts

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow

Get contact information from the simulation. This method can filter contacts for a specific body or return flags for all bodies.

```APIDOC
## contacts(body=None)

### Description
Get contact information from the simulation. This method can filter contacts for a specific body or return flags for all bodies.

### Parameters
#### Query Parameters
- **body** (str | None) - Optional - The name of the body to filter contacts for. If None, returns flags for all bodies.

### Response
#### Success Response (200)
- **Array** - An boolean array of shape (n_worlds,) that is True if any contact is present.

### Response Example
{
  "example": "[true, false, true]"
}
```

--------------------------------

### build_step_fn

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow

Sets up the chain of functions that are called in `Sim.step()`. It constructs a step function at initialization that selects the correct functions based on the settings, and can be used for pure functional style programming.

```APIDOC
## build_step_fn()

### Description
Sets up the chain of functions that are called in `Sim.step()`. This function constructs a step function at initialization that selects the correct functions based on the settings. It modifies `Sim.step()` in-place and also returns the function for pure functional style programming.

### Method
This documentation describes a Python method, not an HTTP endpoint.

### Parameters
This method does not explicitly define parameters in the provided documentation.

### Returns
- **Callable[[SimData, int], SimData]**: The pure JAX function that steps through the simulation. It takes the current `SimData` and the number of steps to simulate, and returns the updated `SimData`.
```

--------------------------------

### Execute Raycasting Script

Source: https://learnsyslab.github.io/crazyflow/examples

This command executes a Python script located at examples/rendering/raycasting.py. It is likely related to the depth sensing examples.

```bash
python examples/rendering/raycasting.py

```

--------------------------------

### Rotor Velocity Control Example

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Shows the lowest-level control by directly commanding each motor's RPM. This requires the `Dynamics.first_principles` model.

```python
import numpy as np
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(n_worlds=1, n_drones=1, control=Control.rotor_vel, dynamics=Dynamics.first_principles)
sim.reset()

# [rpm_motor_0, rpm_motor_1, rpm_motor_2, rpm_motor_3]
cmd = np.full((1, 1, 4), 15_000.0, dtype=np.float32)

sim.rotor_vel_control(cmd)
sim.step(1)

```

--------------------------------

### Inspect Default Step Pipeline

Source: https://learnsyslab.github.io/crazyflow/user-guide/pipelines

This snippet shows how to initialize a Sim object and print the keys of its default step pipeline to see the available stages.

```python
from crazyflow.sim import Sim

sim = Sim()
print(tuple(sim.step_pipeline.keys()))
# ('step_attitude_controller', 'step_force_torque_controller', 'integration', 'increment_steps', 'clip_floor_pos')
```

--------------------------------

### Get Contact Information

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow

Retrieves contact information from the simulation. Use this to check if any contact is present across all worlds or for a specific body.

```python
def contacts(self, body: str | None = None) -> Array:
    """Get contact information from the simulation.

    Args:
        body: Optional body name to filter contacts for. If None, returns flags for all bodies.

    Returns:
        An boolean array of shape (n_worlds,) that is True if any contact is present.
    """
    if body is None:
        return self.mjx_data._impl.contact.dist < 0
    body_id = self.mj_model.body(body).id
    geom_start = self.mj_model.body_geomadr[body_id]
    geom_count = self.mj_model.body_geomnum[body_id]
    return contacts(geom_start, geom_count, self.mjx_data)
```

--------------------------------

### Initialize Sim, Reset All Worlds, Stage Command, and Step

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Initializes a simulation with 4 worlds and 1 drone, resets all worlds, stages a state control command, and advances the simulation by 50 dynamics steps. Controllers fire at their configured rates.

```python
import numpy as np
from crazyflow.sim import Sim
from crazyflow.control import Control

sim = Sim(n_worlds=4, n_drones=1, control=Control.state)
sim.reset()  # reset all worlds

# Stage a command and advance 50 dynamics steps (controllers fire at their rate)
cmd = np.zeros((4, 1, 13), dtype=np.float32)
cmd[..., 2] = 0.5
sim.state_control(cmd)
sim.step(50)
```

--------------------------------

### Use Default Parameters for Drone

Source: https://learnsyslab.github.io/crazyflow/user-guide/control/parametrize

Binds default parameters for a specified drone configuration to a controller function. This example uses the 'cf2x_L250' configuration.

```python
import numpy as np
from crazyflow.control import parametrize
from crazyflow.control.mellinger import state2attitude

ctrl = parametrize(state2attitude, "cf2x_L250")
pos = np.zeros(3)
quat = np.array([0.0, 0.0, 0.0, 1.0])
vel = np.zeros(3)
cmd = np.zeros(13)
rpyt, _ = ctrl(pos, quat, vel, cmd)
```

--------------------------------

### Select and Parametrize Dynamics by Name

Source: https://learnsyslab.github.io/crazyflow/user-guide/dynamics/parametrize

Demonstrates how to retrieve a dynamics function by its name from `available_dynamics` and then parametrize it for a specific drone model.

```python
from crazyflow.dynamics import available_dynamics, parametrize

list(available_dynamics)  # ['first_principles', 'so_rpy', 'so_rpy_rotor', 'so_rpy_rotor_drag']

dynamics = available_dynamics["so_rpy_rotor_drag"]
parametrized_dynamics = parametrize(dynamics, drone="cf2x_T350")
```

--------------------------------

### Get Dynamics Feature Flags

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/dynamics

Retrieve the feature flags declared by a dynamics function. Use this to understand which optional inputs a dynamics function accepts.

```python
from crazyflow.dynamics import dynamics_features
from crazyflow.dynamics.first_principles import dynamics

dynamics_features(dynamics)  # {'rotor_dynamics': True}
```

--------------------------------

### SimCore.create

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/sim/data

Creates a default set of core simulation parameters, including frequency, number of worlds and drones, mocap IDs, RNG key, and device.

```APIDOC
## SimCore.create

### Description
Create a default set of core simulation parameters.

### Method
`staticmethod`

### Signature
`create(freq: int, n_worlds: int, n_drones: int, drone_mocap_ids: Array, rng_key: int | Array, device: Device) -> SimCore`

### Parameters
- **freq** (int) - The simulation frequency.
- **n_worlds** (int) - The number of simulation worlds.
- **n_drones** (int) - The number of drones per world.
- **drone_mocap_ids** (Array) - MuJoCo mocap IDs for the drone bodies.
- **rng_key** (int | Array) - The random number generator key for the simulation.
- **device** (Device) - The device to run the simulation on (e.g., CPU, GPU).

### Returns
- **SimCore** - An instance of SimCore with the initialized simulation parameters.
```

--------------------------------

### ReachPosEnv Initialization

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/envs/reach_pos_env

Initializes the ReachPosEnv with customizable position and velocity bounds, episode time limits, dynamics, frequency, and device. It sets up the observation space to include the difference to the goal.

```python
def __init__(
    self,
    pos_min: Array | None = None,
    pos_max: Array | None = None,
    vel_min: float = -1.0,
    vel_max: float = 1.0,
    num_envs: int = 1,
    max_episode_time: float = 10.0,
    dynamics: Dynamics = Dynamics.so_rpy,
    freq: int = 500,
    device: str = "cpu",
):
    pos_min = jnp.array([-1.0, -1.0, 1.0]) if pos_min is None else pos_min
    pos_max = jnp.array([1.0, 1.0, 2.0]) if pos_max is None else pos_max
    reset_randomization = partial(
        self._reset_randomization, pmin=pos_min, pmax=pos_max, vmin=vel_min, vmax=vel_max
    )
    super().__init__(
        num_envs=num_envs,
        max_episode_time=max_episode_time,
        dynamics=dynamics,
        freq=freq,
        device=device,
        reset_randomization=reset_randomization,
    )
    self.jax_key = jax.device_put(
        jax.random.key(int(self.np_random.random() * 2**32)), self.device
    )
    spec = {k: v for k, v in self.single_observation_space.items()}
    spec["difference_to_goal"] = spaces.Box(-np.inf, np.inf, shape=(3,))
    self.single_observation_space = spaces.Dict(spec)
    self.observation_space = batch_space(self.single_observation_space, self.sim.n_worlds)
    self._goal = jnp.zeros((self.sim.n_worlds, 3), dtype=jnp.float32, device=self.device)
```

--------------------------------

### Initialize Attitude Control Simulation

Source: https://learnsyslab.github.io/crazyflow/user-guide/control

Sets up a simulation environment for attitude control with a specified dynamics model and attitude frequency. The simulation is then reset to its initial state.

```python
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(control=Control.attitude, dynamics=Dynamics.so_rpy, attitude_freq=500)
sim.reset()

```

--------------------------------

### Create Default Simulation Parameters

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/dynamics/first_principles

Creates a default set of simulation parameters for a specified number of worlds and drones. It loads physical constants and initializes them on the specified device.

```python
import jax
import jax.numpy as jnp
from crazyflow.dynamics.first_principles.dynamics import Params, load_params
from crazyflow.types import Device

@staticmethod
def create(n_worlds: int, n_drones: int, drone: str, device: Device) -> Params:
    """Create a default set of parameters for the simulation."""
    p = load_params(dynamics, drone)
    J = jax.device_put(jnp.tile(p["J"][None, None, :, :], (n_worlds, n_drones, 1, 1)), device)
    return Params(
        mass=jnp.full((n_worlds, n_drones, 1), p["mass"], device=device),
        L=jnp.asarray(p["L"], device=device),
        prop_inertia=jnp.asarray(p["prop_inertia"], device=device),
        gravity_vec=jnp.asarray(p["gravity_vec"], device=device),
        J=J,
        J_inv=jnp.linalg.inv(J),
        rpm2thrust=jnp.asarray(p["rpm2thrust"], device=device),
        rpm2torque=jnp.asarray(p["rpm2torque"], device=device),
        mixing_matrix=jnp.asarray(p["mixing_matrix"], device=device),
        drag_matrix=jnp.asarray(p["drag_matrix"], device=device),
        rotor_dyn_coef=jnp.asarray(p["rotor_dyn_coef"], device=device),
    )
```

--------------------------------

### LandingEnv Constructor

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/envs/landing_env

Initializes the LandingEnv with specified parameters for drone landing simulations. It configures the environment for multiple parallel simulations, sets episode time limits, dynamics models, frequency, and device.

```APIDOC
## `crazyflow.envs.landing_env.LandingEnv`

### Description
Drone environment for landing at a target position. This class extends `DroneEnv` and configures the observation space to include the difference to a goal position.

### Parameters
* **num_envs** (int) - Optional - The number of parallel environments to run.
* **max_episode_time** (float) - Optional - The maximum time for each episode in seconds.
* **dynamics** (Dynamics) - Optional - The dynamics model to use for the simulation. Defaults to `Dynamics.so_rpy`.
* **freq** (int) - Optional - The simulation frequency.
* **device** (str) - Optional - The device to run the simulation on (e.g., 'cpu', 'cuda'). Defaults to 'cpu'.

### Source Code
`crazyflow/envs/landing_env.py`
```

--------------------------------

### JIT Compilation with Integral Errors

Source: https://learnsyslab.github.io/crazyflow/user-guide/control/jit

Shows how integral errors, as regular arrays, are handled by JAX's `jit` without special treatment. This example compiles the function once by initializing integral errors to zero.

```python
import jax
import jax.numpy as jnp
from crazyflow.control import parametrize
from crazyflow.control.mellinger import state2attitude

ctrl = parametrize(state2attitude, "cf2x_L250", xp=jnp)
jit_ctrl = jax.jit(ctrl)

pos = jnp.zeros(3)
quat = jnp.array([0.0, 0.0, 0.0, 1.0])
vel = jnp.zeros(3)
cmd = jnp.zeros(13)

pos_err_i = jnp.zeros(3)  # initialise to zero, so the function compiles only once
for _ in range(10):
    rpyt, pos_err_i = jit_ctrl(pos, quat, vel, cmd, pos_err_i=pos_err_i)

```

--------------------------------

### LandingEnv Initialization

Source: https://learnsyslab.github.io/crazyflow/api/crazyflow/envs/landing_env

Initializes the LandingEnv with parameters for simulation, including environment count, episode time, dynamics model, frequency, and device. It customizes the observation space to include the difference to the goal.

```python
def __init__(
    self,
    num_envs: int = 1,
    max_episode_time: float = 10.0,
    dynamics: Dynamics = Dynamics.so_rpy,
    freq: int = 500,
    device: str = "cpu",
):
    super().__init__(
        num_envs=num_envs,
        max_episode_time=max_episode_time,
        dynamics=dynamics,
        freq=freq,
        device=device,
    )
    spec = {k: v for k, v in self.single_observation_space.items()}
    spec["difference_to_goal"] = spaces.Box(-np.inf, np.inf, shape=(3,))
    self.single_observation_space = spaces.Dict(spec)
    self.observation_space = batch_space(self.single_observation_space, self.sim.n_worlds)
    self._goal = jnp.zeros((self.sim.n_worlds, 3), device=self.device)
    self._goal = self._goal.at[..., 2].set(0.1)  # 10cm above ground
```

--------------------------------

### Basic Gymnasium Environment Usage

Source: https://learnsyslab.github.io/crazyflow/user-guide/gymnasium-envs

Demonstrates the fundamental steps of creating, resetting, stepping, and closing a vectorized Gymnasium environment. This is suitable for general RL agent training.

```python
import gymnasium
import crazyflow.envs  # noqa: F401 - registers the environments

env = gymnasium.make_vec("DroneFigureEightTrajectory-v0", num_envs=16)
obs, info = env.reset()

for _ in range(500):
    action = env.action_space.sample()  # random policy for illustration
    obs, reward, terminated, truncated, info = env.step(action)

env.close()

```

--------------------------------

### Attitude Control Example

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Shows how to use attitude control by commanding roll, pitch, yaw setpoints and collective thrust for hovering. This bypasses the position/velocity loop and is suitable for RL agents outputting attitude targets.

```python
import numpy as np
from crazyflow.sim import Sim, Dynamics
from crazyflow.control import Control

sim = Sim(n_worlds=1, n_drones=1, control=Control.attitude, dynamics=Dynamics.so_rpy)
sim.reset()

# [roll, pitch, yaw, collective_thrust_N]
cmd = np.zeros((1, 1, 4), dtype=np.float32)
cmd[0, 0, 3] = float(sim.data.params.mass[0, 0, 0]) * 9.81  # hover thrust

sim.attitude_control(cmd)
sim.step(sim.freq // sim.control_freq)

```

--------------------------------

### Create and Initialize Sim

Source: https://learnsyslab.github.io/crazyflow/user-guide/oo-api

Instantiates the Sim class with various configuration parameters and resets the simulation environment. This is the main entry point for using the OO API.

```python
from crazyflow.sim import Sim, Dynamics
from crazyflow.sim.integration import Integrator
from crazyflow.control import Control

sim = Sim(
    n_worlds=1,
    n_drones=1,
    drone="cf2x_L250",
    dynamics=Dynamics.first_principles,
    control=Control.state,
    integrator=Integrator.rk4,
    freq=500,                       # dynamics update rate, Hz
    state_freq=100,                 # state controller rate, Hz
    attitude_freq=500,              # attitude controller rate, Hz
    device="cpu",
)
sim.reset()

```

--------------------------------

### Drone Attitude Control with RL Agent Output

Source: https://learnsyslab.github.io/crazyflow/examples

This example shows how to command roll, pitch, yaw, and collective thrust directly, bypassing the Mellinger position loop. It's suitable for RL agents outputting attitude targets.

```python
import os

import numpy as np

os.environ["SCIPY_ARRAY_API"] = "1"

from scipy.spatial.transform import Rotation as R

from crazyflow.control import Control
from crazyflow.sim import Sim

kp = np.array([0.4, 0.4, 1.25])
ki = np.array([0.05, 0.05, 0.05])
kd = np.array([0.2, 0.2, 0.4])
g = 9.81


def control(
    t: float, obs: dict[str, np.ndarray], pos_start: np.ndarray, drone_mass: float
) -> np.ndarray:
    des_pos = np.zeros(3)
    des_pos[..., :2] = pos_start[:2] + np.array([np.cos(t) - 1, np.sin(t)])
    des_pos[..., 2] = 0.2 * t
    des_vel = np.zeros_like(des_pos)
    des_yaw = t

    # Calculate the deviations from the desired trajectory
    pos_error = des_pos - np.array(obs["pos"])
    vel_error = des_vel - np.array(obs["vel"])

    # Compute target thrust
    target_thrust = np.zeros(3)
    target_thrust += kp * pos_error
    target_thrust += kd * vel_error
    target_thrust[2] += drone_mass * g

    # Update z_axis to the current orientation of the drone
    z_axis = R.from_quat(obs["quat"]).as_matrix()[:, 2]

    # update current thrust
    thrust_desired = target_thrust.dot(z_axis)

    # update z_axis_desired
    z_axis_desired = target_thrust / np.linalg.norm(target_thrust)
    x_c_des = np.array([np.cos(des_yaw), np.sin(des_yaw), 0.0])
    y_axis_desired = np.cross(z_axis_desired, x_c_des)
    y_axis_desired /= np.linalg.norm(y_axis_desired)
    x_axis_desired = np.cross(y_axis_desired, z_axis_desired)

    R_desired = np.vstack([x_axis_desired, y_axis_desired, z_axis_desired]).T
    euler_desired = R.from_matrix(R_desired).as_euler("xyz", degrees=False)

    action = np.concatenate([euler_desired, [thrust_desired]], dtype=np.float32)

    return action


def main():
    sim = Sim(control=Control.attitude)
    sim.reset()
    duration = 6.5
    fps = 60

    cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))  # [roll, pitch, yaw, thrust]
    pos_start = sim.data.states.pos
    for i in range(int(duration * sim.control_freq)):
        obs = {
            "pos": sim.data.states.pos[0, 0],
            "vel": sim.data.states.vel[0, 0],
            "quat": sim.data.states.quat[0, 0],
        }
        cmd[0, 0, :] = control(
            i / sim.control_freq, obs, pos_start[0, 0], sim.data.params.mass[0, 0, 0]
        )
        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()

```

--------------------------------

### Advancing Simulation Steps with Control

Source: https://learnsyslab.github.io/crazyflow/user-guide/sim-overview

Shows how to set simulation frequency, reset the simulator, apply state control commands, and advance the simulation by multiple dynamics steps while ensuring the controller fires at the correct rate.

```python
import numpy as np
from crazyflow.sim import Sim
from crazyflow.control import Control

sim = Sim(freq=500, control=Control.state)
sim.reset()
cmd = np.zeros((1, 1, 13), dtype=np.float32)
sim.state_control(cmd)
sim.step(sim.freq // sim.control_freq)  # 500 // 100 = 5 dynamics steps, controller fires once

```
Tokens:
Raw