diff --git a/crazyflow/sim/visualize.py b/crazyflow/sim/visualize.py index 2c1c4de..00c5d15 100644 --- a/crazyflow/sim/visualize.py +++ b/crazyflow/sim/visualize.py @@ -69,6 +69,37 @@ 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, + 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. + cylinder: If True, draws a flat-ended cylinder. If False, draws a pill-shaped capsule. + """ + if sim.viewer is None: + return + + 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]) + # 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 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) + + def change_material( sim: Sim, mat_name: str, diff --git a/tests/conftest.py b/tests/conftest.py index 0599b50..752b267 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,3 +31,15 @@ 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") + + +# 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 c54ab64..cf96ece 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_if_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, @@ -272,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 new file mode 100644 index 0000000..cbbf4dd --- /dev/null +++ b/tests/unit/test_visualizations.py @@ -0,0 +1,90 @@ +import numpy as np +import pytest +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) + # 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 +@skip_if_headless +def test_draw_capsule_cylinder(device: str): + """Test drawing a cylinder.""" + 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 +@skip_if_headless +def test_draw_line(device: str): + """Test drawing a line and verify it changes the rendering.""" + 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 +@skip_if_headless +def test_draw_points(device: str): + """Test drawing points and verify it changes the rendering.""" + 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 +@skip_if_headless +def test_draw_combined(device: str): + """Test drawing multiple visualization elements together.""" + 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()