Add capsule draw function

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,

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",
+)

tests/unit/test_sim.py

@@ -2,14 +2,14 @@
 
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING
 
 import jax
 import jax.numpy as jnp
 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"

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()