From b86c77c68d507e7ff66b7c8ea0e90955d48ae813 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Fri, 27 Feb 2026 11:52:21 +0100 Subject: [PATCH 1/5] Add draw capsule --- crazyflow/sim/visualize.py | 44 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index 2c1c4de..ae86e51 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -69,6 +69,50 @@ def draw_points(sim: Sim, points: NDArray, rgba: NDArray | None = None, size: fl ) +def draw_capsule( + sim: Sim, + p1: NDArray, + p2: NDArray, + radius: float = 0.05, + rgba: NDArray | None = None, + is_cylinder: bool = False, +): + """Draw a capsule (pill) or cylinder between two points. + + Args: + sim: The simulation. + p1: Start point [3,] + p2: End point [3,] + radius: The thickness of the geom in [m]. + rgba: The color of the object. + is_cylinder: If True, draws a flat-ended cylinder. + If False, draws a pill-shaped capsule. + """ + if sim.viewer is None: + return + + # 1. Calculate Midpoint (Center of the geom) + pos = (p1 + p2) / 2.0 + + # 2. Calculate Half-length (MuJoCo uses half-extents) + dist = np.linalg.norm(p2 - p1) + half_length = dist / 2.0 + + # 3. Define Size: [radius, radius, half_length] + # Note: For capsules, size[2] is the half-length of the *cylindrical* part. + # MuJoCo adds the hemispherical caps on top of this. + size = np.array([radius, half_length, 0]) + + # 4. Get Rotation (Align Z-axis to the vector p2-p1) + # Using your existing helper (wrapped in a list for the reshape) + mat = _rotation_matrix_from_points(p1[None, :], p2[None, :]).as_matrix().flatten() + + geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER if is_cylinder else mujoco.mjtGeom.mjGEOM_CAPSULE + rgba = rgba if rgba is not None else np.array([0, 1.0, 0, 1]) + + sim.viewer.viewer.add_marker(type=geom_type, pos=pos, size=size, mat=mat, rgba=rgba) + + def change_material( sim: Sim, mat_name: str, From 68d37be7ceccb5a475beb220e27006ed4c97d546 Mon Sep 17 00:00:00 2001 From: Rather1337 Date: Wed, 4 Mar 2026 23:27:14 +0100 Subject: [PATCH 2/5] Add draw capsule function --- crazyflow/sim/visualize.py | 44 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index ae86e51..e02325c 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -6,6 +6,50 @@ from crazyflow.sim import Sim +def draw_capsule( + sim: Sim, + p1: NDArray, + p2: NDArray, + radius: float = 0.05, + rgba: NDArray | None = None, + is_cylinder: bool = False, +): + """Draw a capsule (pill) or cylinder between two points. + + Args: + sim: The simulation. + p1: Start point [3,] + p2: End point [3,] + radius: The thickness of the geom. + rgba: The color of the object. + is_cylinder: If True, draws a flat-ended cylinder. + If False, draws a pill-shaped capsule. + """ + if sim.viewer is None: + return + + # 1. Calculate Midpoint (Center of the geom) + pos = (p1 + p2) / 2.0 + + # 2. Calculate Half-length (MuJoCo uses half-extents) + dist = np.linalg.norm(p2 - p1) + half_length = dist / 2.0 + + # 3. Define Size: [radius, radius, half_length] + # Note: For capsules, size[2] is the half-length of the *cylindrical* part. + # MuJoCo adds the hemispherical caps on top of this. + size = np.array([radius, half_length, 0]) + + # 4. Get Rotation (Align Z-axis to the vector p2-p1) + # Using your existing helper (wrapped in a list for the reshape) + mat = _rotation_matrix_from_points(p1[None, :], p2[None, :]).as_matrix().flatten() + + geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER if is_cylinder else mujoco.mjtGeom.mjGEOM_CAPSULE + rgba = rgba if rgba is not None else np.array([0, 1.0, 0, 1]) + + sim.viewer.viewer.add_marker(type=geom_type, pos=pos, size=size, mat=mat, rgba=rgba) + + def draw_line( sim: Sim, points: NDArray, From 23832739269978ad1b088d894f40c3c005765088 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Tue, 17 Mar 2026 23:14:10 +0100 Subject: [PATCH 3/5] Fix overly verbose drawing --- crazyflow/sim/visualize.py | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index e02325c..a0b17ca 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -119,7 +119,7 @@ def draw_capsule( p2: NDArray, radius: float = 0.05, rgba: NDArray | None = None, - is_cylinder: bool = False, + cylinder: bool = False, ): """Draw a capsule (pill) or cylinder between two points. @@ -129,31 +129,18 @@ def draw_capsule( p2: End point [3,] radius: The thickness of the geom in [m]. rgba: The color of the object. - is_cylinder: If True, draws a flat-ended cylinder. - If False, draws a pill-shaped capsule. + cylinder: If True, draws a flat-ended cylinder. If False, draws a pill-shaped capsule. """ if sim.viewer is None: return - # 1. Calculate Midpoint (Center of the geom) - pos = (p1 + p2) / 2.0 - - # 2. Calculate Half-length (MuJoCo uses half-extents) - dist = np.linalg.norm(p2 - p1) - half_length = dist / 2.0 - - # 3. Define Size: [radius, radius, half_length] - # Note: For capsules, size[2] is the half-length of the *cylindrical* part. - # MuJoCo adds the hemispherical caps on top of this. + pos = (p1 + p2) / 2.0 # Center of the geom + half_length = np.linalg.norm(p2 - p1) / 2.0 # MuJoCo uses half-extents size = np.array([radius, half_length, 0]) - - # 4. Get Rotation (Align Z-axis to the vector p2-p1) - # Using your existing helper (wrapped in a list for the reshape) + # Align the z-axis of the geom to the vector from p1 to p2 mat = _rotation_matrix_from_points(p1[None, :], p2[None, :]).as_matrix().flatten() - - geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER if is_cylinder else mujoco.mjtGeom.mjGEOM_CAPSULE - rgba = rgba if rgba is not None else np.array([0, 1.0, 0, 1]) - + geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER if cylinder else mujoco.mjtGeom.mjGEOM_CAPSULE + rgba = rgba if rgba is not None else np.array([1, 0, 0, 1.0]) sim.viewer.viewer.add_marker(type=geom_type, pos=pos, size=size, mat=mat, rgba=rgba) From 8b681bfb72628fcf32760f7266eb9d474b379824 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Wed, 18 Mar 2026 00:14:16 +0100 Subject: [PATCH 4/5] Add visualization tests --- crazyflow/sim/visualize.py | 44 --------------- tests/conftest.py | 5 ++ tests/unit/test_sim.py | 7 +-- tests/unit/test_visualizations.py | 89 +++++++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 50 deletions(-) create mode 100644 tests/unit/test_visualizations.py diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index a0b17ca..00c5d15 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -6,50 +6,6 @@ from crazyflow.sim import Sim -def draw_capsule( - sim: Sim, - p1: NDArray, - p2: NDArray, - radius: float = 0.05, - rgba: NDArray | None = None, - is_cylinder: bool = False, -): - """Draw a capsule (pill) or cylinder between two points. - - Args: - sim: The simulation. - p1: Start point [3,] - p2: End point [3,] - radius: The thickness of the geom. - rgba: The color of the object. - is_cylinder: If True, draws a flat-ended cylinder. - If False, draws a pill-shaped capsule. - """ - if sim.viewer is None: - return - - # 1. Calculate Midpoint (Center of the geom) - pos = (p1 + p2) / 2.0 - - # 2. Calculate Half-length (MuJoCo uses half-extents) - dist = np.linalg.norm(p2 - p1) - half_length = dist / 2.0 - - # 3. Define Size: [radius, radius, half_length] - # Note: For capsules, size[2] is the half-length of the *cylindrical* part. - # MuJoCo adds the hemispherical caps on top of this. - size = np.array([radius, half_length, 0]) - - # 4. Get Rotation (Align Z-axis to the vector p2-p1) - # Using your existing helper (wrapped in a list for the reshape) - mat = _rotation_matrix_from_points(p1[None, :], p2[None, :]).as_matrix().flatten() - - geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER if is_cylinder else mujoco.mjtGeom.mjGEOM_CAPSULE - rgba = rgba if rgba is not None else np.array([0, 1.0, 0, 1]) - - sim.viewer.viewer.add_marker(type=geom_type, pos=pos, size=size, mat=mat, rgba=rgba) - - def draw_line( sim: Sim, points: NDArray, diff --git a/tests/conftest.py b/tests/conftest.py index 0599b50..bfb2fd6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,3 +31,8 @@ def device() -> str: if "gpu" in available_backends(): return "gpu" return "cpu" + + +def skip_headless(): + if os.environ.get("DISPLAY") is None: + pytest.skip("DISPLAY is not set, skipping test in headless environment") diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index c54ab64..18e65ab 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -2,7 +2,6 @@ from __future__ import annotations -import os from typing import TYPE_CHECKING import jax @@ -10,6 +9,7 @@ import mujoco import numpy as np import pytest +from conftest import skip_headless from jax import Array from crazyflow.control import Control @@ -23,11 +23,6 @@ from typing import Any -def skip_headless(): - if os.environ.get("DISPLAY") is None: - pytest.skip("DISPLAY is not set, skipping test in headless environment") - - def array_meta_assert( x: Array, shape: tuple[int, ...] | None = None, diff --git a/tests/unit/test_visualizations.py b/tests/unit/test_visualizations.py new file mode 100644 index 0000000..4d34a91 --- /dev/null +++ b/tests/unit/test_visualizations.py @@ -0,0 +1,89 @@ +import numpy as np +import pytest +from conftest import skip_headless + +from crazyflow import Sim +from crazyflow.sim.visualize import draw_capsule, draw_line, draw_points + + +@pytest.mark.unit +def test_draw_capsule(device: str): + """Test drawing a capsule and verify it changes the rendering.""" + sim = Sim(device=device) + # Render without drawing + sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer + img_before = sim.render(mode="rgb_array", width=120, height=90) + # Draw a capsule and render + p1 = np.array([0.0, 0.0, 0.5]) + p2 = np.array([0.5, 0.5, 1.0]) + rgba = np.array([1.0, 0.0, 0.0, 1.0]) + draw_capsule(sim, p1, p2, radius=0.05, rgba=rgba) + img_after = sim.render(mode="rgb_array", width=120, height=90) + # Verify that the image changed + assert not np.array_equal(img_before, img_after), "Drawing capsule should change the rendering" + sim.close() + + +@pytest.mark.unit +def test_draw_capsule_cylinder(device: str): + """Test drawing a cylinder.""" + skip_headless() + sim = Sim(device=device) + sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer + img_before = sim.render(mode="rgb_array", width=120, height=90) + # Draw a cylinder instead of capsule + p1 = np.array([0.0, 0.0, 0.5]) + p2 = np.array([0.0, 0.0, 1.0]) + rgba = np.array([0.0, 1.0, 0.0, 1.0]) + draw_capsule(sim, p1, p2, radius=0.05, rgba=rgba, cylinder=True) + img_after = sim.render(mode="rgb_array", width=120, height=90) + assert not np.array_equal(img_before, img_after), "Drawing cylinder should change the rendering" + sim.close() + + +@pytest.mark.unit +def test_draw_line(device: str): + """Test drawing a line and verify it changes the rendering.""" + skip_headless() + sim = Sim(device=device) + sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer + img_before = sim.render(mode="rgb_array", width=120, height=90) + # Draw a line with multiple points + points = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 2.0]]) + rgba = np.array([0.0, 0.0, 1.0, 1.0]) + draw_line(sim, points, rgba=rgba, start_size=5.0, end_size=2.0) + img_after = sim.render(mode="rgb_array", width=120, height=90) + assert not np.array_equal(img_before, img_after), "Drawing line should change the rendering" + sim.close() + + +@pytest.mark.unit +def test_draw_points(device: str): + """Test drawing points and verify it changes the rendering.""" + skip_headless() + sim = Sim(device=device) + sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer + img_before = sim.render(mode="rgb_array", width=120, height=90) + # Draw multiple points + points = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.5], [0.0, 0.0, 1.0]]) + rgba = np.array([1.0, 1.0, 0.0, 1.0]) + draw_points(sim, points, rgba=rgba, size=0.02) + img_after = sim.render(mode="rgb_array", width=120, height=90) + assert not np.array_equal(img_before, img_after), "Drawing points should change the rendering" + sim.close() + + +@pytest.mark.unit +def test_draw_combined(device: str): + """Test drawing multiple visualization elements together.""" + skip_headless() + sim = Sim(device=device) + p1 = np.array([0.0, 0.0, 0.5]) + p2 = np.array([0.2, 0.0, 0.8]) + draw_capsule(sim, p1, p2, radius=0.03, rgba=np.array([1.0, 0.0, 0.0, 1.0])) + line_points = np.array([[0.3, 0.0, 0.5], [0.3, 0.2, 0.7], [0.5, 0.2, 0.9]]) + draw_line(sim, line_points, rgba=np.array([0.0, 1.0, 0.0, 1.0])) + points = np.array([[0.6, 0.0, 0.6], [0.7, 0.1, 0.7], [0.8, 0.0, 0.8]]) + draw_points(sim, points, rgba=np.array([0.0, 0.0, 1.0, 1.0]), size=0.025) + sim.render(mode="rgb_array", width=120, height=90) + sim.close() From df7ff7382ce8e672ed11edd0009b6c79e7096915 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Wed, 18 Mar 2026 00:22:13 +0100 Subject: [PATCH 5/5] Add headless skip marker --- tests/conftest.py | 7 +++++++ tests/unit/test_sim.py | 6 +++--- tests/unit/test_visualizations.py | 11 ++++++----- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index bfb2fd6..752b267 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -36,3 +36,10 @@ def device() -> str: def skip_headless(): if os.environ.get("DISPLAY") is None: pytest.skip("DISPLAY is not set, skipping test in headless environment") + + +# Marker for conditional skip in headless environments +skip_if_headless = pytest.mark.skipif( + os.environ.get("DISPLAY") is None, + reason="DISPLAY is not set, skipping test in headless environment", +) diff --git a/tests/unit/test_sim.py b/tests/unit/test_sim.py index 18e65ab..cf96ece 100644 --- a/tests/unit/test_sim.py +++ b/tests/unit/test_sim.py @@ -9,7 +9,7 @@ import mujoco import numpy as np import pytest -from conftest import skip_headless +from conftest import skip_if_headless from jax import Array from crazyflow.control import Control @@ -267,15 +267,15 @@ def test_sim_state_control_device(device: str): @pytest.mark.render +@skip_if_headless def test_render_human(device: str): sim = Sim(device=device) sim.render() sim.viewer.close() -# Do not mark as render to ensure it runs by default. This function will not open a viewer. +@skip_if_headless def test_render_rgb_array(device: str): - skip_headless() sim = Sim(n_worlds=2, device=device) img = sim.render(mode="rgb_array", width=1024, height=1024) assert isinstance(img, np.ndarray), "Image must be a numpy array" diff --git a/tests/unit/test_visualizations.py b/tests/unit/test_visualizations.py index 4d34a91..cbbf4dd 100644 --- a/tests/unit/test_visualizations.py +++ b/tests/unit/test_visualizations.py @@ -1,12 +1,13 @@ import numpy as np import pytest -from conftest import skip_headless +from conftest import skip_if_headless from crazyflow import Sim from crazyflow.sim.visualize import draw_capsule, draw_line, draw_points @pytest.mark.unit +@skip_if_headless def test_draw_capsule(device: str): """Test drawing a capsule and verify it changes the rendering.""" sim = Sim(device=device) @@ -25,9 +26,9 @@ def test_draw_capsule(device: str): @pytest.mark.unit +@skip_if_headless def test_draw_capsule_cylinder(device: str): """Test drawing a cylinder.""" - skip_headless() sim = Sim(device=device) sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer img_before = sim.render(mode="rgb_array", width=120, height=90) @@ -42,9 +43,9 @@ def test_draw_capsule_cylinder(device: str): @pytest.mark.unit +@skip_if_headless def test_draw_line(device: str): """Test drawing a line and verify it changes the rendering.""" - skip_headless() sim = Sim(device=device) sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer img_before = sim.render(mode="rgb_array", width=120, height=90) @@ -58,9 +59,9 @@ def test_draw_line(device: str): @pytest.mark.unit +@skip_if_headless def test_draw_points(device: str): """Test drawing points and verify it changes the rendering.""" - skip_headless() sim = Sim(device=device) sim.render(mode="rgb_array", width=120, height=90) # Warm up the renderer img_before = sim.render(mode="rgb_array", width=120, height=90) @@ -74,9 +75,9 @@ def test_draw_points(device: str): @pytest.mark.unit +@skip_if_headless def test_draw_combined(device: str): """Test drawing multiple visualization elements together.""" - skip_headless() sim = Sim(device=device) p1 = np.array([0.0, 0.0, 0.5]) p2 = np.array([0.2, 0.0, 0.8])