Extend rangefinder sensor to support cameras.

PiperOrigin-RevId: 848255889
Change-Id: I6e1a9ed13d29d2242558625a47e24e47ad6e87db

src/engine/engine_sensor.c

@@ -435,11 +435,62 @@ void mj_sensorPos(const mjModel* m, mjData* d) {
         break;
 
       case mjSENS_RANGEFINDER:                            // rangefinder
-        rvec[0] = d->site_xmat[9*objid+2];
-        rvec[1] = d->site_xmat[9*objid+5];
-        rvec[2] = d->site_xmat[9*objid+8];
-        d->sensordata[adr] = mj_ray(m, d, d->site_xpos+3*objid, rvec, NULL, 1,
-                                    m->site_bodyid[objid], NULL);
+        if (objtype == mjOBJ_SITE) {
+          rvec[0] = d->site_xmat[9*objid+2];
+          rvec[1] = d->site_xmat[9*objid+5];
+          rvec[2] = d->site_xmat[9*objid+8];
+          d->sensordata[adr] = mj_ray(m, d, d->site_xpos+3*objid, rvec, NULL, 1,
+                                      m->site_bodyid[objid], NULL);
+        } else {
+          // camera-attached rangefinder: depth image
+          const int width = m->cam_resolution[2*objid];
+          const int height = m->cam_resolution[2*objid+1];
+          const int bodyexclude = m->cam_bodyid[objid];
+          const mjtNum* cam_xpos = d->cam_xpos + 3*objid;
+          const mjtNum* cam_xmat = d->cam_xmat + 9*objid;
+          const int projection = m->cam_projection[objid];
+
+          // compute focal length in pixels using helper
+          mjtNum fx, fy, cx, cy, ortho_extent;
+          mju_camIntrinsics(m, objid, &fx, &fy, &cx, &cy, &ortho_extent);
+
+          if (projection == mjPROJ_PERSPECTIVE) {
+            // perspective: all rays share origin, different directions
+            const int npixel = width * height;
+            mj_markStack(d);
+            mjtNum* vec = mjSTACKALLOC(d, 3*npixel, mjtNum);
+            int* geomid = mjSTACKALLOC(d, npixel, int);
+
+            // compute ray directions using helper (normalized)
+            for (int row = 0; row < height; row++) {
+              for (int col = 0; col < width; col++) {
+                int idx = row*width + col;
+                mjtNum origin[3];
+                mju_camPixelRay(origin, vec + 3*idx, cam_xpos, cam_xmat,
+                                col, row, fx, fy, cx, cy, projection, ortho_extent);
+              }
+            }
+
+            // cast all rays
+            mj_multiRay(m, d, cam_xpos, vec, NULL, 1, bodyexclude,
+                        geomid, d->sensordata + adr, npixel, mjMAXVAL);
+
+            mj_freeStack(d);
+          } else {
+            // orthographic: parallel rays, different origins
+            for (int row = 0; row < height; row++) {
+              for (int col = 0; col < width; col++) {
+                int idx = row*width + col;
+                mjtNum origin[3], direction[3];
+                mju_camPixelRay(origin, direction, cam_xpos, cam_xmat,
+                                col, row, fx, fy, cx, cy, projection, ortho_extent);
+
+                d->sensordata[adr + idx] = mj_ray(m, d, origin, direction, NULL, 1,
+                                                  bodyexclude, NULL);
+              }
+            }
+          }
+        }
 
         break;
 

src/engine/engine_support.c

@@ -743,3 +743,43 @@ int mju_condataSize(int dataspec) {
   }
   return size;
 }
+
+
+// compute camera pixel parameters from model, output are:
+//   pixel units: fx, fy (focal lengths), cx, cy (principal point)
+//   length units: extent
+void mju_camIntrinsics(const mjModel* m, int camid,
+                       mjtNum* fx, mjtNum* fy, mjtNum* cx, mjtNum* cy, mjtNum* extent) {
+  const int width = m->cam_resolution[2*camid];
+  const int height = m->cam_resolution[2*camid+1];
+  const float* sensorsize = m->cam_sensorsize + 2*camid;
+  const float* intrinsic = m->cam_intrinsic + 4*camid;
+  const mjtProjection projection = (mjtProjection)m->cam_projection[camid];
+
+  switch (projection) {
+  case mjPROJ_PERSPECTIVE:
+    if (sensorsize[0] && sensorsize[1]) {
+      // intrinsic-based perspective camera
+      *fx = intrinsic[0] / sensorsize[0] * width;
+      *fy = intrinsic[1] / sensorsize[1] * height;
+      *cx = intrinsic[2] / sensorsize[0] * width;
+      *cy = intrinsic[3] / sensorsize[1] * height;
+    } else {
+      // fovy-based perspective camera
+      *fx = *fy = 0.5 / mju_tan(m->cam_fovy[camid] * mjPI / 360.0) * height;
+      *cx = (mjtNum)width / 2.0;
+      *cy = (mjtNum)height / 2.0;
+    }
+    break;
+  case mjPROJ_ORTHOGRAPHIC:
+    // orthographic: normalize pixel offset to [-1, 1]
+    *fx = (mjtNum)width / 2.0;
+    *fy = (mjtNum)height / 2.0;
+    *cx = *fx;
+    *cy = *fy;
+    break;
+  }
+
+  // extent only used for orthographic cameras
+  *extent = m->cam_fovy[camid];
+}

src/engine/engine_support.h

@@ -120,6 +120,12 @@ MJAPI const char* mj_versionString(void);
 // return total size of data fields in a contact sensor bitfield specification
 MJAPI int mju_condataSize(int dataSpec);
 
+// compute camera pixel parameters from model
+// outputs: fx, fy (focal length in pixels), cx, cy (principal point), ortho_extent
+void mju_camIntrinsics(const mjModel* m, int camid,
+                       mjtNum* fx, mjtNum* fy, mjtNum* cx, mjtNum* cy,
+                       mjtNum* ortho_extent);
+
 #ifdef __cplusplus
 }
 #endif

src/engine/engine_util_misc.c

@@ -417,6 +417,8 @@ mjtNum mju_wrap(mjtNum wpnt[6], const mjtNum x0[3], const mjtNum x1[3],
 }
 
 
+//------------------------------ misc geometry -----------------------------------------------------
+
 // all 3 semi-axes of a geom
 void mju_geomSemiAxes(mjtNum semiaxes[3], const mjtNum size[3], mjtGeom type) {
   switch (type) {
@@ -494,7 +496,42 @@ int mju_insideGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3
 }
 
 
-// ----------------------------- Flex interpolation ------------------------------------------------
+// compute ray origin and direction for pixel (col, row) in camera image
+// for perspective: origin is unchanged, direction is computed
+// for orthographic: direction is -Z in camera frame, origin is offset from camera center
+void mju_camPixelRay(mjtNum origin[3], mjtNum direction[3],
+                     const mjtNum cam_xpos[3], const mjtNum cam_xmat[9],
+                     int col, int row, mjtNum fx, mjtNum fy, mjtNum cx, mjtNum cy,
+                     int projection, mjtNum ortho_extent) {
+  // pixel center (row 0 = top of image)
+  mjtNum px = col + 0.5 - cx;
+  mjtNum py = row + 0.5 - cy;
+
+  if (projection == mjPROJ_PERSPECTIVE) {
+    // origin is camera position
+    mju_copy3(origin, cam_xpos);
+
+    // direction in camera frame: (x/fx, -y/fy, -1), then normalized
+    mjtNum dir_cam[3] = {px / fx, -py / fy, -1.0};
+    mju_mulMatVec3(direction, cam_xmat, dir_cam);
+    mju_normalize3(direction);
+  } else {
+    // orthographic: parallel rays, direction is -Z in camera frame
+    direction[0] = -cam_xmat[2];
+    direction[1] = -cam_xmat[5];
+    direction[2] = -cam_xmat[8];
+
+    // origin offset in camera frame (ortho_extent is full height, use half for each side)
+    mjtNum half_extent = ortho_extent / 2;
+    mjtNum offset_cam[3] = {px / fx * half_extent, -py / fy * half_extent, 0};
+    mjtNum offset_world[3];
+    mju_mulMatVec3(offset_world, cam_xmat, offset_cam);
+    mju_add3(origin, cam_xpos, offset_world);
+  }
+}
+
+
+// ----------------------------- flex interpolation ------------------------------------------------
 
 mjtNum static inline phi(mjtNum s, int i, int order) {
   if (order == 1) {

src/engine/engine_util_misc.h

@@ -57,6 +57,13 @@ MJAPI void mju_geomSemiAxes(mjtNum semiaxes[3], const mjtNum size[3], mjtGeom ty
 int mju_insideGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3], mjtGeom type,
                    const mjtNum point[3]);
 
+// compute ray origin and direction for pixel (col, row) in camera image
+// directions are normalized so ray functions return actual 3D distance
+void mju_camPixelRay(mjtNum origin[3], mjtNum direction[3],
+                     const mjtNum cam_xpos[3], const mjtNum cam_xmat[9],
+                     int col, int row, mjtNum fx, mjtNum fy, mjtNum cx, mjtNum cy,
+                     int projection, mjtNum ortho_extent);
+
 // ----------------------------- Flex interpolation ------------------------------------------------
 
 // evaluate the deformation gradient at p using the nodal dof values

src/engine/engine_vis_visualize.c

@@ -2478,22 +2478,64 @@ static void addRangefinderGeoms(const mjModel* m, mjData* d, const mjvOption* vo
 
   for (int i=0; i < m->nsensor; i++) {
     if (m->sensor_type[i] == mjSENS_RANGEFINDER) {
-      // sensor data
-      mjtNum dst = d->sensordata[m->sensor_adr[i]];
-      int sid = m->sensor_objid[i];
+      int objid = m->sensor_objid[i];
+      int adr = m->sensor_adr[i];
 
-      // null output: nothing to render
-      if (dst < 0) {
-        continue;
+      // site-attached rangefinder
+      if (m->sensor_objtype[i] == mjOBJ_SITE) {
+        mjtNum dst = d->sensordata[adr];
+
+        // null output: nothing to render
+        if (dst < 0) {
+          continue;
+        }
+
+        // make ray
+        mjtNum* from = d->site_xpos+3*objid;
+        mjtNum to[3] = {from[0] + d->site_xmat[9*objid+2]*dst,
+                        from[1] + d->site_xmat[9*objid+5]*dst,
+                        from[2] + d->site_xmat[9*objid+8]*dst};
+        addConnector(scn, mjGEOM_LINE, 3, from, to, m->vis.rgba.rangefinder,
+                     i, mjCAT_DECOR, mjOBJ_SENSOR);
       }
 
-      // make ray
-      mjtNum* from = d->site_xpos+3*sid;
-      mjtNum to[3] = {from[0] + d->site_xmat[9*sid+2]*dst,
-                      from[1] + d->site_xmat[9*sid+5]*dst,
-                      from[2] + d->site_xmat[9*sid+8]*dst};
-      addConnector(scn, mjGEOM_LINE, 3, from, to, m->vis.rgba.rangefinder,
-                   i, mjCAT_DECOR, mjOBJ_SENSOR);
+      // camera-attached rangefinder
+      else if (m->sensor_objtype[i] == mjOBJ_CAMERA) {
+        const int width = m->cam_resolution[2*objid];
+        const int height = m->cam_resolution[2*objid+1];
+        const mjtNum* cam_xpos = d->cam_xpos + 3*objid;
+        const mjtNum* cam_xmat = d->cam_xmat + 9*objid;
+        const int projection = m->cam_projection[objid];
+
+        // compute focal length in pixels using helper
+        mjtNum fx, fy, cx, cy, ortho_extent;
+        mju_camIntrinsics(m, objid, &fx, &fy, &cx, &cy, &ortho_extent);
+
+        // draw ray for each pixel
+        for (int row = 0; row < height; row++) {
+          for (int col = 0; col < width; col++) {
+            int idx = row*width + col;
+            mjtNum dst = d->sensordata[adr + idx];
+
+            // null output: nothing to render
+            if (dst < 0) {
+              continue;
+            }
+
+            // compute ray origin and direction
+            mjtNum origin[3], direction[3];
+            mju_camPixelRay(origin, direction, cam_xpos, cam_xmat,
+                            col, row, fx, fy, cx, cy, projection, ortho_extent);
+
+            // compute endpoint
+            mjtNum to[3];
+            mju_addScl3(to, origin, direction, dst);
+
+            addConnector(scn, mjGEOM_LINE, 3, origin, to, m->vis.rgba.rangefinder,
+                         i, mjCAT_DECOR, mjOBJ_SENSOR);
+          }
+        }
+      }
     } else if (m->sensor_type[i] == mjSENS_GEOMFROMTO) {
       // sensor data
       mjtNum* fromto = d->sensordata + m->sensor_adr[i];