diff --git a/crazyflow/drones/cf21B_500.xml b/crazyflow/drones/cf21B_500.xml index 455da94b3..b49d7cb75 100644 --- a/crazyflow/drones/cf21B_500.xml +++ b/crazyflow/drones/cf21B_500.xml @@ -31,7 +31,7 @@ - + @@ -65,8 +65,10 @@ - + + @@ -96,6 +98,7 @@ + diff --git a/crazyflow/drones/cf2x_L250.xml b/crazyflow/drones/cf2x_L250.xml index 9082d9f85..1b0649758 100644 --- a/crazyflow/drones/cf2x_L250.xml +++ b/crazyflow/drones/cf2x_L250.xml @@ -31,7 +31,7 @@ - + @@ -65,8 +65,10 @@ - + + @@ -97,6 +99,7 @@ + diff --git a/crazyflow/drones/cf2x_P250.xml b/crazyflow/drones/cf2x_P250.xml index 40b14bda8..e9370d2ab 100644 --- a/crazyflow/drones/cf2x_P250.xml +++ b/crazyflow/drones/cf2x_P250.xml @@ -31,7 +31,7 @@ - + @@ -65,8 +65,10 @@ - + + @@ -96,6 +98,7 @@ + diff --git a/crazyflow/drones/cf2x_T350.xml b/crazyflow/drones/cf2x_T350.xml index 2b24b91e0..5e3916070 100644 --- a/crazyflow/drones/cf2x_T350.xml +++ b/crazyflow/drones/cf2x_T350.xml @@ -32,7 +32,7 @@ - + @@ -66,8 +66,10 @@ - + + @@ -97,6 +99,7 @@ + diff --git a/crazyflow/drones/hb_x500.xml b/crazyflow/drones/hb_x500.xml index c97e670d2..8043ac4c5 100644 --- a/crazyflow/drones/hb_x500.xml +++ b/crazyflow/drones/hb_x500.xml @@ -49,7 +49,7 @@ https://www.ardusimple.com/wp-content/uploads/2021/02/helical3b_dimensions.jpg --> - + @@ -84,8 +84,10 @@ - + + @@ -180,6 +182,7 @@ + diff --git a/crazyflow/dynamics/utils/data_utils.py b/crazyflow/dynamics/utils/data_utils.py index 7cc86cc16..0fe5048ba 100644 --- a/crazyflow/dynamics/utils/data_utils.py +++ b/crazyflow/dynamics/utils/data_utils.py @@ -5,7 +5,7 @@ from __future__ import annotations -import logging +import warnings from typing import TYPE_CHECKING import numpy as np @@ -19,8 +19,6 @@ if TYPE_CHECKING: from crazyflow._typing import Array # To be changed to array_api_typing later -logger = logging.getLogger(__name__) - def preprocessing(data: dict[str, Array]) -> dict[str, Array]: """Applies preprocessing to collected data. @@ -54,7 +52,7 @@ def preprocessing(data: dict[str, Array]) -> dict[str, Array]: # TODO also check quat for outliers! if n_outliers > 0: - logger.warning(f"{n_outliers} outliers detected. Interpolating") + warnings.warn(f"{n_outliers} outliers detected. Interpolating") time_good = data["time"][~is_outlier] pos_good = data["pos"][~is_outlier] quat_good = data["quat"][~is_outlier] diff --git a/crazyflow/sim/__init__.py b/crazyflow/sim/__init__.py index 6a15efa1a..f143226e1 100644 --- a/crazyflow/sim/__init__.py +++ b/crazyflow/sim/__init__.py @@ -1,4 +1,17 @@ +from enum import StrEnum + from crazyflow.dynamics import Dynamics -from crazyflow.sim.sim import Sim -__all__ = ["Sim", "Dynamics"] + +class CollisionType(StrEnum): + """Collision geometry of the drones, selected with ``select_collision_type``.""" + + sphere = "sphere" + box = "box" + hull = "hull" + default = sphere + + +from crazyflow.sim.sim import Sim # noqa: E402 (sim.py imports CollisionType from this package) + +__all__ = ["Sim", "Dynamics", "CollisionType"] diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index ecd8c1c95..8b3759bf4 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +import warnings import xml.etree.ElementTree as ET from collections import OrderedDict from functools import partial, wraps @@ -30,6 +31,7 @@ from crazyflow.dynamics import Dynamics from crazyflow.dynamics import load_params as load_dynamics_params from crazyflow.exception import ConfigError, NotInitializedError +from crazyflow.sim import CollisionType from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv from crazyflow.sim.dynamics import ( first_principles_dynamics, @@ -345,6 +347,8 @@ def _replicate_drone( asset.append(clone) new_spec = mujoco.MjSpec.from_string(ET.tostring(root, encoding="unicode")) new_spec.copy_during_attach = True # write-only, re-apply after from_string() drops it + for mesh in spec.meshes: # to_xml() drops maxhullvert before MuJoCo 3.16, re-apply it + new_spec.mesh(mesh.name).maxhullvert = mesh.maxhullvert return new_spec def build_mjx_model(self, spec: mujoco.MjSpec) -> tuple[Any, Any, Model, Data]: @@ -509,6 +513,8 @@ def build_mjx(self): self.viewer = None self.spec.numeric("max_geom_pairs").data = [self.max_geom_pairs] self.mj_model, self.mj_data, self.mjx_model, self.mjx_data = self.build_mjx_model(self.spec) + core = self.data.core.replace(mjx_synced=jnp.zeros_like(self.data.core.mjx_synced)) + self.data = self.data.replace(core=core) def init_data( self, @@ -764,26 +770,30 @@ def seed_sim(data: SimData, seed: int, device: Device) -> SimData: return data.replace(core=data.core.replace(rng_key=rng_key)) -def use_box_collision(sim: Sim, enable: bool = True): - """Changes the collision geometry to use boxes or spheres (default). +def select_collision_type(sim: Sim, collision_type: CollisionType = CollisionType.default): + """Select the collision geometry of all drones and rebuild the MJX model. Args: sim: The simulation instance. - enable: If True, use box collision geometry. If False, use sphere collision geometry. + collision_type: ``sphere`` (default) is a bounding sphere, ``box`` an oriented bounding box + and ``hull`` the convex hull of the drone mesh, limited to 32 vertices. The selected + geometry is in geom group 3, which the viewer hides by default (toggle with key 3). Warning: - Using box collision geometry is more computationally expensive than sphere collision - geometry, especially for larger swarms. It is recommended to only enable box collision - geometry for small swarms or when high accuracy is required. + Box and hull collisions are more expensive than spheres. MJX tests all edge pairs of two + hulls, which is cheap on GPU but slows down ``Sim.contacts()`` significantly on CPU for + swarms with many drone pairs. """ + if collision_type == CollisionType.hull and sim.device.platform == "cpu" and sim.n_drones > 16: + warnings.warn( + f"Hull collision with {sim.n_drones} drones slows down `Sim.contacts()` significantly " + "on CPU. Use a GPU or the sphere/box collision type for large swarms." + ) for geom in sim.spec.geoms: - if geom.name.startswith("col_sphere"): - geom.contype = 1 * (not enable) - geom.conaffinity = 1 * (not enable) - geom.rgba[3] = 1 * (not enable) - if geom.name.startswith("col_box"): - geom.contype = 1 * enable - geom.conaffinity = 1 * enable - geom.rgba[3] = 1 * enable - + if not geom.name.startswith(("col_sphere", "col_box", "col_hull")): + continue + enable = geom.name.startswith(f"col_{collision_type}") + geom.contype = 1 * enable + geom.conaffinity = 1 * enable + geom.rgba[3] = 1 * enable sim.build_mjx() diff --git a/docs/examples/index.md b/docs/examples/index.md index c986e5d0c..f58a63dda 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -216,7 +216,7 @@ python examples/rendering/led_deck.py ## Contact queries -The default collision geometry is a sphere around the drone frame. `use_box_collision` replaces it with a tighter oriented box, useful for narrow-gap flight and accurate contact debugging. +The default collision geometry is a sphere around the drone frame. `select_collision_type` switches to a tighter oriented box or to the convex hull of the drone mesh, useful for narrow-gap flight and accurate contact debugging. Hull contacts are cheap on GPU but slow down significantly on CPU for swarms with many drones.
diff --git a/docs/user-guide/mujoco.md b/docs/user-guide/mujoco.md index 5a5c1a279..c3ef56e07 100644 --- a/docs/user-guide/mujoco.md +++ b/docs/user-guide/mujoco.md @@ -168,4 +168,4 @@ in_contact = check_contacts(sim.data, sim.mjx_data.mocap_pos) - [Pipelines](pipelines.md) — inserting custom stages into the step and reset pipelines - [Examples](../examples/index.md#cameras-and-rgbd) — FPV camera and RGBD rendering -- [Examples](../examples/index.md#contact-queries) — contact detection with box collision geometry +- [Examples](../examples/index.md#contact-queries) — contact detection with sphere, box or hull collision geometry diff --git a/examples/contacts/contacts.py b/examples/contacts/contacts.py index 02acd9e2d..e3509d2df 100644 --- a/examples/contacts/contacts.py +++ b/examples/contacts/contacts.py @@ -1,24 +1,82 @@ +import mujoco import numpy as np -from crazyflow.sim import Dynamics, Sim -from crazyflow.sim.sim import use_box_collision +from crazyflow.control import Control +from crazyflow.sim import CollisionType, Dynamics, Sim +from crazyflow.sim.sim import select_collision_type + +RED, GREEN = (1.0, 0.0, 0.0), (0.0, 1.0, 0.0) + + +def add_contact_indicators(sim: Sim) -> list[int]: + """Add a disc above each drone that turns red on contact and return the geom ids.""" + for i in range(sim.n_drones): + geom = sim.spec.body(f"drone:{i}").add_geom() + geom.name = f"indicator:{i}" + geom.type = mujoco.mjtGeom.mjGEOM_CYLINDER + geom.size[:] = [0.05, 0.003, 0.0] + geom.pos[:] = [0.0, 0.0, 0.1] + geom.contype = 0 + geom.conaffinity = 0 + geom.rgba[:] = [*GREEN, 1.0] + sim.build_mjx() + return [sim.mj_model.geom(f"indicator:{i}").id for i in range(sim.n_drones)] + + +def crossing_trajectory(t: float, amplitude: float, period: float) -> tuple[np.ndarray, np.ndarray]: + """Drones 0 and 2 oscillate along x in opposite directions through the hovering drone 1. + + Returns: + Position and velocity setpoints at time ``t`` with shape (3, 3) each. + """ + direction = np.array([1.0, 0.0, -1.0]) + phase = 2 * np.pi * t / period + pos, vel = np.zeros((3, 3)), np.zeros((3, 3)) + pos[:, 0] = amplitude * np.sin(phase) * direction + pos[:, 2] = 1.0 + vel[:, 0] = amplitude * 2 * np.pi / period * np.cos(phase) * direction + return pos, vel def main(): - """Spawn multiple drones in multiple worlds and check for contacts.""" + """Fly the outer drones back and forth through a hovering drone and check for contacts. + + The crossing is repeated with every collision type. Each type reports contacts for a different + fraction of the crossing, from the loose sphere to the tight convex hull. A disc above each + drone turns red while the drone is in contact. Press 3 in the viewer to show the collision + geometry. + """ n_worlds, n_drones = 2, 3 - sim = Sim(n_worlds=n_worlds, n_drones=n_drones, dynamics=Dynamics.so_rpy, device="cpu") - use_box_collision(sim, enable=True) # Enable box collision for all drones + sim = Sim( + n_worlds=n_worlds, + n_drones=n_drones, + control=Control.state, + dynamics=Dynamics.so_rpy, + device="cpu", + ) + indicators = add_contact_indicators(sim) fps = 60 + amplitude, period = 0.3, 2.0 - cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) - cmd[..., 3] = sim.data.params.mass[0] * 9.81 * 1.04 - for i in range(int(2 * sim.control_freq)): - sim.attitude_control(cmd) - sim.step(sim.freq // sim.control_freq) - if ((i * fps) % sim.control_freq) < fps: - sim.render() - print(f"Contacts: {sim.contacts().any()}") + pos, _ = crossing_trajectory(0.0, amplitude, period) # Start on the trajectory + states = sim.data.states + sim.data = sim.data.replace(states=states.replace(pos=states.pos.at[:].set(pos))) + cmd = np.zeros((sim.n_worlds, sim.n_drones, 16)) + cmd[..., 12] = 1.0 # Identity quaternion + for collision_type in CollisionType: + print(f"Collision type: {collision_type}") + select_collision_type(sim, collision_type) # Switches mid-flight, the sim keeps its state + for i in range(int(2 * period * sim.control_freq)): + t = i / sim.control_freq + cmd[..., :3], cmd[..., 3:6] = crossing_trajectory(t, amplitude, period) + sim.state_control(cmd) + sim.step(sim.freq // sim.control_freq) + if ((i * fps) % sim.control_freq) < fps: + for drone, gid in enumerate(indicators): + hit = sim.contacts(f"drone:{drone}")[0].any() + sim.mj_model.geom_rgba[gid, :3] = RED if hit else GREEN + sim.render() + print(f"Contacts: {sim.contacts().any()}") sim.close() diff --git a/examples/contacts/crash.py b/examples/contacts/crash.py index ffa1183ae..e15d03872 100644 --- a/examples/contacts/crash.py +++ b/examples/contacts/crash.py @@ -7,7 +7,6 @@ from crazyflow.control import Control from crazyflow.sim import Sim -from crazyflow.sim.sim import use_box_collision def main(): @@ -15,7 +14,6 @@ def main(): sim = Sim(control=Control.state, freq=500, attitude_freq=500, state_freq=100) sim.reset() fps = 60 - use_box_collision(sim, enable=False) print("Phase 1: Hovering at [0, 0, 0.5] for 3 seconds") hover_duration = 3.0 diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index 5fe1dc9ed..0b09d6f1b 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import jax @@ -16,10 +17,10 @@ from crazyflow.control import Control from crazyflow.dynamics import supported_dynamics from crazyflow.exception import ConfigError -from crazyflow.sim import Dynamics, Sim +from crazyflow.sim import CollisionType, Dynamics, Sim from crazyflow.sim.data import ControlData, SimData from crazyflow.sim.integration import Integrator -from crazyflow.sim.sim import rotor_vel_limits, sync_sim2mjx, use_box_collision +from crazyflow.sim.sim import rotor_vel_limits, select_collision_type, sync_sim2mjx from crazyflow.sim.visualize import change_material if TYPE_CHECKING: @@ -514,17 +515,16 @@ def test_contacts(dynamics: Dynamics): @pytest.mark.unit -@pytest.mark.parametrize("box_collision", [False, True]) -def test_contacts_between_drones(box_collision: bool): - """Two overlapping drones must register a contact. +@pytest.mark.parametrize("collision_type", CollisionType) +def test_contacts_between_drones(collision_type: CollisionType): + """Two overlapping drones must register a contact with every collision geometry. Drones are mocap bodies that are all welded to the world. MuJoCo/MJX filter out collisions between geoms welded into the same kinematic tree, so without explicit contact pairs no drone-drone contact is ever generated. """ sim = Sim(n_drones=2, control=Control.attitude, freq=500, device="cpu") - if box_collision: - use_box_collision(sim, True) + select_collision_type(sim, collision_type) sim.reset() states = sim.data.states # Place both drones at the same position above the floor @@ -536,6 +536,23 @@ def test_contacts_between_drones(box_collision: bool): sim.close() +@pytest.mark.unit +def test_hull_collision_warns_on_cpu(): + """Hull collision for large swarms on CPU must warn, smaller swarms and other types must not.""" + sim = Sim(n_drones=17, device="cpu") + with pytest.warns(UserWarning, match="significantly"): + select_collision_type(sim, CollisionType.hull) + with warnings.catch_warnings(): + warnings.simplefilter("error") + select_collision_type(sim, CollisionType.box) + sim.close() + sim = Sim(n_drones=16, device="cpu") + with warnings.catch_warnings(): + warnings.simplefilter("error") + select_collision_type(sim, CollisionType.hull) + sim.close() + + def _attach_obstacle(sim: Sim, pos: list[float], mocap: bool): """Attach the minimal obstacle to the sim at ``pos``, mirroring downstream track loading.""" OBSTACLE_XML = """