Add additional data fields that can be reported by rangefinder sensors.

PiperOrigin-RevId: 848316991
Change-Id: Idbf7ba81b4da711a22c23302c8782ab2b0b98d82

Author: yuvaltassa

doc/XMLreference.rst

@@ -6382,12 +6382,46 @@ This element creates a rangefinder.
   measurements in this case is equal to product of the camera's width and height
   :ref:`resolutions<body-camera-resolution>`.
 
+.. image:: images/XMLreference/rfcamera.png
+   :width: 45%
+   :align: right
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/sensor/rfcamera.xml
+
 If a ray does not intersect any geom surface, the sensor output is -1. If the origin of the ray is inside a geom, the
 surface is still detected. Geoms attached to the same body as the sensor site/camera are excluded. Invisible geoms,
 defined as geoms whose rgba (or whose material rgba) has alpha=0, are also excluded. Note however that geoms made
 invisible in the visualizer by disabling their geom group are not excluded; this is because sensor calculations are
 independent of the visualizer.
 
+The image on the right (click to see the model being visualized) shows two rangefinder sensors attached to a perspective and
+an orthographic camera, with frustums visualized. Both cameras have 4x4 resolution, for 16 rays each. The rangefinder
+sensors report :at:`data` = :at-val:`"dist point normal"` (see below), so we can see the rays (lines), the intersection
+points (spheres) and the surface normals (arrows).
+
+.. _sensor-rangefinder-data:
+
+:at:`data`: :at-val:`[dist, dir, origin, point, normal, depth], "dist"`
+   By default, the rangefinder outputs a distance measurement, as described above. However, it is also possible to
+   specify a set of output data fields. The :at:`data` attribute can contain **multiple sequential data types**, as long
+   as the relative order---as listed above---is maintained. For example, :at:`data` = :at-val:`"dist point normal"` will
+   return 7 numbers per ray, while :at:`data` = :at-val:`"point origin"` is an error because :at-val:`origin` must come
+   before :at-val:`point`.
+
+   - :at-val:`dist` **real(1)**: The distance from the ray origin to the nearest geom surface, -1 if no surface was hit.
+     If this data type is included, rays will be visualized as lines.
+   - :at-val:`dir` **real(3)**: Normalized direction of the ray, or (0, 0, 0) if no surface was hit.
+   - :at-val:`origin` **real(3)**: The point from which the ray emanates (global frame). For sites and perspective
+     cameras, this is the site/camera xpos. However for orthographic cameras, ray origins are spatially distributed
+     along the image plane.
+   - :at-val:`point` **real(3)**: The point where the ray intersects the nearest geom surface in the global frame, or
+     (0, 0, 0) if no surface was hit. If this data type is included, intersection points will be visualized as spheres.
+   - :at-val:`normal`: **real(3)**: The geom surface normal at the point where the ray intersects it, in the global
+     frame, or (0, 0, 0) if no surface was hit. Note that normals always point towards the outside of the geom surface,
+     regardless of the ray origin. If this data type is included along with either :at-val:`dist` or :at-val:`point`,
+     normals will be visualized as arrows at the intersection points.
+   - :at-val:`depth`: **real(1)**: The distance of the hit point from the camera plane, -1 if no surface was hit. Note
+     that this depth sematic corresponds to depth images in the computer graphics sense.
+
 .. _sensor-rangefinder-name:
 
 .. _sensor-rangefinder-noise:

doc/XMLschema.rst

@@ -940,9 +940,9 @@
 | :ref:`rangefinder                  | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <sensor-rangefinder>`              |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`name<sensor-rangefinder-name>`                            | :ref:`site<sensor-rangefinder-site>`                            | :ref:`camera<sensor-rangefinder-camera>`                        | :ref:`cutoff<sensor-rangefinder-cutoff>`                        | |
+|                                    |    |    | :ref:`name<sensor-rangefinder-name>`                            | :ref:`site<sensor-rangefinder-site>`                            | :ref:`camera<sensor-rangefinder-camera>`                        | :ref:`data<sensor-rangefinder-data>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`noise<sensor-rangefinder-noise>`                          | :ref:`user<sensor-rangefinder-user>`                            |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`cutoff<sensor-rangefinder-cutoff>`                        | :ref:`noise<sensor-rangefinder-noise>`                          | :ref:`user<sensor-rangefinder-user>`                            |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |

doc/changelog.rst

@@ -18,16 +18,24 @@ Upcoming version (not yet released)
 
 General
 ^^^^^^^
+
+.. image:: images/XMLreference/rfcamera.png
+   :width: 45%
+   :align: right
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/sensor/rfcamera.xml
+
 - Camera frustum visualization is now triggered by setting :ref:`resolution<body-camera-resolution>` to values larger
   than 1. Relatedly, frustum visualization also works for :ref:`orthographic<body-camera-projection>` cameras.
 - Rangefinder sensors can now be attached to a camera using the :ref:`ragefinder/camera<sensor-rangefinder-camera>`
-  attribute. In this case, the sensor will cast multiple ray, one for each camera pixel.
+  attribute. In this case, the sensor will cast multiple ray, one for each pixel.
+- Rangefinder sensors can now now report various kinds of information besides ray distances, including surface normals.
+  See :ref:`rangefinder<sensor-rangefinder>` for details.
 - Non-breaking ABI changes:
 
   - The type of the ``sig`` (signature) argument of :ref:`mj_stateSize` and related functions has been changed from
     ``unsigned int`` to ``int``. Before this change, invalid negative arguments passed to this function would result in
     a silent implicit cast, now negativity will trigger an error.
-  - Added a :ref:`depth<mjtRndFlag>` rendering flag
+  - Added a :ref:`depth<mjtRndFlag>` rendering flag.
 
 MJX
 ^^^

doc/images/XMLreference/rfcamera.png

@@ -774,8 +774,18 @@ typedef enum mjtConDataField_ {   // data fields returned by contact sensors
   mjCONDATA_NORMAL,               // contact frame normal
   mjCONDATA_TANGENT,              // contact frame first tangent
 
-  mjNCONDATA          = 7         // number of contact sensor data fields
+  mjNCONDATA                      // number of contact sensor data fields
 } mjtConDataField;
+typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
+  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
+  mjRAYDATA_DIR,                  // normalized ray direction
+  mjRAYDATA_ORIGIN,               // ray origin
+  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
+  mjRAYDATA_NORMAL,               // surface normal at intersection point
+  mjRAYDATA_DEPTH,                // depth along z-axis
+
+  mjNRAYDATA                      // number of rangefinder sensor data fields
+} mjtRayDataField;
 typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
   mjSAMEFRAME_NONE    = 0,        // no alignment
   mjSAMEFRAME_BODY,               // frame is same as body frame

doc/includes/references.h

@@ -400,10 +400,22 @@ typedef enum mjtConDataField_ {   // data fields returned by contact sensors
   mjCONDATA_NORMAL,               // contact frame normal
   mjCONDATA_TANGENT,              // contact frame first tangent
 
-  mjNCONDATA          = 7         // number of contact sensor data fields
+  mjNCONDATA                      // number of contact sensor data fields
 } mjtConDataField;
 
 
+typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
+  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
+  mjRAYDATA_DIR,                  // normalized ray direction
+  mjRAYDATA_ORIGIN,               // ray origin
+  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
+  mjRAYDATA_NORMAL,               // surface normal at intersection point
+  mjRAYDATA_DEPTH,                // depth along z-axis
+
+  mjNRAYDATA                      // number of rangefinder sensor data fields
+} mjtRayDataField;
+
+
 typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
   mjSAMEFRAME_NONE    = 0,        // no alignment
   mjSAMEFRAME_BODY,               // frame is same as body frame

include/mujoco/mjmodel.h

@@ -416,6 +416,20 @@
              ('mjNCONDATA', 7),
          ]),
      )),
+    ('mjtRayDataField',
+     EnumDecl(
+         name='mjtRayDataField',
+         declname='enum mjtRayDataField_',
+         values=dict([
+             ('mjRAYDATA_DIST', 0),
+             ('mjRAYDATA_DIR', 1),
+             ('mjRAYDATA_ORIGIN', 2),
+             ('mjRAYDATA_POINT', 3),
+             ('mjRAYDATA_NORMAL', 4),
+             ('mjRAYDATA_DEPTH', 5),
+             ('mjNRAYDATA', 6),
+         ]),
+     )),
     ('mjtSameFrame',
      EnumDecl(
          name='mjtSameFrame',

python/mujoco/introspect/enums.py

@@ -1662,5 +1662,126 @@ def test_from_zip(self):
         string_spec.compile()
         self.assertEqual(spec.to_xml(), string_spec.to_xml())
 
+  def test_rangefinder_sensor(self):
+    """Test rangefinder sensor with mjSpec, iterative model building."""
+    # Raydata field enum values for dataspec bitfield
+    rd = mujoco.mjtRayDataField
+    dist_val = int(rd.mjRAYDATA_DIST)
+    dir_val = int(rd.mjRAYDATA_DIR)
+    origin_val = int(rd.mjRAYDATA_ORIGIN)
+    point_val = int(rd.mjRAYDATA_POINT)
+    normal_val = int(rd.mjRAYDATA_NORMAL)
+    depth_val = int(rd.mjRAYDATA_DEPTH)
+
+    # Step 1: Create a rangefinder sensor attached to a site, no dataspec set.
+    # Note: site goes on a child body because rangefinder excludes the site's
+    # parent body from ray casting.
+    spec = mujoco.MjSpec()
+    sensor_body = spec.worldbody.add_body(name='sensor_body', pos=[0, 0, 1])
+    sensor_body.add_site(name='rf_site', zaxis=[0, 0, -1])
+    rf_sensor = spec.add_sensor(
+        name='rf',
+        type=mujoco.mjtSensor.mjSENS_RANGEFINDER,
+        objtype=mujoco.mjtObj.mjOBJ_SITE,
+        objname='rf_site',
+    )
+
+    # This should fail: data spec (intprm[0]) must be positive
+    with self.assertRaisesWithPredicateMatch(
+        ValueError,
+        lambda e: 'data spec (intprm[0]) must be positive' in str(e)
+    ):
+      spec.compile()
+
+    # Step 2: Set dataspec to just mjRAYDATA_DIST
+    rf_sensor.intprm[0] = 1 << dist_val
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+
+    # With no geometry, the ray should miss: dist = -1
+    self.assertEqual(model.nsensordata, 1)
+    self.assertEqual(data.bind(rf_sensor).data[0], -1)
+
+    # Step 3: Add all raydata fields and check no-hit values
+    all_fields = (
+        (1 << dist_val) | (1 << dir_val) | (1 << origin_val) |
+        (1 << point_val) | (1 << normal_val) | (1 << depth_val)
+    )
+    rf_sensor.intprm[0] = all_fields
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+
+    # Expected size: dist(1) + dir(3) + origin(3) + point(3) + normal(3) +
+    # depth(1) = 14
+    self.assertEqual(model.nsensordata, 14)
+
+    # No-hit values
+    sd = data.bind(rf_sensor).data
+    self.assertEqual(sd[0], -1)  # dist
+    np.testing.assert_allclose(sd[1:4], [0, 0, 0])  # dir
+    np.testing.assert_allclose(sd[4:7], [0, 0, 1])  # origin
+    np.testing.assert_allclose(sd[7:10], [0, 0, 0])  # point
+    np.testing.assert_allclose(sd[10:13], [0, 0, 0])  # normal
+    self.assertEqual(sd[13], -1)  # depth
+
+    # Step 4: Add a floor plane, now the ray should hit
+    spec.worldbody.add_geom(
+        name='floor',
+        type=mujoco.mjtGeom.mjGEOM_PLANE,
+        size=[10, 10, 0.1],
+    )
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+
+    # Ray starts at z=1 pointing down, hits floor at z=0
+    # For site sensor, depth = dist
+    sd = data.bind(rf_sensor).data
+    self.assertAlmostEqual(sd[0], 1.0, places=6)  # dist
+    np.testing.assert_allclose(sd[1:4], [0, 0, -1], atol=1e-10)  # dir
+    np.testing.assert_allclose(sd[4:7], [0, 0, 1], atol=1e-10)  # origin
+    np.testing.assert_allclose(sd[7:10], [0, 0, 0], atol=1e-10)  # point
+    np.testing.assert_allclose(sd[10:13], [0, 0, 1], atol=1e-10)  # normal
+    self.assertAlmostEqual(sd[13], 1.0, places=6)  # depth
+
+    # Step 5: Add a camera-based rangefinder sensor
+    # Camera also on child body so it doesn't exclude the floor
+    cam_body = spec.worldbody.add_body(name='cam_body', pos=[0, 0, 2])
+    cam_body.add_camera(
+        name='rf_cam',
+        xyaxes=[1, 0, 0, 0, 1, 0],  # z=[0,0,1], looks along -z (down)
+        resolution=[3, 3],
+        fovy=90,
+    )
+    cam_sensor = spec.add_sensor(
+        name='rf_cam_sensor',
+        type=mujoco.mjtSensor.mjSENS_RANGEFINDER,
+        objtype=mujoco.mjtObj.mjOBJ_CAMERA,
+        objname='rf_cam',
+        intprm=[(1 << dist_val) | (1 << depth_val), 0, 0],
+    )
+
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+
+    # Site sensor: 14 values, Camera sensor: (1+1)*9 = 18 values
+    self.assertEqual(model.nsensordata, 14 + 18)
+
+    # Check camera sensor data using bind
+    cam_sd = data.bind(cam_sensor).data
+    stride = 2  # dist + depth per pixel
+    center_pixel = 4  # center of 3x3 = row 1, col 1
+
+    # Center pixel: ray straight down from z=2 to z=0
+    self.assertAlmostEqual(cam_sd[center_pixel * stride], 2.0, places=6)
+    self.assertAlmostEqual(cam_sd[center_pixel * stride + 1], 2.0, places=6)
+
+    # Corner pixel: off-axis ray, dist > depth
+    self.assertGreater(cam_sd[0], cam_sd[1])  # dist > depth
+    self.assertAlmostEqual(cam_sd[1], 2.0, places=6)  # depth is still 2.0
+
 if __name__ == '__main__':
   absltest.main()

python/mujoco/specs_test.py

@@ -662,11 +662,11 @@ static mjtNum mj_rayHfieldNormal(const mjModel* m, const mjData* d, int geomid,
       // triangle normal
       mjtNum normal_tri[3];
 
-      // first triangle
+      // first triangle: swap v1 and v2 for consistent CCW winding (normals point up)
       mjtNum va[3][3] = {
         {dx*c-size[0], dy*r-size[1], data[r*ncol+c]*size[2]},
-        {dx*(c+1)-size[0], dy*(r+1)-size[1], data[(r+1)*ncol+(c+1)]*size[2]},
-        {dx*(c+1)-size[0], dy*(r+0)-size[1], data[(r+0)*ncol+(c+1)]*size[2]}
+        {dx*(c+1)-size[0], dy*(r+0)-size[1], data[(r+0)*ncol+(c+1)]*size[2]},
+        {dx*(c+1)-size[0], dy*(r+1)-size[1], data[(r+1)*ncol+(c+1)]*size[2]}
       };
       mjtNum sol = ray_triangle(va, lpnt, lvec, b0, b1, normal ? normal_tri : NULL);
       if (sol >= 0 && (x < 0 || sol < x)) {
@@ -1458,6 +1458,18 @@ void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
         AABB[3] = mju_max(AABB[3], elevation);
       }
 
+      // add distance-dependent angular margin to account for edge/face curvature
+      // margin = atan(max_half_size / dist) bounds the angular deviation of face centers
+      mjtNum max_half = mju_max(aabb[3], mju_max(aabb[4], aabb[5]));
+      mjtNum dist = mju_dist3(pnt, xpos);
+      if (dist > mjMINVAL) {
+        mjtNum margin = mju_atan2(max_half, dist);
+        AABB[0] -= margin;
+        AABB[1] -= margin;
+        AABB[2] += margin;
+        AABB[3] += margin;
+      }
+
       // azimuth crosses discontinuity, fall back to no angular culling
       if (AABB[2]-AABB[0] > mjPI) {
         AABB[0] = -mjPI;
@@ -1520,8 +1532,23 @@ static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], co
 
       // exclude geom using bounding angles
       if (m->body_bvhadr[b] != -1) {
-        if (azimuth < (geom_ba+4*i)[0] || elevation < (geom_ba+4*i)[1] ||
-            azimuth > (geom_ba+4*i)[2] || elevation > (geom_ba+4*i)[3]) {
+        mjtNum az_min = (geom_ba+4*i)[0];
+        mjtNum az_max = (geom_ba+4*i)[2];
+        mjtNum el_min = (geom_ba+4*i)[1];
+        mjtNum el_max = (geom_ba+4*i)[3];
+
+        // check elevation
+        if (elevation < el_min || elevation > el_max) {
+          continue;
+        }
+
+        // check azimuth with wraparound
+        mjtNum az_center = (az_min + az_max) * 0.5;
+        mjtNum az_half_width = (az_max - az_min) * 0.5;
+        mjtNum az_diff = azimuth - az_center;
+        if (az_diff > mjPI) az_diff -= 2*mjPI;
+        else if (az_diff < -mjPI) az_diff += 2*mjPI;
+        if (mju_abs(az_diff) > az_half_width) {
           continue;
         }
       }