Add height field ray intersection support to MJX ray functions

[achaljhawar]

Description

This PR adds support for height field ray intersection testing to the MJX ray intersection functions. This feature enables depth rendering in environments with terrain (height fields), which is particularly useful for robot locomotion tasks.

Motivation

The mjx.ray function is essential for getting depth images in MJX and is used by the rangefinder sensor, but it currently doesn't support hfield geoms. This PR ports the mj_rayHfield function from MuJoCo to MJX to enable ray intersection tests with height fields.

Implementation Details

• Added ray_hfield function to calculate ray intersection with height fields
• Created _ray_hfield_wrapper to adapt the function to the expected interface
• Updated _RAY_FUNC dictionary to include height field geoms
• Modified the ray function to handle height field geoms similar to mesh geoms
• Updated ray_geom to handle height field cases appropriately

The implementation follows the ray marching approach from the C implementation but adapts it to JAX's functional programming model with vectorized operations where possible.

Testing

Tested with height field environments to ensure correct depth values when using rangefinder sensors over terrain.

Fixes #2155

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[achaljhawar]

I have signed the cla