From 9a52470fab7f9e706e4aeb2e5a81e609f22c76d3 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Wed, 19 Nov 2025 19:50:32 +0100 Subject: [PATCH 1/8] Add change_material function in visualize utils. --- crazyflow/sim/visualize.py | 53 ++++++++++++++++++++++++++++++++++++++ submodules/drone-models | 2 +- 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index ef7f5d5..f79807f 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -68,6 +68,59 @@ def draw_points(sim: Sim, points: NDArray, rgba: NDArray | None = None, size: fl rgba=rgba, ) +def change_material(sim: Sim, mat_name: str, drone_ids: NDArray, rgba: NDArray | None = None, emission: NDArray | None = None): + """Change the material of all drones matching the mask. + + Args: + sim: The simulation. + mat_name: The name of the material to change. + drone_ids: Array of drone indices to modify, shape (n,), dtype=int. + rgba: The RGBA color to set, should be of shape (n, 4) or (4,) to be auto-broadcasted. + emission: The emission value of material, should be of shape (n,) or scalar. + """ + if drone_ids.ndim != 1: + raise ValueError(f"drone_ids must be 1D array, got shape {drone_ids.shape}") + if np.any(drone_ids < 0) or np.any(drone_ids >= sim.n_drones): + raise ValueError(f"drone_ids must be in range [0, {sim.n_drones - 1}], got {drone_ids}") + + if rgba is not None: + rgba = np.asarray(rgba, dtype=float) + try: + # this returns itself if rgba is already the right shape + rgba = np.broadcast_to(rgba, (len(drone_ids), 4)).copy() + except Exception: + raise ValueError( + f"rgba must be shape (4,) or ({len(drone_ids)}, 4), got {rgba.shape}" + ) + rgba = np.clip(rgba, 0.0, 1.0) + + if emission is not None: + emission = np.asarray(emission, dtype=float) + try: + emission = np.broadcast_to(emission, (len(drone_ids),)).copy() + except Exception: + raise ValueError( + f"emission must be scalar or shape ({len(drone_ids)},), got {emission.shape}" + ) + emission = np.maximum(emission, 0.0) + + mj_model = sim.mj_model + + for i, id in enumerate(drone_ids): + full_mat_name = f"{mat_name}:{id}" + mat_id = mujoco.mj_name2id( + mj_model, + mujoco.mjtObj.mjOBJ_MATERIAL, + full_mat_name, + ) + if mat_id < 0: + raise ValueError(f"Material '{full_mat_name}' not found in MuJoCo model.") + + if rgba is not None: + mj_model.mat_rgba[mat_id, :] = rgba[i] + + if emission is not None: + mj_model.mat_emission[mat_id] = emission[i] def _rotation_matrix_from_points(p1: NDArray, p2: NDArray) -> R: """Generate rotation matrices that align their z-axis to p2-p1.""" diff --git a/submodules/drone-models b/submodules/drone-models index 73a083a..4eff0cc 160000 --- a/submodules/drone-models +++ b/submodules/drone-models @@ -1 +1 @@ -Subproject commit 73a083a4bb13cd192c1653229e3d984615879a44 +Subproject commit 4eff0cc7b0ed3789228d9c6c55bfb38c19160280 From 29398f4e445befe4fc5edfc7f7e948434e05c989 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Wed, 19 Nov 2025 20:26:26 +0100 Subject: [PATCH 2/8] Fix linting. --- crazyflow/sim/visualize.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index f79807f..a46b5d0 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -68,7 +68,14 @@ def draw_points(sim: Sim, points: NDArray, rgba: NDArray | None = None, size: fl rgba=rgba, ) -def change_material(sim: Sim, mat_name: str, drone_ids: NDArray, rgba: NDArray | None = None, emission: NDArray | None = None): + +def change_material( + sim: Sim, + mat_name: str, + drone_ids: NDArray, + rgba: NDArray | None = None, + emission: NDArray | None = None, +): """Change the material of all drones matching the mask. Args: @@ -82,16 +89,14 @@ def change_material(sim: Sim, mat_name: str, drone_ids: NDArray, rgba: NDArray | raise ValueError(f"drone_ids must be 1D array, got shape {drone_ids.shape}") if np.any(drone_ids < 0) or np.any(drone_ids >= sim.n_drones): raise ValueError(f"drone_ids must be in range [0, {sim.n_drones - 1}], got {drone_ids}") - + if rgba is not None: rgba = np.asarray(rgba, dtype=float) try: # this returns itself if rgba is already the right shape rgba = np.broadcast_to(rgba, (len(drone_ids), 4)).copy() except Exception: - raise ValueError( - f"rgba must be shape (4,) or ({len(drone_ids)}, 4), got {rgba.shape}" - ) + raise ValueError(f"rgba must be shape (4,) or ({len(drone_ids)}, 4), got {rgba.shape}") rgba = np.clip(rgba, 0.0, 1.0) if emission is not None: @@ -108,11 +113,7 @@ def change_material(sim: Sim, mat_name: str, drone_ids: NDArray, rgba: NDArray | for i, id in enumerate(drone_ids): full_mat_name = f"{mat_name}:{id}" - mat_id = mujoco.mj_name2id( - mj_model, - mujoco.mjtObj.mjOBJ_MATERIAL, - full_mat_name, - ) + mat_id = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, full_mat_name) if mat_id < 0: raise ValueError(f"Material '{full_mat_name}' not found in MuJoCo model.") @@ -122,6 +123,7 @@ def change_material(sim: Sim, mat_name: str, drone_ids: NDArray, rgba: NDArray | if emission is not None: mj_model.mat_emission[mat_id] = emission[i] + def _rotation_matrix_from_points(p1: NDArray, p2: NDArray) -> R: """Generate rotation matrices that align their z-axis to p2-p1.""" p1, p2 = p1.copy(), p2.copy() # Make sure we don't modify the original arrays From f1df3469f85f7adbecc36a517890932405ce6878 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Wed, 19 Nov 2025 20:55:22 +0100 Subject: [PATCH 3/8] Add example for rendering led decks. --- examples/led_deck.py | 56 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 examples/led_deck.py diff --git a/examples/led_deck.py b/examples/led_deck.py new file mode 100644 index 0000000..e89ee87 --- /dev/null +++ b/examples/led_deck.py @@ -0,0 +1,56 @@ +import numpy as np + +from crazyflow.control.control import Control +from crazyflow.sim import Physics, Sim +from crazyflow.sim.visualize import change_material + + +def main(): + """Spawn 25 drones in one world and activate led decks.""" + n_worlds, n_drones = 1, 25 + sim = Sim( + n_worlds=n_worlds, + n_drones=n_drones, + drone_model="cf21B_500", + control=Control.state, + physics=Physics.so_rpy, + device="cpu", + ) + fps = 60 + cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) + cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 + rgbas = np.random.default_rng(0).uniform(0, 1, (n_drones, 4)) + rgbas[..., 3] = 1.0 + + init_pos = np.array(sim.data.states.pos[0, :, :]) + cmd = np.zeros((sim.n_worlds, sim.n_drones, 13)) + cmd[:, :, :3] = init_pos + cmd[:, :, 2] += 1.5 + + for i in range(int(10 * sim.control_freq)): + sim.state_control(cmd) + sim.step(sim.freq // sim.control_freq) + if ((i * fps) % sim.control_freq) < fps: + even_ids = np.arange(0, n_drones, 2) + odd_ids = np.arange(1, n_drones, 2) + emission = np.sin(i / sim.control_freq * np.pi) + 1.0 + change_material( + sim, + mat_name="led_top", + drone_ids=even_ids, + rgba=rgbas[even_ids, :], + emission=emission, + ) + change_material( + sim, + mat_name="led_bot", + drone_ids=odd_ids, + rgba=rgbas[odd_ids, :], + emission=emission, + ) + sim.render() + sim.close() + + +if __name__ == "__main__": + main() From d5ca60db973d9a2b139750421eed73c5f14c5eba Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Thu, 20 Nov 2025 09:21:21 +0100 Subject: [PATCH 4/8] Simplify code. --- crazyflow/sim/visualize.py | 35 ++++++++++++----------------------- examples/led_deck.py | 20 ++++++-------------- 2 files changed, 18 insertions(+), 37 deletions(-) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index a46b5d0..48cfe8f 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -76,7 +76,7 @@ def change_material( rgba: NDArray | None = None, emission: NDArray | None = None, ): - """Change the material of all drones matching the mask. + """Change the material of specified drones. Args: sim: The simulation. @@ -91,37 +91,26 @@ def change_material( raise ValueError(f"drone_ids must be in range [0, {sim.n_drones - 1}], got {drone_ids}") if rgba is not None: - rgba = np.asarray(rgba, dtype=float) - try: - # this returns itself if rgba is already the right shape - rgba = np.broadcast_to(rgba, (len(drone_ids), 4)).copy() - except Exception: - raise ValueError(f"rgba must be shape (4,) or ({len(drone_ids)}, 4), got {rgba.shape}") - rgba = np.clip(rgba, 0.0, 1.0) + # this returns itself if rgba is already the right shape + rgba = np.broadcast_to(rgba, (len(drone_ids), 4)) if emission is not None: - emission = np.asarray(emission, dtype=float) - try: - emission = np.broadcast_to(emission, (len(drone_ids),)).copy() - except Exception: - raise ValueError( - f"emission must be scalar or shape ({len(drone_ids)},), got {emission.shape}" - ) - emission = np.maximum(emission, 0.0) + emission = np.broadcast_to(emission, (len(drone_ids),)) mj_model = sim.mj_model - - for i, id in enumerate(drone_ids): - full_mat_name = f"{mat_name}:{id}" + mat_ids = [] + for i, drone_id in enumerate(drone_ids): + full_mat_name = f"{mat_name}:{drone_id}" mat_id = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, full_mat_name) if mat_id < 0: raise ValueError(f"Material '{full_mat_name}' not found in MuJoCo model.") + mat_ids.append(mat_id) - if rgba is not None: - mj_model.mat_rgba[mat_id, :] = rgba[i] + if rgba is not None: + mj_model.mat_rgba[mat_ids, :] = rgba - if emission is not None: - mj_model.mat_emission[mat_id] = emission[i] + if emission is not None: + mj_model.mat_emission[mat_ids] = emission def _rotation_matrix_from_points(p1: NDArray, p2: NDArray) -> R: diff --git a/examples/led_deck.py b/examples/led_deck.py index e89ee87..0ee1fb0 100644 --- a/examples/led_deck.py +++ b/examples/led_deck.py @@ -1,25 +1,17 @@ import numpy as np from crazyflow.control.control import Control -from crazyflow.sim import Physics, Sim +from crazyflow.sim import Sim from crazyflow.sim.visualize import change_material def main(): """Spawn 25 drones in one world and activate led decks.""" - n_worlds, n_drones = 1, 25 - sim = Sim( - n_worlds=n_worlds, - n_drones=n_drones, - drone_model="cf21B_500", - control=Control.state, - physics=Physics.so_rpy, - device="cpu", - ) + sim = Sim(n_drones=25, drone_model="cf21B_500", control=Control.state) fps = 60 cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 - rgbas = np.random.default_rng(0).uniform(0, 1, (n_drones, 4)) + rgbas = np.random.default_rng(0).uniform(0, 1, (sim.n_drones, 4)) rgbas[..., 3] = 1.0 init_pos = np.array(sim.data.states.pos[0, :, :]) @@ -31,9 +23,9 @@ def main(): sim.state_control(cmd) sim.step(sim.freq // sim.control_freq) if ((i * fps) % sim.control_freq) < fps: - even_ids = np.arange(0, n_drones, 2) - odd_ids = np.arange(1, n_drones, 2) - emission = np.sin(i / sim.control_freq * np.pi) + 1.0 + even_ids = np.arange(0, sim.n_drones, 2) + odd_ids = np.arange(1, sim.n_drones, 2) + emission = np.sin(i / sim.control_freq * np.pi) change_material( sim, mat_name="led_top", From 33c844c4a254b0e9c31483dcccf595f1b1e2ab87 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Thu, 20 Nov 2025 10:22:13 +0100 Subject: [PATCH 5/8] Use lagecy drones in example. --- examples/led_deck.py | 2 +- submodules/drone-models | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/led_deck.py b/examples/led_deck.py index 0ee1fb0..3a928b0 100644 --- a/examples/led_deck.py +++ b/examples/led_deck.py @@ -7,7 +7,7 @@ def main(): """Spawn 25 drones in one world and activate led decks.""" - sim = Sim(n_drones=25, drone_model="cf21B_500", control=Control.state) + sim = Sim(n_drones=25, control=Control.state) fps = 60 cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 diff --git a/submodules/drone-models b/submodules/drone-models index 4eff0cc..cce86cf 160000 --- a/submodules/drone-models +++ b/submodules/drone-models @@ -1 +1 @@ -Subproject commit 4eff0cc7b0ed3789228d9c6c55bfb38c19160280 +Subproject commit cce86cf6f1696296bcc71377ac89a149905be524 From b3e3d5e9d145708a289191f2dc5f8acec301ada0 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Thu, 20 Nov 2025 11:23:40 +0100 Subject: [PATCH 6/8] Add test cases for change_material. Use new drones in example. --- examples/led_deck.py | 2 +- tests/unit/test_sim.py | 63 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/examples/led_deck.py b/examples/led_deck.py index 3a928b0..0ee1fb0 100644 --- a/examples/led_deck.py +++ b/examples/led_deck.py @@ -7,7 +7,7 @@ def main(): """Spawn 25 drones in one world and activate led decks.""" - sim = Sim(n_drones=25, control=Control.state) + sim = Sim(n_drones=25, drone_model="cf21B_500", control=Control.state) fps = 60 cmd = np.zeros((sim.n_worlds, sim.n_drones, 4)) cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index dcf0c81..297f1aa 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -7,6 +7,7 @@ import jax import jax.numpy as jnp +import mujoco import numpy as np import pytest from jax import Array @@ -16,6 +17,7 @@ from crazyflow.sim import Physics, Sim from crazyflow.sim.data import ControlData from crazyflow.sim.sim import sync_sim2mjx +from crazyflow.sim.visualize import change_material if TYPE_CHECKING: from typing import Any @@ -493,3 +495,64 @@ def test_scan_results(physics: Physics): assert np.all(pos_loop_steps[..., 2] > 0.1), "Drones should have moved" assert np.allclose(pos_scan_steps, pos_loop_steps), "Scan results should be identical" sim.close() + + +@pytest.mark.unit +@pytest.mark.parametrize("drone_model", ["cf2x_L250", "cf2x_P250", "cf2x_T350", "cf21B_500"]) +@pytest.mark.parametrize("mat_name", ["led_top", "led_bot"]) +def test_change_material(device: str, drone_model: str, mat_name: str): + """change_material should broadcast RGBA/emission and update MuJoCo materials appropriately.""" + n_drones = 2 + + sim = Sim(n_drones=n_drones, drone_model=drone_model, device=device) + + drone_ids = np.array([0, 1], dtype=int) + rgba = 0.42 * np.ones((n_drones, 4), dtype=float) + emission = 0.42 * np.ones((n_drones,), dtype=float) + + change_material(sim, mat_name=mat_name, drone_ids=drone_ids, rgba=rgba, emission=emission) + + mj_model = sim.mj_model + mat0 = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, f"{mat_name}:0") + mat1 = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, f"{mat_name}:1") + expected_rgba = 0.42 * np.ones((4,), dtype=float) + expected_emission = 0.42 + np.testing.assert_allclose(mj_model.mat_rgba[mat0], expected_rgba) + np.testing.assert_allclose(mj_model.mat_rgba[mat1], expected_rgba) + assert mj_model.mat_emission[mat0] == pytest.approx(expected_emission) + assert mj_model.mat_emission[mat1] == pytest.approx(expected_emission) + + +@pytest.mark.unit +def test_change_material_errors(device: str): + """Test that change_material raises the expected errors for bad inputs.""" + n_drones = 2 + sim = Sim(n_drones=n_drones, device=device) + + # ---------- 1) Missing material name: mat_name="bad_mat" ---------- + drone_ids_ok = np.array([0, 1], dtype=int) + rgba_ok = 0.42 * np.ones((n_drones, 4), dtype=float) + emission_ok = 0.42 * np.ones((n_drones,), dtype=float) + + with pytest.raises(ValueError, match=r"Material 'bad_mat:0' not found in MuJoCo model\."): + change_material( + sim, mat_name="bad_mat", drone_ids=drone_ids_ok, rgba=rgba_ok, emission=emission_ok + ) + + with pytest.raises(ValueError, match=r"drone_ids must be 1D array"): + change_material( + sim, + mat_name="led_top", + drone_ids=np.array(2, dtype=int), + rgba=rgba_ok, + emission=emission_ok, + ) + + with pytest.raises(ValueError, match=r"drone_ids must be in range \[0, 1\]"): + change_material( + sim, + mat_name="led_top", + drone_ids=np.arange(3, dtype=int), + rgba=rgba_ok, + emission=emission_ok, + ) From bb59cff8daee357e50ef755a9d502dab2a7b8f38 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Thu, 20 Nov 2025 11:41:36 +0100 Subject: [PATCH 7/8] Clean up test error function. --- tests/unit/test_sim.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index 297f1aa..af51f66 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -529,14 +529,13 @@ def test_change_material_errors(device: str): n_drones = 2 sim = Sim(n_drones=n_drones, device=device) - # ---------- 1) Missing material name: mat_name="bad_mat" ---------- - drone_ids_ok = np.array([0, 1], dtype=int) - rgba_ok = 0.42 * np.ones((n_drones, 4), dtype=float) - emission_ok = 0.42 * np.ones((n_drones,), dtype=float) + drone_ids = np.array([0, 1], dtype=int) + rgba = np.ones((n_drones, 4), dtype=float) + emission = np.ones((n_drones,), dtype=float) - with pytest.raises(ValueError, match=r"Material 'bad_mat:0' not found in MuJoCo model\."): + with pytest.raises(ValueError): change_material( - sim, mat_name="bad_mat", drone_ids=drone_ids_ok, rgba=rgba_ok, emission=emission_ok + sim, mat_name="bad_mat", drone_ids=drone_ids, rgba=rgba, emission=emission ) with pytest.raises(ValueError, match=r"drone_ids must be 1D array"): @@ -544,8 +543,8 @@ def test_change_material_errors(device: str): sim, mat_name="led_top", drone_ids=np.array(2, dtype=int), - rgba=rgba_ok, - emission=emission_ok, + rgba=rgba, + emission=emission, ) with pytest.raises(ValueError, match=r"drone_ids must be in range \[0, 1\]"): @@ -553,6 +552,6 @@ def test_change_material_errors(device: str): sim, mat_name="led_top", drone_ids=np.arange(3, dtype=int), - rgba=rgba_ok, - emission=emission_ok, + rgba=rgba, + emission=emission, ) From c0c4c969b9337cb757de04e626776cfdd0f4fe46 Mon Sep 17 00:00:00 2001 From: yufei4hua Date: Fri, 21 Nov 2025 09:55:05 +0100 Subject: [PATCH 8/8] Final cleaning up. --- crazyflow/sim/visualize.py | 3 +-- tests/unit/test_sim.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index 48cfe8f..2c1c4de 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -91,7 +91,6 @@ def change_material( raise ValueError(f"drone_ids must be in range [0, {sim.n_drones - 1}], got {drone_ids}") if rgba is not None: - # this returns itself if rgba is already the right shape rgba = np.broadcast_to(rgba, (len(drone_ids), 4)) if emission is not None: @@ -107,7 +106,7 @@ def change_material( mat_ids.append(mat_id) if rgba is not None: - mj_model.mat_rgba[mat_ids, :] = rgba + mj_model.mat_rgba[mat_ids] = rgba if emission is not None: mj_model.mat_emission[mat_ids] = emission diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index af51f66..c6f25ed 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -538,7 +538,7 @@ def test_change_material_errors(device: str): sim, mat_name="bad_mat", drone_ids=drone_ids, rgba=rgba, emission=emission ) - with pytest.raises(ValueError, match=r"drone_ids must be 1D array"): + with pytest.raises(ValueError, match="drone_ids must be 1D array"): change_material( sim, mat_name="led_top",