Add automatic free-joint alignment.

PiperOrigin-RevId: 671010083
Change-Id: Ib7d60f75c7f545b1886703715663db8de77feeef

src/user/user_init.c

@@ -104,6 +104,7 @@ void mjs_defaultJoint(mjsJoint* joint) {
   joint->axis[2] = 1;
   joint->limited = mjLIMITED_AUTO;
   joint->actfrclimited = mjLIMITED_AUTO;
+  joint->align = mjALIGNFREE_AUTO;
   mj_defaultSolRefImp(joint->solref_limit, joint->solimp_limit);
   mj_defaultSolRefImp(joint->solref_friction, joint->solimp_friction);
 }

src/user/user_objects.cc

@@ -1579,6 +1579,33 @@ void mjCBody::Compile(void) {
     margin = std::max(margin, geoms[i]->margin);
   }
 
+  // check conditions for free-joint alignment
+  bool align_free = (joints.size() == 1                  &&  // only one joint AND
+                     joints[0]->spec.type == mjJNT_FREE  &&  // it is a free joint AND
+                     bodies.empty()                      &&  // no child bodies AND
+                     (joints[0]->spec.align == 1 ||          // either joint.align="true"
+                      (joints[0]->spec.align == 2        &&  // or joint.align="auto"
+                       model->alignfree)));                  //  and compiler.align="true"
+
+  // free-joint alignment, phase 1 (this body + child geoms)
+  double ipos_inverse[3], iquat_inverse[4];
+  if (align_free) {
+    // accumulate iframe transformation to body frame
+    mjuu_frameaccum(pos, quat, ipos, iquat);
+
+    // compute inverse iframe transformation
+    mjuu_frameinvert(ipos_inverse, iquat_inverse, ipos, iquat);
+
+    // save iframe, set it to null
+    mjuu_setvec(ipos, 0, 0, 0);
+    mjuu_setvec(iquat, 1, 0, 0, 0);
+
+    // apply inverse iframe transformation to all child geoms
+    for (int i=0; i<geoms.size(); i++) {
+      mjuu_frameaccumChild(ipos_inverse, iquat_inverse, geoms[i]->pos, geoms[i]->quat);
+    }
+  }
+
   // compute bounding volume hierarchy
   ComputeBVH();
 
@@ -1652,6 +1679,28 @@ void mjCBody::Compile(void) {
       }
     }
   }
+
+  // free joint alignment, phase 2 (transform sites, cameras and lights)
+  if (align_free) {
+    // frames have already been compiled and applied to children
+
+    // sites
+    for (int i=0; i<sites.size(); i++) {
+      mjuu_frameaccumChild(ipos_inverse, iquat_inverse, sites[i]->pos, sites[i]->quat);
+    }
+
+    // cameras
+    for (int i=0; i<cameras.size(); i++) {
+      mjuu_frameaccumChild(ipos_inverse, iquat_inverse, cameras[i]->pos, cameras[i]->quat);
+    }
+
+    // lights
+    for (int i=0; i<lights.size(); i++) {
+      double qunit[4]= {1, 0, 0, 0};
+      mjuu_frameaccumChild(ipos_inverse, iquat_inverse, lights[i]->pos, qunit);
+      mjuu_rotVecQuat(lights[i]->dir, lights[i]->dir, iquat_inverse);
+    }
+  }
 }
 
 

src/xml/xml_native_reader.cc

@@ -97,11 +97,11 @@ static void UpdateString(string& psuffix, int count, int i) {
 const char* MJCF[nMJCF][mjXATTRNUM] = {
 {"mujoco", "!", "1", "model"},
 {"<"},
-    {"compiler", "*", "20", "autolimits", "boundmass", "boundinertia", "settotalmass",
+    {"compiler", "*", "21", "autolimits", "boundmass", "boundinertia", "settotalmass",
         "balanceinertia", "strippath", "coordinate", "angle", "fitaabb", "eulerseq",
         "meshdir", "texturedir", "discardvisual", "convexhull", "usethread",
         "fusestatic", "inertiafromgeom", "inertiagrouprange", "exactmeshinertia",
-        "assetdir"},
+        "assetdir", "alignfree"},
     {"<"},
         {"lengthrange", "?", "10", "mode", "useexisting", "uselimit",
             "accel", "maxforce", "timeconst", "timestep",
@@ -271,7 +271,7 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
             "solreflimit", "solimplimit", "solreffriction", "solimpfriction",
             "stiffness", "range", "actuatorfrcrange", "actuatorgravcomp", "margin", "ref",
             "springref", "armature", "damping", "frictionloss", "user"},
-        {"freejoint", "*", "2", "name", "group"},
+        {"freejoint", "*", "3", "name", "group", "align"},
         {"geom", "*", "33", "name", "class", "type", "contype", "conaffinity", "condim",
             "group", "priority", "size", "material", "friction", "mass", "density",
             "shellinertia", "solmix", "solref", "solimp",
@@ -1021,6 +1021,9 @@ void mjXReader::Compiler(XMLElement* section, mjSpec* spec) {
   if (MapValue(section, "exactmeshinertia", &n, bool_map, 2)){
     spec->exactmeshinertia = (n==1);
   }
+  if (MapValue(section, "alignfree", &n, bool_map, 2)) {
+    spec->alignfree = (n==1);
+  }
 
   // lengthrange subelement
   XMLElement* elem = FindSubElem(section, "lengthrange");
@@ -3423,6 +3426,7 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
         mjs_setString(joint->name, name.c_str());
       }
       ReadAttrInt(elem, "group", &joint->group);
+      MapValue(elem, "align", &joint->align, TFAuto_map, 3);
     }
 
     // geom sub-element

src/xml/xml_native_writer.cc

@@ -924,6 +924,9 @@ void mjXWriter::Compiler(XMLElement* root) {
   if (model->boundinertia) {
     WriteAttr(section, "boundinertia", 1, &model->boundinertia);
   }
+  if (model->alignfree) {
+    WriteAttrTxt(section, "alignfree", "true");
+  }
   if (!model->autolimits) {
     WriteAttrTxt(section, "autolimits", "false");
   }