[MJX] Best method for using mjx.ray in RL multi-environment training without running out of memory

[AlexS28]

Intro

Hi!

I am a PostDoct using MuJoCo for writing research papers in the area of manipulation.

My setup

Nvidia 4090 GPU

My question

I copy and pasted my code below which initializes a ray function which calls mjx.ray, and then I also copy and pasted the function which I call in my get observation function used in my RL pipeline (mjx), called compute_camera_ray_hits_mjx. Currently, my assumed resolution is low (40,40), and number of environments is 128 (also quite low). I flatten the dists variable and include that in my observation space. Note, even if I change my resolution to something dramatic such as (2,2) or lower number of environments to 2, I still get memory issues. I must be doing some wrong logic in my implementation. It does work if I make my number of environments equal to 1, so there's likely some vmap or jit issue? Any help would be greatly appreciated. I also tried without Jit and vmap, but it took too long to initialize so I am not sure if I would have run into memory issue.

2025-06-05 20:30:32.817821: W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3021] Can't reduce memory use below 16.64GiB (17871886491 bytes) by rematerialization; only reduced to 26.34TiB (28956816731264 bytes), down from 26.34TiB (28958311494848 bytes) originally

Minimal model and/or code that explain my question

def _initialize_raycaster(self, model, data, cam_id, resolution: Tuple[int, int] = (40, 40),):
        self.H, self.W = resolution

        # Camera intrinsics
        fovy = model.cam_fovy[cam_id] * jnp.pi / 180.0
        fy = 0.5 * self.H / jnp.tan(0.5 * fovy)
        fx = fy

        # Pixel grid
        i = jnp.arange(self.W)
        j = jnp.arange(self.H)
        ii, jj = jnp.meshgrid(i, j, indexing="xy")  # shape (H, W)

        x = (ii - self.W / 2) / fx
        y = -(jj - self.H / 2) / fy
        z = -jnp.ones_like(x)

        dirs = jnp.stack([x, y, z], axis=-1)  # shape (H, W, 3)
        dirs = dirs / jnp.linalg.norm(dirs, axis=-1, keepdims=True)
        dirs_flat = dirs.reshape(-1, 3)

        # Camera pose
        cam_pos = data.cam_xpos[cam_id]  # shape (3,)
        cam_mat = data.cam_xmat[cam_id].reshape(3, 3)  # shape (3, 3)

        self.dirs_world = dirs_flat @ cam_mat.T  # shape (H*W, 3)
        self.origins_world = jnp.repeat(cam_pos[None, :], self.dirs_world.shape[0], axis=0)

        # Ray function with static geomgroup
        self.ray_fn = jax.jit(mjx.ray, static_argnums=(4,))

def compute_camera_ray_hits_mjx(
        self,
        model: mjx.Model,
        data: mjx.Data,
        geomgroup: Tuple[int, int, int, int, int, int] = (1, 1, 1, 1, 1, 1),
) -> Tuple[jax.Array, jax.Array]:
    """
    Perform fully JAX-compatible raycasting from a camera into the MJX scene.

    Args:
        model: mjx.Model
        data: mjx.Data
        cam_id: camera index
        resolution: image resolution (H, W)
        geomgroup: tuple of 6 binary values to filter geom groups

    Returns:
        dists: [H, W] distances to hit points (0.0 if no hit)
        geom_ids: [H, W] geom IDs hit (-1 if no hit)
    """

    dist, geomid = jax.vmap(lambda o, d: self.ray_fn(model, data, o, d, geomgroup))(
        self.origins_world, self.dirs_world
    )

    dists = dist.reshape(self.H, self.W)
    geom_ids = geomid.reshape(self.H, self.W)

    # Replace -1 (no hit) in distances with 0.0
    dists = jnp.where(dists < 0, 0.0, dists)

    return dists, geom_ids

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.