Improve raycasting performance

Author: amacati

crazyflow/sim/sensors.py

@@ -1,48 +1,92 @@
 from functools import partial
+from typing import Callable
 
 import jax
 import jax.numpy as jnp
 import mujoco.mjx as mjx
-import numpy as np
 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)) -> Array:
+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
     """
-    local_rays = camera_rays(resolution=resolution, fov_y=np.pi / 4)[None, ...]
-    rays = to_mjx_frame(local_rays, sim.mjx_data.cam_xmat[:, camera])
-    dist, _ = ray_fn(sim.mjx_model, sim.mjx_data, sim.mjx_data.cam_xpos[:, camera], rays)
-    return dist
+    return _render_depth(
+        mjx_data=sim.mjx_data,
+        mjx_model=sim.mjx_model,
+        camera=camera,
+        resolution=resolution,
+        include_drone=include_drone,
+    )
 
 
-@jax.jit
-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 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
 
 
-_geomgroup = (1, 1, 1, 0, 1, 1, 1, 1)  # Exclude collision geoms
-ray_fn = jax.jit(
-    jax.vmap(
-        jax.vmap(
-            jax.vmap(partial(mjx.ray, geomgroup=_geomgroup), in_axes=(None, None, None, 0)),
-            in_axes=(None, None, None, 0),
-        ),
+@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]
 
 
-@partial(jax.jit, static_argnames=("resolution", "fov_y"))
-def camera_rays(resolution: tuple[int, int] = (100, 100), fov_y: float = jnp.pi / 4) -> Array:
+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:

examples/raycasting.py

@@ -2,22 +2,35 @@
 import matplotlib.pyplot as plt
 
 from crazyflow.sim import Sim
-from crazyflow.sim.sensors import render_depth
+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))
     )
-    dist = render_depth(sim, camera=0, resolution=(100, 100))
+    # 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 not plot:
-        return
-    plt.imshow(dist[0], cmap="viridis")
-    plt.colorbar(label="Distance (m)")
-    plt.title("Raycast Distance from Camera")
-    plt.show()
+    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__":