Add batched raycasting sensor support to crazyflow (#48)

Author: amacati

.github/copilot-instructions.md

@@ -0,0 +1,7 @@
+# General
+- Prefer failing fast over capturing exceptions and continuing execution.
+- Use pure functions where possible.
+- Prefer protocols over inheritance.
+- Use pathlib instead of os.path for file path manipulations.
+- Use napoleon-type docstrings for documenting functions and classes.
+- Use type hints for function signatures. Do not put types into the docstrings.

crazyflow/sim/sensors.py

@@ -0,0 +1,103 @@
+from functools import partial
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+import mujoco.mjx as mjx
+from jax import Array
+
+from crazyflow.sim.sim import Sim, requires_mujoco_sync
+
+
+@requires_mujoco_sync
+def render_depth(
+    sim: Sim, camera: int = 0, resolution: tuple[int, int] = (100, 100), include_drone: bool = False
+) -> Array:
+    """Render depth images using raycasting.
+
+    Note:
+        Code has been adoped from
+        https://github.com/Andrew-Luo1/jax_shac/blob/main/vision/2dof_ball.ipynb
+    """
+    return _render_depth(
+        mjx_data=sim.mjx_data,
+        mjx_model=sim.mjx_model,
+        camera=camera,
+        resolution=resolution,
+        include_drone=include_drone,
+    )
+
+
+def build_render_depth_fn(
+    mjx_model: mjx.Model,
+    camera: int = 0,
+    resolution: tuple[int, int] = (100, 100),
+    geomgroup: tuple[int, ...] = (1, 1, 0, 0, 1, 1, 1, 1),
+) -> Callable[[Sim], Array]:
+    """Build a depth renderer function for given camera and resolution.
+
+    Compiles the mjx model and rays directly into the rendering function for higher performance. The
+    returned function takes a Sim object as input and returns depth images.
+    """
+    rays = _camera_rays(resolution=resolution, fov_y=jnp.pi / 4)[None, ...]
+    ray_fn = jax.jit(
+        partial(_render_rays, mjx_model=mjx_model, camera=camera, geomgroup=geomgroup, rays=rays),
+        static_argnames=("mjx_model", "camera", "geomgroup", "rays"),
+    )
+
+    @requires_mujoco_sync
+    def render_depth_fn(sim: Sim) -> Array:
+        return ray_fn(mjx_data=sim.mjx_data)
+
+    return render_depth_fn
+
+
+@jax.jit(static_argnames=("camera", "resolution", "include_drone"))
+def _render_depth(
+    mjx_data: mjx.Data,
+    mjx_model: mjx.Model,
+    camera: int,
+    resolution: tuple[int, int],
+    include_drone: bool = False,
+) -> Array:
+    """Accelerates the dynamic rendering of depth images."""
+    local_rays = _camera_rays(resolution=resolution, fov_y=jnp.pi / 4)[None, ...]
+    geomgroup = (1, 1, 1, 0, 1, 1, 1, 1) if include_drone else (1, 1, 0, 0, 1, 1, 1, 1)
+    return _render_rays(
+        mjx_data=mjx_data, mjx_model=mjx_model, camera=camera, rays=local_rays, geomgroup=geomgroup
+    )
+
+
+def _render_rays(
+    mjx_data: mjx.Data, mjx_model: mjx.Model, camera: int, rays: Array, geomgroup: tuple[int, ...]
+) -> Array:
+    """Render a given ray array using MuJoCo's raycasting."""
+    rays = _to_mjx_frame(rays, mjx_data.cam_xmat[:, camera])
+    ray_ax = (None, None, None, 0)
+    ray = jax.vmap(
+        jax.vmap(jax.vmap(partial(mjx.ray, geomgroup=geomgroup), in_axes=ray_ax), in_axes=ray_ax),
+        in_axes=(None, 0, 0, 0),
+    )
+    return ray(mjx_model, mjx_data, mjx_data.cam_xpos[:, camera], rays)[0]
+
+
+def _to_mjx_frame(x: Array, xmat: Array) -> Array:
+    """Transform points to a different frame given its rotation matrix."""
+    return (xmat[:, None, None, ...] @ x[..., None])[..., 0]
+
+
+def _camera_rays(resolution: tuple[int, int] = (100, 100), fov_y: float = jnp.pi / 4) -> Array:
+    """Create an array of rays with a given field of view and resolution.
+
+    Args:
+        resolution: Image resolution as (width, height).
+        fov_y: Vertical field of view in radians.
+    """
+    image_height = jnp.tan(fov_y / 2) * 2
+    image_width = image_height * (resolution[0] / resolution[1])  # Square pixels.
+    delta = image_width / (2 * resolution[0])
+    x = jnp.linspace(-image_width / 2 + delta, image_width / 2 - delta, resolution[0])
+    y = jnp.flip(jnp.linspace(-image_height / 2 + delta, image_height / 2 - delta, resolution[1]))
+    X, Y = jnp.meshgrid(x, y)
+    rays = jnp.stack([X, Y, -jnp.ones_like(X)], axis=-1)
+    return rays / jnp.linalg.norm(rays, axis=-1, keepdims=True)

crazyflow/sim/sim.py

@@ -232,11 +232,10 @@ def build_mjx_spec(self) -> mujoco.MjSpec:
         spec.copy_during_attach = True
         drone_spec = mujoco.MjSpec.from_file(str(self.drone_path))
         frame = spec.worldbody.add_frame(name="world")
+        if (drone_body := drone_spec.body("drone")) is None:
+            raise ValueError("Drone body not found in drone spec")
         # Add drones and their actuators
         for i in range(self.n_drones):
-            drone_body = drone_spec.body("drone")
-            if drone_body is None:
-                raise ValueError("Drone body not found in drone spec")
             drone = frame.attach_body(drone_body, "", f":{i}")
             drone.add_freejoint()
         return spec
@@ -534,6 +533,8 @@ def sync_sim2mjx(data: SimData, mjx_data: Data, mjx_model: Model) -> tuple[SimDa
     qvel = rearrange(jnp.concat([vel, ang_vel], axis=-1), "w d qvel -> w (d qvel)")
     mjx_data = mjx_data.replace(qpos=qpos, qvel=qvel)
     mjx_data = jax.vmap(mjx.kinematics, in_axes=(None, 0))(mjx_model, mjx_data)
+    # Required for rendering w. ray casting
+    mjx_data = jax.vmap(mjx.camlight, in_axes=(None, 0))(mjx_model, mjx_data)
     mjx_data = jax.vmap(mjx.collision, in_axes=(None, 0))(mjx_model, mjx_data)
     data = data.replace(core=data.core.replace(mjx_synced=True))
     return data, mjx_data

examples/raycasting.py

@@ -0,0 +1,37 @@
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+
+from crazyflow.sim import Sim
+from crazyflow.sim.sensors import build_render_depth_fn, render_depth
+
+
+def main(plot: bool = False):
+    sim = Sim()
+    sim.data = sim.data.replace(
+        states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.2))
+    )
+    # The easiest way to get depth images is to use the render_depth function
+    dist = render_depth(sim, camera=0, resolution=(100, 100), include_drone=False)
+    dist = dist.at[dist > 1.5].set(jnp.nan)  # Cap max distance for better visualization
+    if plot:
+        plt.imshow(dist[0], cmap="viridis")
+        plt.colorbar(label="Distance (m)")
+        plt.title("Raycast Distance from Camera")
+        plt.show()
+    # We can also build a depth renderer function for better performance if we need maximum speed or
+    # more fine-grained control. Here we only render the drone collision geometry to avoid expensive
+    # raycasting against the high-poly visual mesh of the drone.
+    render_depth_fn = build_render_depth_fn(
+        sim.mjx_model, camera=0, resolution=(200, 200), geomgroup=(1, 1, 0, 1, 1, 1, 1, 1)
+    )
+    dist_fn = render_depth_fn(sim)
+    dist_fn = dist_fn.at[dist_fn > 1.5].set(jnp.nan)  # Cap max distance for better visualization
+    if plot:
+        plt.imshow(dist_fn[0], cmap="viridis")
+        plt.colorbar(label="Distance (m)")
+        plt.title("Raycast Distance from Camera (Compiled)")
+        plt.show()
+
+
+if __name__ == "__main__":
+    main(plot=True)

tests/unit/test_sensors.py

@@ -0,0 +1,52 @@
+"""Unit tests for the sensors module."""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from crazyflow.sim import Sim
+from crazyflow.sim.sensors import _camera_rays, build_render_depth_fn, render_depth
+
+
+@pytest.mark.unit
+def test_camera_rays():
+    """Test that camera_rays produces arrays with correct shape and device."""
+    resolution = (64, 48)
+    rays = _camera_rays(resolution=resolution)
+    # Check shape: should be (height, width, 3)
+    expected_shape = (resolution[1], resolution[0], 3)
+    assert rays.shape == expected_shape, f"Expected shape {expected_shape}, got {rays.shape}"
+    # Check that rays are normalized
+    norm = jnp.linalg.norm(rays, axis=-1)
+    assert jnp.allclose(norm, 1.0, atol=1e-6), "Rays should be normalized"
+    # Check that rays respect the FOV
+    rays_narrow = _camera_rays(fov_y=np.pi / 6)
+    rays_wide = _camera_rays(fov_y=np.pi / 3)
+    # Corner rays should have different angles for different FOV
+    # Check the top corner ray y-component (wider FOV should have larger y-component)
+    corner_y_narrow = abs(rays_narrow[0, 0, 1])  # Top-left corner
+    corner_y_wide = abs(rays_wide[0, 0, 1])
+    assert corner_y_wide > corner_y_narrow, "Wider FOV should produce rays with larger y-components"
+
+
+@pytest.mark.unit
+def test_render_depth(device: str):
+    """Test render_depth with different resolutions."""
+    sim = Sim(n_worlds=2, device=device)
+    dist = render_depth(sim, camera=0, resolution=(10, 10))
+    assert dist.shape == (2, 10, 10), f"Expected shape (2, 10, 10), got {dist.shape}"
+    assert dist.device == jax.devices(device)[0], f"Expected device {device}, got {dist.device}"
+
+
+@pytest.mark.unit
+def test_build_render_depth_fn():
+    """Test build_render_depth_fn produces a callable that returns correct shapes."""
+    sim = Sim(n_worlds=3)
+    render_depth_fn = build_render_depth_fn(
+        sim.mjx_model, camera=0, resolution=(20, 15), geomgroup=(1, 1, 0, 1, 1, 1, 1, 1)
+    )
+    dist = render_depth_fn(sim)
+    assert dist.shape == (3, 15, 20), f"Expected shape (3, 15, 20), got {dist.shape}"