Refactor hfield-convex collision code. This should be a no-op.

PiperOrigin-RevId: 837045492
Change-Id: I82851615c7f69774da83cab361316e5a6e9982c6

src/engine/engine_collision_convex.c

@@ -750,6 +750,12 @@ void mjc_initCCDObj(mjCCDObj* obj, const mjModel* m, const mjData* d, int g, mjt
     case mjGEOM_HFIELD:
       obj->center = mjc_prism_center;
       obj->support = mjc_prism_support;
+
+      int hid = m->geom_dataid[g];
+      obj->hfield_nrow = m->hfield_nrow[hid];
+      obj->hfield_ncol = m->hfield_ncol[hid];
+      obj->size = m->hfield_size + 4*hid;
+      obj->hfield_data = m->hfield_data + m->hfield_adr[hid];
       break;
     default:
       obj->support = NULL;
@@ -1141,8 +1147,8 @@ static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
 }
 
 
-// add vertex to prism, count vertices
+// add vertex to prism
-static void addVert(int* nvert, mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) {
+static inline void addVert(mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) {
   // move old data
   mju_copy3(obj->prism[0], obj->prism[1]);
   mju_copy3(obj->prism[1], obj->prism[2]);
@@ -1153,125 +1159,140 @@ static void addVert(int* nvert, mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) {
   obj->prism[2][0] = obj->prism[5][0] = x;
   obj->prism[2][1] = obj->prism[5][1] = y;
   obj->prism[5][2] = z;
+}
 
-  // count
+
-  (*nvert)++;
+// add vertex to prism
+static inline void addPrismVert(mjCCDObj* obj, int r, int c, int i, mjtNum dx, mjtNum dy, mjtNum margin) {
+  // move old data
+  mju_copy3(obj->prism[0], obj->prism[1]);
+  mju_copy3(obj->prism[1], obj->prism[2]);
+  mju_copy3(obj->prism[3], obj->prism[4]);
+  mju_copy3(obj->prism[4], obj->prism[5]);
+
+  int dr = 1 - i;
+
+  // add new vertex at last position
+  obj->prism[2][0] = obj->prism[5][0] = dx*c - obj->size[0];
+  obj->prism[2][1] = obj->prism[5][1] = dy*(r + dr) - obj->size[1];
+  obj->prism[5][2] = obj->hfield_data[(r + dr)*obj->hfield_ncol + c]*obj->size[2];
+
+  // factor in margin
+  obj->prism[5][2] += margin;
 }
 
 
 // entry point for heightfield collisions
 int mjc_ConvexHField(const mjModel* m, const mjData* d,
                      mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO_HFIELD
+  // hfield frame
-  mjtNum mat[9], savemat2[9], savepos2[3], pos[3], vec[3], r2, dx, dy;
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
-  mjtNum xmin, xmax, ymin, ymax, zmin, zmax;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+
+  // geom2 frame
+  mjtNum* pos2 = d->geom_xpos + 3*g2;
+  mjtNum* mat2 = d->geom_xmat + 9*g2;
+
+  // hfield data
   int hid = m->geom_dataid[g1];
   int nrow = m->hfield_nrow[hid];
   int ncol = m->hfield_ncol[hid];
-  int dr[2], cnt, rmin, rmax, cmin, cmax;
+  mjtNum size0 = m->hfield_size[4*hid + 0], size1 = m->hfield_size[4*hid + 1];
-  const float* data = m->hfield_data + m->hfield_adr[hid];
+  mjtNum size2 = m->hfield_size[4*hid + 2], size3 = m->hfield_size[4*hid + 3];
 
-  // ccd-related
+  // try early return using box-sphere test
+
+  // express geom2 pos in hfield frame
+  mjtNum pos[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]};
+  mju_mulMatTVec3(pos, mat1, pos);
+
+  // sphere radius is geom2 rbound + margin
+  mjtNum radius = m->geom_rbound[g2] + margin;
+
+  // box-sphere test
+  if ((size0 < pos[0] - radius) || (-size0 > pos[0] + radius) ||
+      (size1 < pos[1] - radius) || (-size1 > pos[1] + radius) ||
+      (size2 < pos[2] - radius) || (-size3 > pos[2] + radius)) {
+    return 0;
+  }
+
+  // ccd set up
   mjCCDObj obj1, obj2;
   mjc_initCCDObj(&obj1, m, d, g1, 0);
   mjc_initCCDObj(&obj2, m, d, g2, 0);
 
-  ccd_vec3_t dirccd, vecccd;
-  ccd_real_t depth;
-  ccd_t ccd;
 
-  // point size1 to hfield size instead of geom1 size
+  // try early return using AABB box-box test
-  size1 = m->hfield_size + 4*hid;
 
-  //------------------------------------- frame alignment, box-sphere test
+  // express geom2 mat in hfield frame
-
+  mjtNum mat[9];
-  // express geom2 pos in heightfield frame
-  mju_sub3(vec, pos2, pos1);
-  mju_mulMatTVec(pos, mat1, vec, 3, 3);
-
-  // get geom2 rbound
-  r2 = m->geom_rbound[g2];
-
-  // box-sphere test: horizontal plane
-  for (int i=0; i < 2; i++) {
-    if ((size1[i] < pos[i]-r2-margin) || (-size1[i] > pos[i]+r2+margin)) {
-      return 0;
-    }
-  }
-
-  // box-sphere test in: vertical direction
-  if (size1[2] < pos[2]-r2-margin) {   // up
-    return 0;
-  }
-  if (-size1[3] > pos[2]+r2+margin) {  // down
-    return 0;
-  }
-
-  // express geom2 mat in heightfield frame
   mju_mulMatTMat3(mat, mat1, mat2);
 
-  //------------------------------------- AABB computation, box-box test
-
   // save mat2 and pos2, replace with relative frame
-  mju_copy(savemat2, mat2, 9);
+  mjtNum savemat2[9], savepos2[3];
+  mju_copy9(savemat2, mat2);
   mju_copy3(savepos2, pos2);
-  mju_copy(mat2, mat, 9);
+  mju_copy9(mat2, mat);
   mju_copy3(pos2, pos);
 
+  mjtNum dir[3] = {0, 0, 0}, res[3];
+
   // get support point in +X
-  ccdVec3Set(&dirccd, 1, 0, 0);
+  dir[0] = 1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  xmax = vecccd.v[0];
+  mjtNum xmax = res[0];
 
   // get support point in -X
-  ccdVec3Set(&dirccd, -1, 0, 0);
+  dir[0] = -1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  xmin = vecccd.v[0];
+  mjtNum xmin = res[0];
+
+  dir[0] = 0;
 
   // get support point in +Y
-  ccdVec3Set(&dirccd, 0, 1, 0);
+  dir[1] = 1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  ymax = vecccd.v[1];
+  mjtNum ymax = res[1];
 
   // get support point in -Y
-  ccdVec3Set(&dirccd, 0, -1, 0);
+  dir[1] = -1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  ymin = vecccd.v[1];
+  mjtNum ymin = res[1];
+
+  dir[1] = 0;
 
   // get support point in +Z
-  ccdVec3Set(&dirccd, 0, 0, 1);
+  dir[2] = 1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  zmax = vecccd.v[2];
+  mjtNum zmax = res[2];
 
   // get support point in -Z
-  ccdVec3Set(&dirccd, 0, 0, -1);
+  dir[2] = -1;
-  mjccd_support(&obj2, &dirccd, &vecccd);
+  obj2.support(res, &obj2, dir);
-  zmin = vecccd.v[2];
+  mjtNum zmin = res[2];
 
-  // box-box test
+  // AABB box-box test
-  if ((xmin-margin > size1[0]) || (xmax+margin < -size1[0]) ||
+  if ((xmin - margin > size0) || (xmax + margin < -size0) ||
-      (ymin-margin > size1[1]) || (ymax+margin < -size1[1]) ||
+      (ymin - margin > size1) || (ymax + margin < -size1) ||
-      (zmin-margin > size1[2]) || (zmax+margin < -size1[3])) {
+      (zmin - margin > size2) || (zmax + margin < -size3)) {
-    // restore mat2 and pos2
+    mju_copy9(mat2, savemat2);
-    mju_copy(mat2, savemat2, 9);
     mju_copy3(pos2, savepos2);
-
     return 0;
   }
 
   // compute sub-grid bounds
-  cmin = (int) mju_floor((xmin + size1[0]) / (2*size1[0]) * (ncol-1));
+  int cmin = (int) mju_floor((xmin + size0) / (2.0*size0) * (ncol-1));
-  cmax = (int) mju_ceil ((xmax + size1[0]) / (2*size1[0]) * (ncol-1));
+  int cmax = (int) mju_ceil ((xmax + size0) / (2.0*size0) * (ncol-1));
-  rmin = (int) mju_floor((ymin + size1[1]) / (2*size1[1]) * (nrow-1));
+  int rmin = (int) mju_floor((ymin + size1) / (2.0*size1) * (nrow-1));
-  rmax = (int) mju_ceil ((ymax + size1[1]) / (2*size1[1]) * (nrow-1));
+  int rmax = (int) mju_ceil ((ymax + size1) / (2.0*size1) * (nrow-1));
   cmin = mjMAX(0, cmin);
   cmax = mjMIN(ncol-1, cmax);
   rmin = mjMAX(0, rmin);
   rmax = mjMIN(nrow-1, rmax);
 
-  //------------------------------------- collision testing
+  // CCD collision testing
 
-  // init ccd structure
+  ccd_t ccd;
   mjc_initCCD(&ccd, m);
   ccd.first_dir = prism_firstdir;
   ccd.center1 = mjccd_prism_center;
@@ -1282,51 +1303,46 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   // geom margin needed for actual collision test
   obj2.margin = margin;
 
-  // compute real-valued grid step, and triangulation direction
+  // compute real-valued grid step
-  dx = (2.0*size1[0]) / (ncol-1);
+  mjtNum dx = (2.0*size0) / (ncol-1);
-  dy = (2.0*size1[1]) / (nrow-1);
+  mjtNum dy = (2.0*size1) / (nrow-1);
-  dr[0] = 1;
-  dr[1] = 0;
 
   // set zbottom value using base size
-  obj1.prism[0][2] = obj1.prism[1][2] = obj1.prism[2][2] = -size1[3];
+  obj1.prism[0][2] = obj1.prism[1][2] = obj1.prism[2][2] = -size3;
 
-  // process all prisms in sub-grid
+  // process all prisms in subgrid
-  cnt = 0;
+  int ncon = 0;
   for (int r=rmin; r < rmax; r++) {
-    int nvert = 0;
+    addPrismVert(&obj1, r, cmin, 0, dx, dy, margin);
-    for (int c=cmin; c <= cmax; c++) {
+    addPrismVert(&obj1, r, cmin, 1, dx, dy, margin);
+    for (int c=cmin + 1; c <= cmax; c++) {
       for (int i=0; i < 2; i++) {
         // send vertex to prism constructor
-        addVert(&nvert, &obj1, dx*c-size1[0], dy*(r+dr[i])-size1[1],
+        addPrismVert(&obj1, r, c, i, dx, dy, margin);
-                data[(r+dr[i])*ncol+c]*size1[2]+margin);
 
-        // check for enough vertices
+        // prism height test
-        if (nvert > 2) {
+        if (obj1.prism[3][2] < zmin && obj1.prism[4][2] < zmin && obj1.prism[5][2] < zmin) {
-          // prism height test
+          continue;
-          if (obj1.prism[3][2] < zmin && obj1.prism[4][2] < zmin
+        }
-              && obj1.prism[5][2] < zmin) {
-            continue;
-          }
 
-          // run penetration function, save contact
+        // run penetration function, save contact
-          if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0
+        ccd_vec3_t dirccd, vecccd;
-              && !ccdVec3Eq(&dirccd, ccd_vec3_origin)) {
+        ccd_real_t depth;
-            // fill in contact data, transform to global coordinates
+        if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0
-            con[cnt].dist = -depth;
+            && !ccdVec3Eq(&dirccd, ccd_vec3_origin)) {
-            mju_mulMatVec3(con[cnt].frame, mat1, dirccd.v);
+          // fill in contact data, transform to global coordinates
-            mju_mulMatVec3(con[cnt].pos, mat1, vecccd.v);
+          con[ncon].dist = -depth;
-            mju_addTo3(con[cnt].pos, pos1);
+          mju_mulMatVec3(con[ncon].frame, mat1, dirccd.v);
-            mju_zero3(con[cnt].frame+3);
+          mju_mulMatVec3(con[ncon].pos, mat1, vecccd.v);
+          mju_addTo3(con[ncon].pos, pos1);
+          mju_zero3(con[ncon].frame+3);
 
-            // count, stop if max number reached
+          // force out of all loops if max contacts reached
-            cnt++;
+          if (++ncon >= mjMAXCONPAIR) {
-            if (cnt >= mjMAXCONPAIR) {
+            r = rmax+1;
-              r = rmax+1;
+            c = cmax+1;
-              c = cmax+1;
+            i = 3;
-              i = 3;
+            break;
-              break;
-            }
           }
         }
       }
@@ -1334,17 +1350,17 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   }
 
   // restore mat2 and pos2
-  mju_copy(mat2, savemat2, 9);
+  mju_copy9(mat2, savemat2);
   mju_copy3(pos2, savepos2);
 
   if (mjDISABLED(mjDSBL_NATIVECCD)) {
     // fix contact normals
-    for (int i=0; i < cnt; i++) {
+    for (int i=0; i < ncon; i++) {
       mjc_fixNormal(m, d, con+i, g1, g2);
     }
   }
 
-  return cnt;
+  return ncon;
 }
 
 
@@ -1737,11 +1753,11 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
     for (int c=cmin; c <= cmax; c++) {
       for (int k=0; k < 2; k++) {
         // send vertex to prism constructor
-        addVert(&nvert, &obj1, dx*c-hsize[0], dy*(r+dr[k])-hsize[1],
+        addVert(&obj1, dx*c-hsize[0], dy*(r+dr[k])-hsize[1],
                 hdata[(r+dr[k])*ncol+c]*hsize[2]+margin);
 
         // check for enough vertices
-        if (nvert > 2) {
+        if (++nvert > 2) {
           // prism height test
           if (obj1.prism[3][2] < zmin && obj1.prism[4][2] < zmin && obj1.prism[5][2] < zmin) {
             continue;

src/engine/engine_collision_convex.h

@@ -28,14 +28,6 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 
-#define mjGETINFO_HFIELD \
-    const mjtNum* pos1  = d->geom_xpos + 3*g1; \
-    const mjtNum* mat1  = d->geom_xmat + 9*g1; \
-    const mjtNum* size1 = m->geom_size + 3*g1; \
-          mjtNum* pos2  = d->geom_xpos + 3*g2; \
-          mjtNum* mat2  = d->geom_xmat + 9*g2;
-// mjc_ConvexHField modifies and then restores pos2 and mat2
-
 // minimum number of vertices to use hill-climbing in mesh support
 #define mjMESH_HILLCLIMB_MIN 10
 
@@ -58,7 +50,13 @@ struct _mjCCDObj {
   mjtNum rotate[4];
   void (*center)(mjtNum res[3], const struct _mjCCDObj* obj);
   void (*support)(mjtNum res[3], struct _mjCCDObj* obj, const mjtNum dir[3]);
-  mjtNum prism[6][3];  // for hfield
+
+  // for hfield
+  mjtNum prism[6][3];
+  const mjtNum* size;
+  const float* hfield_data;
+  int hfield_nrow;
+  int hfield_ncol;
 };
 typedef struct _mjCCDObj mjCCDObj;
 

src/engine/engine_core_smooth.c

@@ -377,7 +377,7 @@ void mj_camlight(const mjModel* m, mjData* d) {
     case mjCAMLIGHT_TRACK:
     case mjCAMLIGHT_TRACKCOM:
       // fixed global orientation
-      mju_copy(d->cam_xmat+9*i, m->cam_mat0+9*i, 9);
+      mju_copy9(d->cam_xmat+9*i, m->cam_mat0+9*i);
 
       // position: track camera body
       if (m->cam_mode[i] == mjCAMLIGHT_TRACK) {

src/engine/engine_core_util.c

@@ -686,7 +686,7 @@ void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
 
     // orientation of the COM (inertial) frame of b-th body
     mjtNum ximat[9];
-    mju_copy(ximat, d->ximat+9*b, 9);
+    mju_copy9(ximat, d->ximat+9*b);
 
     // save the inertia matrix of b-th body
     mjtNum inertia[9] = {0};
@@ -885,11 +885,11 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
       break;
     case mjSAMEFRAME_BODY:
     case mjSAMEFRAME_BODYROT:
-      mju_copy(xmat, d->xmat+9*body, 9);
+      mju_copy9(xmat, d->xmat+9*body);
       break;
     case mjSAMEFRAME_INERTIA:
     case mjSAMEFRAME_INERTIAROT:
-      mju_copy(xmat, d->ximat+9*body, 9);
+      mju_copy9(xmat, d->ximat+9*body);
       break;
     }
   }

src/engine/engine_derivative.c

@@ -247,7 +247,7 @@ void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db
   mju_addToScl(Da_tmp, KK, coef, 9);
 
   if (Da) {
-    mju_copy(Da, Da_tmp, 9);
+    mju_copy9(Da, Da_tmp);
   }
 
   if (Db) {  // Db = -Da^T
@@ -309,7 +309,7 @@ void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
     if (Dquat)          Dquat[i] = a*eye[i] + b*cross[i] + c*outer[i];
     if (Dvel || Dscale) Dvel_[i] = b*eye[i] + c*cross[i] + d*outer[i];
   }
-  if (Dvel) mju_copy(Dvel, Dvel_, 9);
+  if (Dvel) mju_copy9(Dvel, Dvel_);
   if (Dscale) mju_mulMatVec3(Dscale, Dvel_, vel);
 }
 

src/engine/engine_sensor.c

@@ -1064,7 +1064,7 @@ void mj_sensorAcc(const mjModel* m, mjData* d) {
               int conid = match[j].id;
               mj_contactForce(m, d, conid, wrench + 6*j);
               mju_copy3(pos + 3*j, d->contact[conid].pos);
-              mju_copy(frame + 9*j, d->contact[conid].frame, 9);
+              mju_copy9(frame + 9*j, d->contact[conid].frame);
               if (match[j].flip) {
                 mju_scl(wrench + 6*j , wrench + 6*j, -1, 6);
               }

src/engine/engine_setconst.c

@@ -578,7 +578,7 @@ static void set0(mjModel* m, mjData* d) {
     mju_sub3(m->cam_poscom0+3*i, d->cam_xpos+3*i, d->subtree_com+ (id1 >= 0 ? 3*id1 : 3*id));
 
     // copy mat
-    mju_copy(m->cam_mat0+9*i, d->cam_xmat+9*i, 9);
+    mju_copy9(m->cam_mat0+9*i, d->cam_xmat+9*i);
   }
 
   // light compos0, pos0, dir0

src/engine/engine_util_blas.c

@@ -46,10 +46,17 @@ int mju_equal3(const mjtNum vec1[3], const mjtNum vec2[3]) {
 
 
 // res = vec
-void mju_copy3(mjtNum res[3], const mjtNum data[3]) {
+void mju_copy3(mjtNum res[3], const mjtNum vec[3]) {
-  res[0] = data[0];
+  res[0] = vec[0];
-  res[1] = data[1];
+  res[1] = vec[1];
-  res[2] = data[2];
+  res[2] = vec[2];
+}
+
+// res = mat
+void mju_copy9(mjtNum res[9], const mjtNum mat[9]) {
+  res[0] = mat[0];  res[1] = mat[1];  res[2] = mat[2];
+  res[3] = mat[3];  res[4] = mat[4];  res[5] = mat[5];
+  res[6] = mat[6];  res[7] = mat[7];  res[8] = mat[8];
 }
 
 

src/engine/engine_util_blas.h

@@ -71,7 +71,10 @@ MJAPI void mju_zero3(mjtNum res[3]);
 MJAPI int mju_equal3(const mjtNum vec1[3], const mjtNum vec2[3]);
 
 // res = vec
-MJAPI void mju_copy3(mjtNum res[3], const mjtNum data[3]);
+MJAPI void mju_copy3(mjtNum res[3], const mjtNum vec[3]);
+
+// res = mat
+void mju_copy9(mjtNum res[9], const mjtNum mat[9]);
 
 // res = vec*scl
 MJAPI void mju_scl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl);