Initial open sourcing of MuJoCo.

PiperOrigin-RevId: 450374687
Change-Id: Ie3225a46ce095fc28ae8e63c326a640261f562bb

src/engine/CMakeLists.txt

@@ -0,0 +1,74 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set(MUJOCO_ENGINE_SRCS
+    engine_array_safety.h
+    engine_callback.c
+    engine_callback.h
+    engine_collision_box.c
+    engine_collision_convex.c
+    engine_collision_convex.h
+    engine_collision_driver.c
+    engine_collision_driver.h
+    engine_collision_primitive.c
+    engine_collision_primitive.h
+    engine_core_constraint.c
+    engine_core_constraint.h
+    engine_core_smooth.c
+    engine_core_smooth.h
+    engine_crossplatform.h
+    engine_file.c
+    engine_file.h
+    engine_forward.c
+    engine_forward.h
+    engine_inverse.c
+    engine_inverse.h
+    engine_io.c
+    engine_io.h
+    engine_macro.h
+    engine_print.c
+    engine_print.h
+    engine_ray.c
+    engine_ray.h
+    engine_sensor.c
+    engine_sensor.h
+    engine_setconst.c
+    engine_setconst.h
+    engine_solver.c
+    engine_solver.h
+    engine_support.c
+    engine_support.h
+    engine_util_blas.c
+    engine_util_blas.h
+    engine_util_errmem.c
+    engine_util_errmem.h
+    engine_util_misc.c
+    engine_util_misc.h
+    engine_util_solve.c
+    engine_util_solve.h
+    engine_util_sparse.c
+    engine_util_sparse.h
+    engine_util_spatial.c
+    engine_util_spatial.h
+    engine_vfs.c
+    engine_vfs.h
+    engine_vis_init.c
+    engine_vis_init.h
+    engine_vis_interact.c
+    engine_vis_interact.h
+    engine_vis_visualize.c
+    engine_vis_visualize.h
+)
+
+target_sources(mujoco PRIVATE ${MUJOCO_ENGINE_SRCS})

src/engine/engine_array_safety.h

@@ -0,0 +1,31 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_ARRAY_SAFETY_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_ARRAY_SAFETY_H_
+
+#include <stdio.h>
+#include <string.h>
+
+// Evaluates to sizeof(arr) if arr is a char array, and emits a compiler error
+// otherwise. In particular, emits a compiler error if arr is a char*.
+#define mjSIZEOFARRAY(arr) _Generic(&(arr), char(*)[sizeof(arr)]: sizeof(arr))
+
+#define mjSNPRINTF(dest, ...) snprintf(dest, mjSIZEOFARRAY(dest), __VA_ARGS__)
+
+#define mjSTRNCAT(dest, src) strncat(dest, src, mjSIZEOFARRAY(dest) - strlen(dest) - 1)
+
+#define mjSTRNCPY(dest, src) mju_strncpy(dest, src, mjSIZEOFARRAY(dest))
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_ARRAY_SAFETY_H_

src/engine/engine_callback.c

@@ -0,0 +1,42 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_callback.h"
+
+#include <mujoco/mjdata.h>
+
+//------------------------- global callback pointers -----------------------------------------------
+
+mjfGeneric mjcb_passive  = 0;
+mjfGeneric mjcb_control  = 0;
+mjfConFilt mjcb_contactfilter = 0;
+mjfSensor  mjcb_sensor   = 0;
+mjfTime    mjcb_time     = 0;
+mjfAct     mjcb_act_bias = 0;
+mjfAct     mjcb_act_gain = 0;
+mjfAct     mjcb_act_dyn  = 0;
+
+
+
+// reset callbacks to defauls
+void mj_resetCallbacks(void) {
+  mjcb_passive  = 0;
+  mjcb_control  = 0;
+  mjcb_contactfilter = 0;
+  mjcb_sensor   = 0;
+  mjcb_time     = 0;
+  mjcb_act_bias = 0;
+  mjcb_act_gain = 0;
+  mjcb_act_dyn  = 0;
+}

src/engine/engine_callback.h

@@ -0,0 +1,42 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_CALLBACK_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_CALLBACK_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// global callback function pointers
+MJAPI extern mjfGeneric  mjcb_passive;
+MJAPI extern mjfGeneric  mjcb_control;
+MJAPI extern mjfConFilt  mjcb_contactfilter;
+MJAPI extern mjfSensor   mjcb_sensor;
+MJAPI extern mjfTime     mjcb_time;
+MJAPI extern mjfAct      mjcb_act_bias;
+MJAPI extern mjfAct      mjcb_act_gain;
+MJAPI extern mjfAct      mjcb_act_dyn;
+
+
+// reset callbacks to defaults
+MJAPI void mj_resetCallbacks(void);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_CALLBACK_H_

src/engine/engine_collision_convex.h

@@ -0,0 +1,64 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_COLLISION_CONVEX_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_COLLISION_CONVEX_H_
+
+// libCCD has an unconditional `#define _CRT_SECURE_NO_WARNINGS` on Windows.
+// TODO(stunya): Remove once https://github.com/danfis/libccd/pull/77 is merged
+#ifdef _CRT_SECURE_NO_WARNINGS
+#undef _CRT_SECURE_NO_WARNINGS
+#endif
+
+#include <ccd/vec3.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// ccd general object type
+struct _mjtCCD {
+  const mjModel* model;
+  const mjData* data;
+  int geom;
+  int meshindex;
+  mjtNum margin;
+  mjtNum rotate[4];
+};
+typedef struct _mjtCCD mjtCCD;
+
+
+// ccd support function
+void mjccd_support(const void *obj, const ccd_vec3_t *dir, ccd_vec3_t *vec);
+
+
+// pairwise collision functions using ccd
+int mjc_PlaneConvex   (const mjModel* m, const mjData* d,
+                       mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_ConvexHField  (const mjModel* m, const mjData* d,
+                       mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_Convex        (const mjModel* m, const mjData* d,
+                       mjContact* con, int g1, int g2, mjtNum margin);
+
+
+// fix contact frame normal
+void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, int g2);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_COLLISION_CONVEX_H_

src/engine/engine_collision_driver.c

@@ -0,0 +1,784 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_collision_driver.h"
+
+#include <stddef.h>
+#include <string.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_callback.h"
+#include "engine/engine_collision_convex.h"
+#include "engine/engine_collision_primitive.h"
+#include "engine/engine_core_constraint.h"
+#include "engine/engine_crossplatform.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_solve.h"
+#include "engine/engine_util_spatial.h"
+
+// table of pair-wise collision functions
+mjfCollision mjCOLLISIONFUNC[mjNGEOMTYPES][mjNGEOMTYPES] = {
+
+  /*               PLANE  HFIELD  SPHERE              CAPSULE             ELLIPSOID           CYLINDER            BOX                 MESH      */
+  /*PLANE     */  {0,     0,      mjc_PlaneSphere,    mjc_PlaneCapsule,   mjc_PlaneConvex,    mjc_PlaneCylinder,  mjc_PlaneBox,       mjc_PlaneConvex},
+  /*HFIELD    */  {0,     0,      mjc_ConvexHField,   mjc_ConvexHField,   mjc_ConvexHField,   mjc_ConvexHField,   mjc_ConvexHField,   mjc_ConvexHField},
+  /*SHPERE    */  {0,     0,      mjc_SphereSphere,   mjc_SphereCapsule,  mjc_Convex,         mjc_Convex,         mjc_SphereBox,      mjc_Convex},
+  /*CAPSULE   */  {0,     0,      0,                  mjc_CapsuleCapsule, mjc_Convex,         mjc_Convex,         mjc_CapsuleBox,     mjc_Convex},
+  /*ELLIPSOID */  {0,     0,      0,                  0,                  mjc_Convex,         mjc_Convex,         mjc_Convex,         mjc_Convex},
+  /*CYLINDER  */  {0,     0,      0,                  0,                  0,                  mjc_Convex,         mjc_Convex,         mjc_Convex},
+  /*BOX       */  {0,     0,      0,                  0,                  0,                  0,                  mjc_BoxBox,         mjc_Convex},
+  /*MESH      */  {0,     0,      0,                  0,                  0,                  0,                  0,                  mjc_Convex}
+
+};
+
+
+
+//----------------------------- collision detection entry point ------------------------------------
+
+void mj_collision(const mjModel* m, mjData* d) {
+  int g1, g2, signature, merged, b1 = 0, b2 = 0, exadr = 0, pairadr = 0, startadr;
+  int nexclude = m->nexclude, npair = m->npair, nbodypair = ((m->nbody-1)*m->nbody)/2;
+  int *broadphasepair = 0;
+  mjMARKSTACK;
+
+  // clear size
+  d->ncon = 0;
+
+  // return if disabled
+  if (mjDISABLED(mjDSBL_CONSTRAINT) || mjDISABLED(mjDSBL_CONTACT)
+      || m->nconmax==0 || m->nbody < 2) {
+    return;
+  }
+
+  // predefined only; ignore exclude
+  if (m->opt.collision==mjCOL_PAIR) {
+    for (pairadr=0; pairadr<npair; pairadr++) {
+      mj_collideGeoms(m, d, pairadr, -1, 0, 0);
+    }
+  }
+
+  // dynamic only or merge; apply exclude
+  else {
+    // call broadphase collision detector
+    broadphasepair = (int*)mj_stackAlloc(d, (m->nbody*(m->nbody-1))/2);
+    nbodypair = mj_broadphase(m, d, broadphasepair, (m->nbody*(m->nbody-1))/2);
+
+    // loop over body pairs (broadphase or all)
+    for (int i=0; i<nbodypair; i++) {
+      // reconstruct body pair ids
+      b1 = (broadphasepair[i]>>16) & 0xFFFF;
+      b2 = broadphasepair[i] & 0xFFFF;
+
+      // compute signature for this body pair
+      signature = ((b1+1)<<16) + (b2+1);
+
+      // merge predefined pairs
+      merged = 0;
+      startadr = pairadr;
+      if (npair && m->opt.collision==mjCOL_ALL) {
+        // test all predefined pairs for which pair_signature<=signature
+        while (pairadr<npair && m->pair_signature[pairadr]<=signature) {
+          if (m->pair_signature[pairadr]==signature) {
+            merged = 1;
+          }
+          mj_collideGeoms(m, d, pairadr++, -1, 0, 0);
+        }
+      }
+
+      // handle exclusion
+      if (nexclude) {
+        // advance exadr while exclude_signature < signature
+        while (m->exclude_signature[exadr]<signature && exadr<nexclude) {
+          exadr++;
+        }
+
+        // skip this body pair if its signature is found in exclude array
+        if (exadr<nexclude && m->exclude_signature[exadr]==signature) {
+          continue;
+        }
+      }
+
+      // test all geom pairs within this body pair
+      if (m->body_geomnum[b1] && m->body_geomnum[b2]) {
+        for (g1=m->body_geomadr[b1]; g1<m->body_geomadr[b1]+m->body_geomnum[b1]; g1++) {
+          for (g2=m->body_geomadr[b2]; g2<m->body_geomadr[b2]+m->body_geomnum[b2]; g2++) {
+            // merged: make sure geom pair is not repeated
+            if (merged) {
+              // find matching pair
+              int found = 0;
+              for (int k=startadr; k<pairadr; k++) {
+                if ((m->pair_geom1[k]==g1 && m->pair_geom2[k]==g2) ||
+                    (m->pair_geom1[k]==g2 && m->pair_geom2[k]==g1)) {
+                  found = 1;
+                  break;
+                }
+              }
+
+              // not found: test
+              if (!found) {
+                mj_collideGeoms(m, d, g1, g2, 0, 0);
+              }
+            }
+
+            // not merged: always test
+            else {
+              mj_collideGeoms(m, d, g1, g2, 0, 0);
+            }
+          }
+        }
+      }
+    }
+
+    // finish merging predefined pairs
+    if (npair && m->opt.collision==mjCOL_ALL)
+      while (pairadr<npair) {
+        mj_collideGeoms(m, d, pairadr++, -1, 0, 0);
+      }
+  }
+
+  mjFREESTACK;
+}
+
+
+
+//----------------------------- broad-phase collision detection ------------------------------------
+
+// helper structure for SAP sorting
+struct _mjtBroadphase {
+  float value;
+  int body_ismax;
+};
+typedef struct _mjtBroadphase mjtBroadphase;
+
+
+// make AABB for one body
+static void makeAABB(const mjModel* m, mjData* d, mjtNum* aabb, int body, const mjtNum* frame) {
+  int geom;
+  mjtNum _aabb[6], cen;
+
+  // no geoms attached to body: set to 0
+  if (m->body_geomnum[body]==0) {
+    mju_zero(aabb, 6);
+    return;
+  }
+
+  // process all body geoms
+  for (int i=0; i<m->body_geomnum[body]; i++) {
+    // get geom id
+    geom = m->body_geomadr[body]+i;
+
+    // set _aabb for this geom
+    for (int j=0; j<3; j++) {
+      cen = mju_dot3(d->geom_xpos+3*geom, frame+3*j);
+      _aabb[2*j] = cen - m->geom_rbound[geom] - m->geom_margin[geom];
+      _aabb[2*j+1] = cen + m->geom_rbound[geom] + m->geom_margin[geom];
+    }
+
+    // update body aabb
+    if (i==0) {
+      mju_copy(aabb, _aabb, 6);
+    } else {
+      for (int j=0; j<3; j++) {
+        aabb[2*j] = mju_min(aabb[2*j], _aabb[2*j]);
+        aabb[2*j+1] = mju_max(aabb[2*j+1], _aabb[2*j+1]);
+      }
+    }
+  }
+}
+
+
+
+// return 1 if body has plane or hfield geom, 0 otherwise
+static int has_plane_or_hfield(const mjModel* m, int body) {
+  int start = m->body_geomadr[body];
+  int end = m->body_geomadr[body] + m->body_geomnum[body];
+
+  // scan geoms belonging to body
+  int g;
+  for (g=start; g<end; g++) {
+    if (m->geom_type[g]==mjGEOM_PLANE || m->geom_type[g]==mjGEOM_HFIELD) {
+      return 1;
+    }
+  }
+
+  return 0;
+}
+
+
+
+// add body pair in buffer
+static void add_pair(const mjModel* m, int b1, int b2, int* npair, int* pair, int maxpair) {
+  // add pair if there is room in buffer
+  if ((*npair)<maxpair) {
+    // exlude based on contype and conaffinity
+    if (m && m->body_geomnum[b1]==1 && m->body_geomnum[b2]==1) {
+      // get contypes and conaffinities
+      int contype1 = m->geom_contype[m->body_geomadr[b1]];
+      int conaffinity1 = m->geom_conaffinity[m->body_geomadr[b1]];
+      int contype2 = m->geom_contype[m->body_geomadr[b2]];
+      int conaffinity2 = m->geom_conaffinity[m->body_geomadr[b2]];
+
+      // compatibility check
+      if (!(contype1 & conaffinity2) && !(contype2 & conaffinity1)) {
+        return;
+      }
+    }
+
+    // add pair
+    if (b1<b2) {
+      pair[*npair] = (b1<<16) + b2;
+    } else {
+      pair[*npair] = (b2<<16) + b1;
+    }
+
+    (*npair)++;
+  } else {
+    mju_error("Broadphase buffer full");
+  }
+}
+
+
+
+// comparison function for broadphase
+quicksortfunc(broadcompare, context, el1, el2) {
+  mjtBroadphase* b1 = (mjtBroadphase*)el1;
+  mjtBroadphase* b2 = (mjtBroadphase*)el2;
+
+  if (b1->value<b2->value) {
+    return -1;
+  } else if (b1->value==b2->value) {
+    return 0;
+  } else {
+    return 1;
+  }
+}
+
+
+
+// comparison function for pair sorting
+quicksortfunc(paircompare, context, el1, el2) {
+  int signature1 = *(int*)el1;
+  int signature2 = *(int*)el2;
+
+  if (signature1<signature2) {
+    return -1;
+  } else if (signature1==signature2) {
+    return 0;
+  } else {
+    return 1;
+  }
+}
+
+
+
+// does body have collidable geoms
+static int can_collide(const mjModel* m, int b) {
+  int g;
+
+  // scan geoms; return if collidable
+  for (g=0; g<m->body_geomnum[b]; g++) {
+    int ind = m->body_geomadr[b] + g;
+    if (m->geom_contype[ind] || m->geom_conaffinity[ind]) {
+      return 1;
+    }
+  }
+
+  // none found
+  return 0;
+}
+
+
+
+// broadphase collision detector
+int mj_broadphase(const mjModel* m, mjData* d, int* pair, int maxpair) {
+  int i, j, b1, b2, toremove, cnt, npair = 0, nbody = m->nbody, ngeom = m->ngeom;
+  mjtNum cov[9], cen[3], dif[3], eigval[3], frame[9], quat[4];
+  mjtBroadphase *sortbuf, *activebuf;
+  mjtNum *aabb;
+  mjMARKSTACK;
+
+  // world with geoms, and body with plane or hfield, can collide all bodies
+  for (b1=0; b1<nbody; b1++) {
+    // cannot colide
+    if (!can_collide(m, b1)) {
+      continue;
+    }
+
+    // world with geoms, or welded body with plane or hfield
+    if ((b1==0 && m->body_geomnum[b1]>0) || (m->body_weldid[b1]==0 && has_plane_or_hfield(m, b1))) {
+      for (b2=0; b2<nbody; b2++) {
+        if (b1!=b2) {
+          add_pair(NULL, b1, b2, &npair, pair, maxpair);
+        }
+      }
+    }
+  }
+
+  // find center of non-world geoms; return if none
+  cnt = 0;
+  mju_zero3(cen);
+  for (i=0; i<ngeom; i++) {
+    if (m->geom_bodyid[i]) {
+      mju_addTo3(cen, d->geom_xpos+3*i);
+      cnt++;
+    }
+  }
+  if (cnt==0) {
+    return npair;
+  } else {
+    for (i=0; i<3; i++) {
+      cen[i] /= cnt;
+    }
+  }
+
+  // compute covariance
+  mju_zero(cov, 9);
+  for (i=0; i<ngeom; i++) {
+    if (m->geom_bodyid[i]) {
+      mju_sub3(dif, d->geom_xpos+3*i, cen);
+      mjtNum D00 = dif[0]*dif[0];
+      mjtNum D01 = dif[0]*dif[1];
+      mjtNum D02 = dif[0]*dif[2];
+      mjtNum D11 = dif[1]*dif[1];
+      mjtNum D12 = dif[1]*dif[2];
+      mjtNum D22 = dif[2]*dif[2];
+      cov[0] += D00;
+      cov[1] += D01;
+      cov[2] += D02;
+      cov[3] += D01;
+      cov[4] += D11;
+      cov[5] += D12;
+      cov[6] += D02;
+      cov[7] += D12;
+      cov[8] += D22;
+    }
+  }
+  for (i=0; i<9; i++) {
+    cov[i] /= cnt;
+  }
+
+  // construct covariance-aligned 3D frame
+  mju_eig3(eigval, frame, quat, cov);
+
+  // allocate AABB; clear world entry (not used)
+  aabb = mj_stackAlloc(d, 6*nbody);
+  mju_zero(aabb, 6);
+
+  // construct body AABB for the aligned frame, count collidable
+  int bufcnt = 0;
+  for (i=1; i<nbody; i++) {
+    makeAABB(m, d, aabb+6*i, i, frame);
+
+    if (can_collide(m, i)) {
+      bufcnt++;
+    }
+  }
+
+  // nothing collidable
+  if (!bufcnt) {
+    goto endbroad;
+  }
+
+  // allocate sort buffer
+  i = sizeof(mjtBroadphase)/sizeof(mjtNum);
+  j = sizeof(mjtBroadphase)%sizeof(mjtNum);
+  sortbuf = (mjtBroadphase*)mj_stackAlloc(d, 2*bufcnt*(i + (j ? 1 : 0)));
+  activebuf = (mjtBroadphase*)mj_stackAlloc(d, 2*bufcnt*(i + (j ? 1 : 0)));
+
+  // init sortbuf with axis0
+  j = 0;
+  for (i=1; i<nbody; i++) {
+    // cannot colide
+    if (!can_collide(m, i)) {
+      continue;
+    }
+
+    // init
+    sortbuf[2*j].body_ismax = i;
+    sortbuf[2*j].value = (float)aabb[6*i];
+    sortbuf[2*j+1].body_ismax = i + 0x10000;
+    sortbuf[2*j+1].value = (float)aabb[6*i+1];
+    j++;
+  }
+
+  // sanity check; SHOULD NOT OCCUR
+  if (j!=bufcnt) {
+    mju_error("Internal error in broadphase: unexpected bufcnt");
+  }
+
+  // sort along axis0
+  mjQUICKSORT(sortbuf, 2*bufcnt, sizeof(mjtBroadphase), broadcompare, 0);
+
+  // sweep and prune
+  cnt = 0;    // size of active list
+  for (i=0; i<2*bufcnt; i++) {
+    // min value: collide with all in list, add
+    if (!(sortbuf[i].body_ismax & 0x10000)) {
+      for (j=0; j<cnt; j++) {
+        // get body ids
+        b1 = activebuf[j].body_ismax;
+        b2 = sortbuf[i].body_ismax;
+
+        // use the other two axes to prune if possible
+        if (aabb[6*b1+2] > aabb[6*b2+3] ||
+            aabb[6*b1+3] < aabb[6*b2+2] ||
+            aabb[6*b1+4] > aabb[6*b2+5] ||
+            aabb[6*b1+5] < aabb[6*b2+4]) {
+          continue;
+        }
+
+        // add body pair if there is room in buffer
+        add_pair(m, b1, b2, &npair, pair, maxpair);
+      }
+
+      // add to list
+      activebuf[cnt] = sortbuf[i];
+      cnt++;
+    }
+
+    // max value: remove corresponding min value from list
+    else {
+      toremove = sortbuf[i].body_ismax & 0xFFFF;
+      for (j=0; j<cnt; j++) {
+        if (activebuf[j].body_ismax==toremove) {
+          if (j<cnt-1) {
+            memmove(activebuf+j, activebuf+j+1, sizeof(mjtBroadphase)*(cnt-1-j));
+          }
+          cnt--;
+          break;
+        }
+      }
+    }
+  }
+
+endbroad:
+
+  // sort pairs by signature
+  if (npair) {
+    mjQUICKSORT(pair, npair, sizeof(int), paircompare, 0);
+  }
+
+  mjFREESTACK;
+  return npair;
+}
+
+
+
+//----------------------------- narrow-phase collision detection -----------------------------------
+
+// plane : geom_center distance, assuming g1 is plane
+static mjtNum plane_geom(const mjModel* m, mjData* d, int g1, int g2) {
+  mjtNum* mat1 = d->geom_xmat + 9*g1;
+  mjtNum norm[3] = {mat1[2], mat1[5], mat1[8]};
+  mjtNum dif[3];
+
+  mju_sub3(dif, d->geom_xpos + 3*g2, d->geom_xpos + 3*g1);
+  return mju_dot3(dif, norm);
+}
+
+
+// test two geoms for collision, apply filters, add to contact list
+//  flg_user disables filters and uses usermargin
+void mj_collideGeoms(const mjModel* m, mjData* d, int g1, int g2, int flg_user, mjtNum usermargin) {
+  int i, num, type1, type2, b1, b2, weld1, weld2, condim;
+  mjtNum margin, gap, mix, friction[5], solref[mjNREF], solimp[mjNIMP];
+  mjContact con[mjMAXCONPAIR];
+  int ipair = (g2<0 ? g1 : -1);
+
+  // get explicit geom ids from pair
+  if (ipair>=0) {
+    g1 = m->pair_geom1[ipair];
+    g2 = m->pair_geom2[ipair];
+  }
+
+  // order geoms by type
+  if (m->geom_type[g1] > m->geom_type[g2]) {
+    i = g1;
+    g1 = g2;
+    g2 = i;
+  }
+
+  // copy types and bodies
+  type1 = m->geom_type[g1];
+  type2 = m->geom_type[g2];
+  b1 = m->geom_bodyid[g1];
+  b2 = m->geom_bodyid[g2];
+  weld1 = m->body_weldid[b1];
+  weld2 = m->body_weldid[b2];
+
+  // return if no collision function
+  if (!mjCOLLISIONFUNC[type1][type2]) {
+    return;
+  }
+
+  // apply filters if not predefined pair and not flg_user
+  if (ipair<0 && !flg_user) {
+    // user filter if defined
+    if (mjcb_contactfilter) {
+      if (mjcb_contactfilter(m, d, g1, g2)) {
+        return;
+      }
+    }
+
+    // otherwise built-in filter
+    else if (mj_contactFilter(
+               type1, m->geom_contype[g1], m->geom_conaffinity[g1],
+               weld1, m->body_weldid[m->body_parentid[weld1]],
+               type2, m->geom_contype[g2], m->geom_conaffinity[g2],
+               weld2, m->body_weldid[m->body_parentid[weld2]],
+               !mjDISABLED(mjDSBL_FILTERPARENT) && weld1 && weld2)) {
+      return;
+    }
+  }
+
+  // set margin, gap, condim: dynamic
+  if (ipair<0) {
+    // margin and gap: max
+    margin = mju_max(m->geom_margin[g1], m->geom_margin[g2]);
+    gap = mju_max(m->geom_gap[g1], m->geom_gap[g2]);
+
+    // condim: priority or max
+    if (m->geom_priority[g1]!=m->geom_priority[g2]) {
+      int gp = (m->geom_priority[g1]>m->geom_priority[g2] ? g1 : g2);
+      condim = m->geom_condim[gp];
+    } else {
+      condim = mjMAX(m->geom_condim[g1], m->geom_condim[g2]);
+    }
+  }
+
+  // set margin, gap, condim: pair
+  else {
+    margin = m->pair_margin[ipair];
+    gap = m->pair_gap[ipair];
+    condim = m->pair_dim[ipair];
+  }
+
+  // adjust margin
+  if (flg_user) {
+    margin = usermargin;
+  } else {
+    margin = mj_assignMargin(m, margin);
+  }
+
+  // bounding sphere filter
+  if (m->geom_rbound[g1]>0 && m->geom_rbound[g2]>0 &&
+      (mju_dist3(d->geom_xpos+3*g1, d->geom_xpos+3*g2) >
+       m->geom_rbound[g1] + m->geom_rbound[g2] + margin)) {
+    return;
+  }
+
+  // plane : bounding sphere filter
+  if (m->geom_type[g1]==mjGEOM_PLANE && m->geom_rbound[g2]>0
+      && plane_geom(m, d, g1, g2) > margin+m->geom_rbound[g2]) {
+      return;
+  }
+  if (m->geom_type[g2]==mjGEOM_PLANE && m->geom_rbound[g1]>0
+      && plane_geom(m, d, g2, g1) > margin+m->geom_rbound[g1]) {
+      return;
+  }
+
+  // call collision detector to generate contacts
+  num = mjCOLLISIONFUNC[type1][type2](m, d, con, g1, g2, margin);
+
+  // no contacts from near-phase
+  if (!num) {
+    return;
+  }
+
+  // check number of contacts, SHOULD NOT OCCUR
+  if (num>mjMAXCONPAIR) {
+    mju_error("Too many contacts returned by collision function");
+  }
+
+  // remove repeated contacts in box-box
+  if (type1==mjGEOM_BOX && type2==mjGEOM_BOX) {
+    // use dim field to mark: -1: bad, 0: good
+    for (i=0; i<num; i++) {
+      con[i].dim = 0;
+    }
+
+    // find bad
+    int j;
+    for (i=0; i<num-1; i++) {
+      for (j=i+1; j<num; j++) {
+        if (con[i].pos[0]==con[j].pos[0] &&
+            con[i].pos[1]==con[j].pos[1] &&
+            con[i].pos[2]==con[j].pos[2]) {
+          con[i].dim = -1;
+          break;
+        }
+      }
+    }
+
+    // consolidate good
+    i = 0;
+    for (j=0; j<num; j++) {
+      if (con[j].dim==0) {
+        // different: copy
+        if (i<j) {
+          con[i] = con[j];
+        }
+
+        // advance either way
+        i++;
+      }
+    }
+
+    // adjust size
+    num = i;
+  }
+
+  // set friction, solref, solimp: dynamic
+  if (ipair<0) {
+    // different priority
+    if (m->geom_priority[g1]!=m->geom_priority[g2]) {
+      int gp = (m->geom_priority[g1]>m->geom_priority[g2] ? g1 : g2);
+
+      // friction
+      for (i=0; i<3; i++) {
+        friction[2*i] = m->geom_friction[3*gp+i];
+      }
+
+      // reference
+      mju_copy(solref, m->geom_solref+mjNREF*gp, mjNREF);
+
+      // impedance
+      mju_copy(solimp, m->geom_solimp+mjNIMP*gp, mjNIMP);
+    }
+
+    // same priority
+    else {
+      // friction: max
+      for (i=0; i<3; i++) {
+        friction[2*i] = mju_max(m->geom_friction[3*g1+i], m->geom_friction[3*g2+i]);
+      }
+
+      // solver mix factor
+      if (m->geom_solmix[g1]>=mjMINVAL && m->geom_solmix[g2]>=mjMINVAL) {
+        mix = m->geom_solmix[g1] / (m->geom_solmix[g1] + m->geom_solmix[g2]);
+      } else if (m->geom_solmix[g1]<mjMINVAL && m->geom_solmix[g2]<mjMINVAL) {
+        mix = 0.5;
+      } else if (m->geom_solmix[g1]<mjMINVAL) {
+        mix = 0.0;
+      } else {
+        mix = 1.0;
+      }
+
+      // reference standard: mix
+      if (m->geom_solref[mjNREF*g1]>0 && m->geom_solref[mjNREF*g2]>0) {
+        for (i=0; i<mjNREF; i++) {
+          solref[i] = mix*m->geom_solref[mjNREF*g1+i] + (1-mix)*m->geom_solref[mjNREF*g2+i];
+        }
+      }
+
+      // reference direct: min
+      else {
+        for (i=0; i<mjNREF; i++) {
+          solref[i] = mju_min(m->geom_solref[mjNREF*g1+i], m->geom_solref[mjNREF*g2+i]);
+        }
+      }
+
+      // impedance: mix
+      mju_scl(solimp, m->geom_solimp+mjNIMP*g1, mix, mjNIMP);
+      mju_addToScl(solimp, m->geom_solimp+mjNIMP*g2, 1-mix, mjNIMP);
+    }
+
+    // unpack 5D friction
+    friction[1] = friction[0];
+    friction[3] = friction[4];
+  }
+
+  // set friction, solref, solimp: pair
+  else {
+    // friction
+    for (i=0; i<5; i++) {
+      friction[i] = m->pair_friction[5*ipair+i];
+    }
+
+    // reference
+    mju_copy(solref, m->pair_solref+mjNREF*ipair, mjNREF);
+
+    // impedance
+    mju_copy(solimp, m->pair_solimp+mjNIMP*ipair, mjNIMP);
+  }
+
+  // clamp friction to mjMINMU
+  for (i=0; i<5; i++) {
+    friction[i] = mju_max(mjMINMU, friction[i]);
+  }
+
+  // add contact returned by collision detector
+  for (i=0; i<num; i++) {
+    // set contact data
+    if (condim > 6 || condim < 0) {  // SHOULD NOT OCCUR
+      mju_error_i("Invalid condim value: %d", i);
+    }
+    con[i].dim = condim;
+    con[i].geom1 = g1;
+    con[i].geom2 = g2;
+    con[i].includemargin = margin-gap;
+    mju_copy(con[i].friction, friction, 5);
+    mj_assignRef(m, con[i].solref, solref);
+    mj_assignImp(m, con[i].solimp, solimp);
+
+    // exclude in gap
+    if (con[i].dist<con[i].includemargin) {
+      con[i].exclude = 0;
+    } else {
+      con[i].exclude = 1;
+    }
+
+    // complete frame
+    mju_makeFrame(con[i].frame);
+
+    // clear fields that are computed later
+    con[i].efc_address = -1;
+    con[i].mu = 0;
+    mju_zero(con[i].H, 36);
+    // add to mjData, abort if too many contacts
+    if (mj_addContact(m, d, con + i)) {
+      return;
+    }
+  }
+}
+
+
+
+// filter contacts: 1- discard, 0- proceed
+int mj_contactFilter(int type1, int contype1, int conaffinity1, int weldbody1, int weldparent1,
+                     int type2, int contype2, int conaffinity2, int weldbody2, int weldparent2,
+                     int filterparent) {
+  // compatibility check
+  if (!(contype1 & conaffinity2) && !(contype2 & conaffinity1)) {
+    return 1;
+  }
+
+  // same weldbody check
+  if (weldbody1==weldbody2) {
+    return 1;
+  }
+
+  // weldparent check
+  if (filterparent && (weldbody1==weldparent2 || weldbody2==weldparent1)) {
+    return 1;
+  }
+
+  // all tests passed
+  return 0;
+}

src/engine/engine_collision_driver.h

@@ -0,0 +1,49 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_COLLISION_DRIVER_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_COLLISION_DRIVER_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// collision function pointers and max contact pairs
+MJAPI extern mjfCollision mjCOLLISIONFUNC[mjNGEOMTYPES][mjNGEOMTYPES];
+
+// collision detection entry point
+MJAPI void mj_collision(const mjModel* m, mjData* d);
+
+// broad phase collistion detection; return list of body pairs for narrow phase
+int mj_broadphase(const mjModel* m, mjData* d, int* bodypair, int maxpair);
+
+// test two geoms for collision, apply filters, add to contact list
+//  flg_user disables filters and uses usermargin
+void mj_collideGeoms(const mjModel* m, mjData* d,
+                     int g1, int g2, int flg_user, mjtNum usermargin);
+
+// number of possible collisions based on fitlers and geom types
+int mj_contactFilter(int type1, int contype1, int conaffinity1, int weldbody1, int weldparent1,
+                     int type2, int contype2, int conaffinity2, int weldbody2, int weldparent2,
+                     int filterparent);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_COLLISION_DRIVER_H_

src/engine/engine_collision_primitive.c

@@ -0,0 +1,460 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_collision_primitive.h"
+
+#include <math.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_spatial.h"
+
+
+//--------------------------- plane collisions -----------------------------------------------------
+
+// plane : sphere (actual implementation, can be called with modified parameters)
+static int _PlaneSphere(mjContact* con, mjtNum margin,
+                        mjtNum* pos1, mjtNum* mat1, mjtNum* size1,
+                        mjtNum* pos2, mjtNum* mat2, mjtNum* size2) {
+  mjtNum tmp[3];
+  mjtNum cdist;
+
+  // set normal
+  con[0].frame[0] = mat1[2];
+  con[0].frame[1] = mat1[5];
+  con[0].frame[2] = mat1[8];
+
+  // compute distance, return if too large
+  mju_sub3(tmp, pos2, pos1);
+  cdist = mju_dot3(tmp, con[0].frame);
+  if (cdist > margin + size2[0]) {
+    return 0;
+  }
+
+  // depth and position
+  con[0].dist = cdist - size2[0];
+  mju_scl3(tmp, con[0].frame, -con[0].dist/2 - size2[0]);
+  mju_add3(con[0].pos, pos2, tmp);
+
+  mju_zero3(con[0].frame+3);
+  return 1;
+}
+
+
+
+// plane : sphere
+int mjc_PlaneSphere(const mjModel* m, const mjData* d,
+                    mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  return _PlaneSphere(con, margin, pos1, mat1, size1, pos2, mat2, size2);
+}
+
+
+
+// plane : capsule
+int mjc_PlaneCapsule(const mjModel* m, const mjData* d,
+                     mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  mjtNum pos[3], axis[3], segment[3];
+  int n1, n2;
+
+  // get capsule axis, segment = scaled axis
+  axis[0] = mat2[2];
+  axis[1] = mat2[5];
+  axis[2] = mat2[8];
+  mju_scl3(segment, axis, size2[1]);
+
+  // get point 1, do sphere-plane test
+  mju_add3(pos, pos2, segment);
+  n1 = _PlaneSphere(con, margin, pos1, mat1, size1, pos, mat2, size2);
+
+  // get point 2, do sphere-plane test
+  mju_sub3(pos, pos2, segment);
+  n2 = _PlaneSphere(con+n1, margin, pos1, mat1, size1, pos, mat2, size2);
+
+  // align contact frames with capsule axis
+  if (n1) {
+    mju_copy3(con->frame+3, axis);
+  }
+  if (n2) {
+    mju_copy3((con+n1)->frame+3, axis);
+  }
+
+  return n1+n2;
+}
+
+
+
+// plane : cylinder
+int mjc_PlaneCylinder(const mjModel* m, const mjData* d,
+                      mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  mjtNum normal[3] = {mat1[2], mat1[5], mat1[8]};
+  mjtNum axis[3] = {mat2[2], mat2[5], mat2[8]};
+  mjtNum vec[3], vec1[3];
+  mjtNum len, scl, dist0, prjaxis, prjvec, prjvec1;
+  int cnt = 0;
+
+  // project, make sure axis points towards plane
+  prjaxis = mju_dot3(normal, axis);
+  if (prjaxis > 0) {
+    mju_scl3(axis, axis, -1);
+    prjaxis = -prjaxis;
+  }
+
+  // compute normal distance to cylinder center
+  mju_sub3(vec, pos2, pos1);
+  dist0 = mju_dot3(vec, normal);
+
+  // remove component of -normal along axis, compute length
+  mju_scl3(vec, axis, prjaxis);
+  mju_subFrom3(vec, normal);
+  len = mju_norm3(vec);
+
+  // general configuration: normalize vector, scale by radius
+  if (len >= mjMINVAL) {
+    scl = size2[0]/len;
+    vec[0] *= scl;
+    vec[1] *= scl;
+    vec[2] *= scl;
+  }
+
+  // disk parallel to plane: pick x-axis of cylinder, scale by radius
+  else {
+    vec[0] = mat2[0]*size2[0];
+    vec[1] = mat2[3]*size2[0];
+    vec[2] = mat2[6]*size2[0];
+  }
+
+  // project vector on normal
+  prjvec = mju_dot3(vec, normal);
+
+  // scale axis by half-length
+  mju_scl3(axis, axis, size2[1]);
+  prjaxis *= size2[1];
+
+  // check first point, construct contact
+  if (dist0 + prjaxis + prjvec <= margin) {
+    con[cnt].dist = dist0 + prjaxis + prjvec;
+    mju_add3(con[cnt].pos, pos2, vec);
+    mju_addTo3(con[cnt].pos, axis);
+    mju_addToScl3(con[cnt].pos, normal, -con[cnt].dist*0.5);
+    mju_copy3(con[cnt].frame, normal);
+    mju_zero3(con[cnt].frame+3);
+    cnt++;
+  } else {
+    return 0;  // nearest point is above margin: no contacts
+  }
+
+  // check second point, construct contact
+  if (dist0 - prjaxis + prjvec <= margin) {
+    con[cnt].dist = dist0 - prjaxis + prjvec;
+    mju_add3(con[cnt].pos, pos2, vec);
+    mju_subFrom3(con[cnt].pos, axis);
+    mju_addToScl3(con[cnt].pos, normal, -con[cnt].dist*0.5);
+    mju_copy3(con[cnt].frame, normal);
+    mju_zero3(con[cnt].frame+3);
+    cnt++;
+  }
+
+  // try to add triangle points on side closer to plane
+  prjvec1 = -prjvec*0.5;
+  if (dist0 + prjaxis + prjvec1 <= margin) {
+    // compute sideways vector: vec1
+    mju_cross(vec1, vec, axis);
+    mju_normalize3(vec1);
+    mju_scl3(vec1, vec1, size2[0]*mju_sqrt(3.0)/2);
+
+    // add point A
+    con[cnt].dist = dist0 + prjaxis + prjvec1;
+    mju_add3(con[cnt].pos, pos2, vec1);
+    mju_addTo3(con[cnt].pos, axis);
+    mju_addToScl3(con[cnt].pos, vec, -0.5);
+    mju_addToScl3(con[cnt].pos, normal, -con[cnt].dist*0.5);
+    mju_copy3(con[cnt].frame, normal);
+    mju_zero3(con[cnt].frame+3);
+    cnt++;
+
+    // add point B
+    con[cnt].dist = dist0 + prjaxis + prjvec1;
+    mju_sub3(con[cnt].pos, pos2, vec1);
+    mju_addTo3(con[cnt].pos, axis);
+    mju_addToScl3(con[cnt].pos, vec, -0.5);
+    mju_addToScl3(con[cnt].pos, normal, -con[cnt].dist*0.5);
+    mju_copy3(con[cnt].frame, normal);
+    mju_zero3(con[cnt].frame+3);
+    cnt++;
+  }
+
+  return cnt;
+}
+
+
+
+// plane : box
+int mjc_PlaneBox(const mjModel* m, const mjData* d,
+                 mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  int cnt = 0;
+
+  // get normal, difference between centers, normal distance
+  mjtNum norm[3] = {mat1[2], mat1[5], mat1[8]};
+  mjtNum dif[3], vec[3], corner[3], dist, ldist;
+  mju_sub3(dif, pos2, pos1);
+  dist = mju_dot3(dif, norm);
+
+  // test all corners, pick bottom 4
+  for (int i=0; i<8; i++) {
+    // get corner in local coordinates
+    vec[0] = (i&1 ? size2[0] : -size2[0]);
+    vec[1] = (i&2 ? size2[1] : -size2[1]);
+    vec[2] = (i&4 ? size2[2] : -size2[2]);
+
+    // get corner in global coordinates relative to box center
+    mju_rotVecMat(corner, vec, mat2);
+
+    // compute distance to plane, skip if too far or pointing up
+    ldist = mju_dot3(norm, corner);
+    if (dist + ldist > margin || ldist > 0) {
+      continue;
+    }
+
+    // construct contact
+    con[cnt].dist = dist + ldist;
+    mju_copy3(con[cnt].frame, norm);
+    mju_zero3(con[cnt].frame+3);
+    mju_addTo3(corner, pos2);
+    mju_scl3(vec, norm, -con[cnt].dist/2);
+    mju_add3(con[cnt].pos, corner, vec);
+
+    // count; max is 4
+    if (++cnt >= 4) {
+      return 4;
+    }
+  }
+
+  return cnt;
+}
+
+
+
+//--------------------------- sphere and capsule collisions ----------------------------------------
+
+// sphere : sphere (actual implementation, can be called with modified parameters)
+static int _SphereSphere(mjContact* con, mjtNum margin,
+                         mjtNum* pos1, mjtNum* mat1, mjtNum* size1,
+                         mjtNum* pos2, mjtNum* mat2, mjtNum* size2) {
+  mjtNum len, cdist;
+  mjtNum axis1[3], axis2[3];
+
+  // check bounding spheres (this is called from other functions)
+  cdist = mju_dist3(pos1, pos2);
+  if (cdist > margin + size1[0] + size2[0]) {
+    return 0;
+  }
+
+  // depth and normal
+  con[0].dist = cdist - size1[0] - size2[0];
+  mju_sub3(con[0].frame, pos2, pos1);
+  len = mju_normalize3(con[0].frame);
+
+  // if centers are the same, norm = cross-product of z axes
+  //  if z axes are parallel, norm = [1;0;0]
+  if (len < mjMINVAL) {
+    axis1[0] = mat1[2];
+    axis1[1] = mat1[5];
+    axis1[2] = mat1[8];
+    axis2[0] = mat2[2];
+    axis2[1] = mat2[5];
+    axis2[2] = mat2[8];
+    mju_cross(con[0].frame, axis1, axis2);
+    mju_normalize3(con[0].frame);
+  }
+
+  // position
+  mju_scl3(con[0].pos, con[0].frame, size1[0] + con[0].dist/2);
+  mju_addTo3(con[0].pos, pos1);
+
+  mju_zero3(con[0].frame+3);
+  return 1;
+}
+
+
+
+// sphere : sphere
+int mjc_SphereSphere(const mjModel* m, const mjData* d,
+                     mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  return  _SphereSphere(con, margin, pos1, mat1, size1, pos2, mat2, size2);
+}
+
+
+
+// sphere : capsule
+int mjc_SphereCapsule(const mjModel* m, const mjData* d,
+                      mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  mjtNum x, axis[3], vec[3];
+
+  // get capsule axis (scaled)
+  axis[0] = mat2[2] * size2[1];
+  axis[1] = mat2[5] * size2[1];
+  axis[2] = mat2[8] * size2[1];
+
+  // find projection, clip to segment
+  mju_sub3(vec, pos1, pos2);
+  x = mju_dot3(axis, vec) / mju_dot3(axis, axis);
+  if (x > 1) {
+    x = 1;
+  } else if (x < -1) {
+    x = -1;
+  }
+
+  // find nearest point on segment, do sphere-sphere test
+  mju_scl3(vec, axis, x);
+  mju_addTo3(vec, pos2);
+  return _SphereSphere(con, margin, pos1, mat1, size1, vec, mat2, size2);
+}
+
+
+
+// capsule : capsule
+int mjc_CapsuleCapsule(const mjModel* m, const mjData* d,
+                       mjContact* con, int g1, int g2, mjtNum margin) {
+  mjGETINFO
+  mjtNum axis1[3], axis2[3], dif[3], vec1[3], vec2[3];
+  mjtNum ma, mb, mc, u, v, det, x1, x2;
+  int n1, n2, n3, n4;
+
+  // get capsule axes (scaled) and center difference
+  axis1[0] = mat1[2] * size1[1];
+  axis1[1] = mat1[5] * size1[1];
+  axis1[2] = mat1[8] * size1[1];
+  axis2[0] = mat2[2] * size2[1];
+  axis2[1] = mat2[5] * size2[1];
+  axis2[2] = mat2[8] * size2[1];
+  mju_sub3(dif, pos1, pos2);
+
+  // compute matrix coefficients and determinant
+  ma =  mju_dot3(axis1, axis1);
+  mb = -mju_dot3(axis1, axis2);
+  mc =  mju_dot3(axis2, axis2);
+  u =  -mju_dot3(axis1, dif);
+  v =   mju_dot3(axis2, dif);
+  det = ma*mc - mb*mb;
+
+  // general configuration (non-parallel axes)
+  if (fabs(det) >= mjMINVAL) {
+    // find projections, clip to segments
+    x1 = (mc*u - mb*v) / det;
+    x2 = (ma*v - mb*u) / det;
+
+    if (x1 > 1) {
+      x1 = 1;
+      x2 = (v-mb)/mc;
+    } else if (x1 < -1) {
+      x1 = -1;
+      x2 = (v+mb)/mc;
+    }
+    if (x2 > 1) {
+      x2 = 1;
+      x1 = (u-mb)/ma;
+      if (x1 > 1) {
+        x1 = 1;
+      } else if (x1 < -1) {
+        x1 = -1;
+      }
+    } else if (x2 < -1) {
+      x2 = -1;
+      x1 = (u+mb)/ma;
+      if (x1 > 1) {
+        x1 = 1;
+      } else if (x1 < -1) {
+        x1 = -1;
+      }
+    }
+
+    // find nearest points, do sphere-sphere test
+    mju_scl3(vec1, axis1, x1);
+    mju_addTo3(vec1, pos1);
+    mju_scl3(vec2, axis2, x2);
+    mju_addTo3(vec2, pos2);
+
+    return _SphereSphere(con, margin, vec1, mat1, size1, vec2, mat2, size2);
+  }
+
+  // parallel axes
+  else {
+    // x1 = 1
+    mju_add3(vec1, pos1, axis1);
+    x2 = (v - mb) / mc;
+    if (x2 > 1) {
+      x2 = 1;
+    } else if (x2 < -1) {
+      x2 = -1;
+    }
+    mju_scl3(vec2, axis2, x2);
+    mju_addTo3(vec2, pos2);
+    n1 = _SphereSphere(con, margin, vec1, mat1, size1, vec2, mat2, size2);
+
+    // x1 = -1
+    mju_sub3(vec1, pos1, axis1);
+    x2 = (v + mb) / mc;
+    if (x2 > 1) {
+      x2 = 1;
+    } else if (x2 < -1) {
+      x2 = -1;
+    }
+    mju_scl3(vec2, axis2, x2);
+    mju_addTo3(vec2, pos2);
+    n2 = _SphereSphere(con+n1, margin, vec1, mat1, size1, vec2, mat2, size2);
+
+    // return if two contacts already found
+    if (n1+n2>=2) {
+      return n1+n2;
+    }
+
+    // x2 = 1
+    mju_add3(vec2, pos2, axis2);
+    x1 = (u - mb) / ma;
+    if (x1 > 1) {
+      x1 = 1;
+    } else if (x1 < -1) {
+      x1 = -1;
+    }
+    mju_scl3(vec1, axis1, x1);
+    mju_addTo3(vec1, pos1);
+    n3 = _SphereSphere(con+n1+n2, margin, vec1, mat1, size1, vec2, mat2, size2);
+
+    // return if two contacts already found
+    if (n1+n2+n3>=2) {
+      return n1+n2+n3;
+    }
+
+    // x2 = -1
+    mju_sub3(vec2, pos2, axis2);
+    x1 = (u + mb) / ma;
+    if (x1 > 1) {
+      x1 = 1;
+    } else if (x1 < -1) {
+      x1 = -1;
+    }
+    mju_scl3(vec1, axis1, x1);
+    mju_addTo3(vec1, pos1);
+    n4 = _SphereSphere(con+n1+n2+n3, margin, vec1, mat1, size1, vec2, mat2, size2);
+
+    return n1+n2+n3+n4;
+  }
+}

src/engine/engine_collision_primitive.h

@@ -0,0 +1,64 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_COLLISION_PRIMITIVE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_COLLISION_PRIMITIVE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+
+// define and extract geom info
+#define mjGETINFO \
+    mjtNum* pos1 = d->geom_xpos + 3*g1; \
+    mjtNum* mat1 = d->geom_xmat + 9*g1; \
+    mjtNum* size1= m->geom_size + 3*g1; \
+    mjtNum* pos2 = d->geom_xpos + 3*g2; \
+    mjtNum* mat2 = d->geom_xmat + 9*g2; \
+    mjtNum* size2= m->geom_size + 3*g2; \
+    (void) size1; (void) size2;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// plane collisions
+int mjc_PlaneSphere     (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_PlaneCapsule    (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_PlaneCylinder   (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_PlaneBox        (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+
+// sphere and capsule collisions
+int mjc_SphereSphere    (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_SphereCapsule   (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_CapsuleCapsule  (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+
+// box collisions: from boxcollisions.c
+int mjc_CapsuleBox      (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_SphereBox       (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_BoxBox          (const mjModel* m, const mjData* d,
+                         mjContact* con, int g1, int g2, mjtNum margin);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_COLLISION_PRIMITIVE_H_

src/engine/engine_core_constraint.h

@@ -0,0 +1,115 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_CORE_CONSTRAINT_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_CORE_CONSTRAINT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//-------------------------- Jacobian-related ------------------------------------------------------
+
+// determine type of friction cone
+MJAPI int mj_isPyramidal(const mjModel* m);
+
+// determine type of constraint Jacobian
+MJAPI int mj_isSparse(const mjModel* m);
+
+// determine type of solver
+MJAPI int mj_isDual(const mjModel* m);
+
+// multiply Jacobian by vector
+MJAPI void mj_mulJacVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec);
+
+// multiply JacobianT by vector
+MJAPI void mj_mulJacTVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec);
+
+
+//-------------------------- utility functions -----------------------------------------------------
+
+// assign/override solver reference parameters
+void mj_assignRef(const mjModel* m, mjtNum* target, const mjtNum* source);
+
+// assign/override solver impedance parameters
+void mj_assignImp(const mjModel* m, mjtNum* target, const mjtNum* source);
+
+// assign/override geom/limit/tendon margin
+mjtNum mj_assignMargin(const mjModel* m, mjtNum source);
+
+// add contact to d->contact list; return 0 if success; 1 if buffer full
+MJAPI int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
+
+// add #size rows to constraint Jacobian; set pos, margin, frictionloss, type, id
+//  result: 0=success; 1=buffer full
+int mj_addConstraint(const mjModel* m, mjData* d,
+                     const mjtNum* jac, const mjtNum* pos,
+                     const mjtNum* margin, mjtNum frictionloss,
+                     int size, int type, int id, int NV, const int* chain);
+
+// merge dof chains for two bodies
+int mj_mergeChain(const mjModel* m, int* dofid, int b1, int b2);
+
+// merge dof chains for two simple bodies
+int mj_mergeChainSimple(const mjModel* m, int* dofid, int b1, int b2);
+
+
+//-------------------------- constraint instantiation ----------------------------------------------
+
+// equality constraints
+void mj_instantiateEquality(const mjModel* m, mjData* d);
+
+// frictional dofs and tendons
+void mj_instantiateFriction(const mjModel* m, mjData* d);
+
+// joint and tendon limits
+void mj_instantiateLimit(const mjModel* m, mjData* d);
+
+// frictionelss and frictional contacts
+void mj_instantiateContact(const mjModel* m, mjData* d);
+
+
+//------------------------ parameter computation/extraction ----------------------------------------
+
+// compute efc_diagApprox
+void mj_diagApprox(const mjModel* m, mjData* d);
+
+// compute efc_R, efc_D, efc_KDIP, adjust diagApprox
+void mj_makeImpedance(const mjModel* m, mjData* d);
+
+
+//---------------------------- top-level API for constraint construction ---------------------------
+
+// main driver: call all functions above
+MJAPI void mj_makeConstraint(const mjModel* m, mjData* d);
+
+// compute efc_AR
+MJAPI void mj_projectConstraint(const mjModel* m, mjData* d);
+
+// compute efc_vel, efc_aref
+MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d);
+
+// compute efc_state, efc_force, qfrc_constraint
+// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians
+MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
+                               mjtNum cost[1], int flg_coneHessian);
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_CORE_CONSTRAINT_H_

src/engine/engine_core_smooth.h

@@ -0,0 +1,127 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_CORE_SMOOTH_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_CORE_SMOOTH_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+//-------------------------- position --------------------------------------------------------------
+
+// forward kinematics
+MJAPI void mj_kinematics(const mjModel* m, mjData* d);
+
+// map inertias and motion dofs to global frame centered at CoM
+MJAPI void mj_comPos(const mjModel* m, mjData* d);
+
+// compute camera and light positions and orientations
+MJAPI void mj_camlight(const mjModel* m, mjData* d);
+
+// compute tendon lengths, velocities and moment arms
+MJAPI void mj_tendon(const mjModel* m, mjData* d);
+
+// compute actuator transmission lengths and moments
+MJAPI void mj_transmission(const mjModel* m, mjData* d);
+
+
+//-------------------------- inertia ---------------------------------------------------------------
+
+// composite rigid body inertia algorithm, with skip
+void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple);
+
+// composite rigid body inertia algorithm
+MJAPI void mj_crb(const mjModel* m, mjData* d);
+
+// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
+MJAPI void mj_factorM(const mjModel* m, mjData* d);
+
+// sparse backsubstitution:  x = inv(L'*D*L)*y
+MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+
+// half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
+MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+
+
+//-------------------------- velocity --------------------------------------------------------------
+
+// compute cvel, cdof_dot
+MJAPI void mj_comVel(const mjModel* m, mjData* d);
+
+// passive forces
+MJAPI void mj_passive(const mjModel* m, mjData* d);
+
+// subtree linear velocity and angular momentum
+MJAPI void mj_subtreeVel(const mjModel* m, mjData* d);
+
+
+//------------------------- fluid model ------------------------------------------------------------
+
+
+void mj_inertiaBoxFluidModel(const mjModel* m, mjData* d, int i);
+
+void mj_ellipsoidFluidModel(const mjModel* m, mjData* d, int bodyid);
+
+// compute forces due to added mass (potential flow)
+void mj_addedMassForces(
+    const mjtNum local_vels[6], const mjtNum local_accels[6],
+    const mjtNum fluid_density, const mjtNum virtual_mass[3],
+    const mjtNum virtual_inertia[3], mjtNum local_force[6]);
+
+// compute forces due to viscous effects
+void mj_viscousForces(
+    const mjtNum local_vels[6], const mjtNum fluid_density,
+    const mjtNum fluid_viscosity, const mjtNum size[3],
+    const mjtNum magnus_lift_coef, const mjtNum kutta_lift_coef,
+    const mjtNum blunt_drag_coef, const mjtNum slender_drag_coef,
+    const mjtNum ang_drag_coef, mjtNum local_force[6]);
+
+void readFluidGeomInteraction(const mjtNum * geom_fluid_coefs,
+                              mjtNum * geom_fluid_coef,
+                              mjtNum * blunt_drag_coef,
+                              mjtNum * slender_drag_coef,
+                              mjtNum * ang_drag_coef,
+                              mjtNum * kutta_lift_coef,
+                              mjtNum * magnus_lift_coef,
+                              mjtNum virtual_mass[3],
+                              mjtNum virtual_inertia[3]);
+
+void writeFluidGeomInteraction (mjtNum * geom_fluid_coefs,
+                                const mjtNum * geom_fluid_coef,
+                                const mjtNum * blunt_drag_coef,
+                                const mjtNum * slender_drag_coef,
+                                const mjtNum * ang_drag_coef,
+                                const mjtNum * kutta_lift_coef,
+                                const mjtNum * magnus_lift_coef,
+                                const mjtNum virtual_mass[3],
+                                const mjtNum virtual_inertia[3]);
+
+
+//-------------------------- RNE -------------------------------------------------------------------
+
+// RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term
+MJAPI void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result);
+
+// RNE with complete data: compute cacc, cfrc_ext, cfrc_int
+MJAPI void mj_rnePostConstraint(const mjModel* m, mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_CORE_SMOOTH_H_

src/engine/engine_crossplatform.h

@@ -0,0 +1,49 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_CROSSPLATFORM_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_CROSSPLATFORM_H_
+
+#include <stdlib.h>
+
+// Windows
+#ifdef _WIN32
+//  #define isnan _isnan
+  #define strcasecmp _stricmp
+  #define strncasecmp _strnicmp
+
+  #define mjQUICKSORT(buf, elnum, elsz, func, context) \
+      qsort_s(buf, elnum, elsz, func, context)
+  #define quicksortfunc(name, context, el1, el2) \
+      static int name(void* context, const void* el1, const void* el2)
+
+// Unix-common
+#else
+  // Apple
+  #ifdef __APPLE__
+    #define mjQUICKSORT(buf, elnum, elsz, func, context) \
+        qsort_r(buf, elnum, elsz, context, func)
+    #define quicksortfunc(name, context, el1, el2) \
+        static int name(void* context, const void* el1, const void* el2)
+
+  // non-Apple
+  #else
+    #define mjQUICKSORT(buf, elnum, elsz, func, context) \
+        qsort_r(buf, elnum, elsz, func, context)
+    #define quicksortfunc(name, context, el1, el2) \
+        static int name(const void* el1, const void* el2, void* context)
+  #endif
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_CROSSPLATFORM_H_

src/engine/engine_file.c

@@ -0,0 +1,80 @@
+// Copyright 2022 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_file.h"
+
+#include <stdio.h>
+#include <limits.h>
+
+#include "engine/engine_util_errmem.h"
+
+void* mju_fileToMemory(const char* filename, int* filesize) {
+  // open file
+  *filesize = 0;
+  FILE* fp = fopen(filename, "rb");
+  if (!fp) {
+    return NULL;
+  }
+
+  // find size
+  if (fseek(fp, 0, SEEK_END) != 0) {
+    fclose(fp);
+    mju_warning_s("Failed to calculate size for '%s'", filename);
+    return NULL;
+  }
+
+  // ensure file size fits in int
+  long long_filesize = ftell(fp);
+  if (long_filesize > INT_MAX) {
+    fclose(fp);
+    mju_warning_s("File size over 2GB is not supported. File: '%s'", filename);
+    return NULL;
+  } else if (long_filesize < 0) {
+    fclose(fp);
+    mju_warning_s("Failed to calculate size for '%s'", filename);
+    return NULL;
+  }
+  *filesize = long_filesize;
+
+  // go back to start of file
+  if (fseek(fp, 0, SEEK_SET) != 0) {
+    fclose(fp);
+    mju_warning_s("Read error while reading '%s'", filename);
+    return NULL;
+  }
+
+  // allocate and read
+  void* buffer = mju_malloc(*filesize);
+  if (!buffer) {
+    mju_error("mjFileToMemory: could not allocate memory");
+  }
+  size_t bytes_read = fread(buffer, 1, *filesize, fp);
+
+  // check that read data matches file size
+  if (bytes_read != *filesize) { // SHOULD NOT OCCUR
+    if (ferror(fp)) {
+      fclose(fp);
+      mju_free(buffer);
+      *filesize = 0;
+      mju_warning_s("Read error while reading '%s'", filename);
+      return NULL;
+    } else if (feof(fp)) {
+      *filesize = bytes_read;
+    }
+  }
+
+  // close file, return contents
+  fclose(fp);
+  return buffer;
+}

src/engine/engine_file.h

@@ -0,0 +1,29 @@
+// Copyright 2022 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_FILE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_FILE_H_
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// read file into memory buffer (allocated here with mju_malloc)
+void* mju_fileToMemory(const char* filename, int* filesize);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_FILE_H_

src/engine/engine_forward.c

@@ -0,0 +1,768 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_forward.h"
+
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_callback.h"
+#include "engine/engine_collision_driver.h"
+#include "engine/engine_core_constraint.h"
+#include "engine/engine_core_smooth.h"
+#include "engine/engine_inverse.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_sensor.h"
+#include "engine/engine_solver.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_sparse.h"
+
+//--------------------------- check values ---------------------------------------------------------
+
+// check positions, reset if bad
+void mj_checkPos(const mjModel* m, mjData* d) {
+  for (int i=0; i<m->nq; i++) {
+    if (mju_isBad(d->qpos[i])) {
+      mj_warning(d, mjWARN_BADQPOS, i);
+      mj_resetData(m, d);
+      d->warning[mjWARN_BADQPOS].number++;
+      d->warning[mjWARN_BADQPOS].lastinfo = i;
+      return;
+    }
+  }
+}
+
+
+
+// check velocities, reset if bad
+void mj_checkVel(const mjModel* m, mjData* d) {
+  for (int i=0; i<m->nv; i++) {
+    if (mju_isBad(d->qvel[i])) {
+      mj_warning(d, mjWARN_BADQVEL, i);
+      mj_resetData(m, d);
+      d->warning[mjWARN_BADQVEL].number++;
+      d->warning[mjWARN_BADQVEL].lastinfo = i;
+      return;
+    }
+  }
+}
+
+
+
+// check accelerations, reset if bad
+void mj_checkAcc(const mjModel* m, mjData* d) {
+  for (int i=0; i<m->nv; i++) {
+    if (mju_isBad(d->qacc[i])) {
+      mj_warning(d, mjWARN_BADQACC, i);
+      mj_resetData(m, d);
+      d->warning[mjWARN_BADQACC].number++;
+      d->warning[mjWARN_BADQACC].lastinfo = i;
+      mj_forward(m, d);
+      return;
+    }
+  }
+}
+
+
+
+//-------------------------- solver components -----------------------------------------------------
+
+// position-dependent computations
+void mj_fwdPosition(const mjModel* m, mjData* d) {
+  TM_START1;
+
+  TM_START;
+  mj_kinematics(m, d);
+  mj_comPos(m, d);
+  mj_camlight(m, d);
+  mj_tendon(m, d);
+  mj_transmission(m, d);
+  TM_END(mjTIMER_POS_KINEMATICS);
+
+  TM_RESTART;
+  mj_crb(m, d);
+  mj_factorM(m, d);
+  TM_END(mjTIMER_POS_INERTIA);
+
+  TM_RESTART;
+  mj_collision(m, d);
+  TM_END(mjTIMER_POS_COLLISION);
+
+  TM_RESTART;
+  mj_makeConstraint(m, d);
+  TM_END(mjTIMER_POS_MAKE);
+
+  TM_RESTART;
+  mj_projectConstraint(m, d);
+  TM_END(mjTIMER_POS_PROJECT);
+
+  TM_END1(mjTIMER_POSITION);
+}
+
+
+
+// velocity-dependent computations
+void mj_fwdVelocity(const mjModel* m, mjData* d) {
+  TM_START;
+
+  // tendon velocity: dense or sparse
+  if (mj_isSparse(m)) {
+    mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon,
+                        d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL);
+  } else {
+    mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv);
+  }
+
+  // actuator velocity
+  mju_mulMatVec(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu, m->nv);
+
+  // standard velocity computations
+  mj_comVel(m, d);
+  mj_passive(m, d);
+  mj_referenceConstraint(m, d);
+
+  // compute qfrc_bias with abbreviated RNE (without acceleration)
+  mj_rne(m, d, 0, d->qfrc_bias);
+
+  TM_END(mjTIMER_VELOCITY);
+}
+
+
+
+// (qpos, qvel, crtl, act) => (qfrc_actuator, actuator_force, act_dot)
+void mj_fwdActuation(const mjModel* m, mjData* d) {
+  TM_START;
+  int nv = m->nv, nu = m->nu, na = m->na;
+  mjtNum gain, bias, tau;
+  mjtNum *prm, *moment = d->actuator_moment, *force = d->actuator_force;
+
+  // clear results
+  mju_zero(d->qfrc_actuator, nv);
+  if (nu) {
+    mju_zero(d->actuator_force, nu);
+  }
+
+  // check controls, set to 0 if any are bad
+  for (int i=0; i<nu; i++) {
+    if (mju_isBad(d->ctrl[i])) {
+      mj_warning(d, mjWARN_BADCTRL, i);
+      mju_zero(d->ctrl, nu);
+      break;
+    }
+  }
+
+  // disabled or no actuation: return
+  if (nu==0 || mjDISABLED(mjDSBL_ACTUATION)) {
+    return;
+  }
+
+  // force = gain .* [ctrl/act] + bias
+  for (int i=0; i<nu; i++) {
+    // clamp ctrl
+    if (m->actuator_ctrllimited[i] && !mjDISABLED(mjDSBL_CLAMPCTRL)) {
+      if (d->ctrl[i] < m->actuator_ctrlrange[2*i]) {
+        d->ctrl[i] = m->actuator_ctrlrange[2*i];
+      } else if (d->ctrl[i] > m->actuator_ctrlrange[2*i+1]) {
+        d->ctrl[i] = m->actuator_ctrlrange[2*i+1];
+      }
+    }
+
+    // extract gain info
+    prm = m->actuator_gainprm + mjNGAIN*i;
+
+    // handle according to gain type
+    switch (m->actuator_gaintype[i]) {
+    case mjGAIN_FIXED:              // fixed gain: prm = gain
+      gain = prm[0];
+      break;
+
+    case mjGAIN_MUSCLE:             // muscle gain
+      gain = mju_muscleGain(d->actuator_length[i],
+                            d->actuator_velocity[i],
+                            m->actuator_lengthrange+2*i,
+                            m->actuator_acc0[i],
+                            prm);
+      break;
+
+    default:                        // user gain
+      if (mjcb_act_gain) {
+        gain = mjcb_act_gain(m, d, i);
+      } else {
+        gain = 1;
+      }
+    }
+
+    // set force = gain .* [ctrl/act]
+    if (m->actuator_dyntype[i]==mjDYN_NONE) {
+      force[i] = gain * d->ctrl[i];
+    } else {
+      force[i] = gain * d->act[i-(nu-na)];
+    }
+
+    // extract bias info
+    prm = m->actuator_biasprm + mjNBIAS*i;
+
+    // handle according to bias type
+    switch (m->actuator_biastype[i]) {
+    case mjBIAS_NONE:               // none
+      bias = 0.0;
+      break;
+
+    case mjBIAS_AFFINE:             // affine: prm = [const, kp, kv]
+      bias = prm[0] + prm[1]*d->actuator_length[i] + prm[2]*d->actuator_velocity[i];
+      break;
+
+    case mjBIAS_MUSCLE:             // muscle passive force
+      bias =  mju_muscleBias(d->actuator_length[i],
+                             m->actuator_lengthrange+2*i,
+                             m->actuator_acc0[i],
+                             prm);
+      break;
+
+    default:                        // user bias
+      if (mjcb_act_bias) {
+        bias = mjcb_act_bias(m, d, i);
+      } else {
+        bias = 0;
+      }
+    }
+
+    // add bias
+    force[i] += bias;
+  }
+
+  // clamp actuator_force
+  for (int i=0; i<nu; i++) {
+    if (m->actuator_forcelimited[i]) {
+      if (force[i]<m->actuator_forcerange[2*i]) {
+        force[i] = m->actuator_forcerange[2*i];
+      } else if (force[i]>m->actuator_forcerange[2*i+1]) {
+        force[i] = m->actuator_forcerange[2*i+1];
+      }
+    }
+  }
+
+  // qfrc_actuator = moment' * force
+  mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv);
+
+  // act_dot for stateful actuators
+  for (int i=nu-na; i<nu; i++) {
+    // extract info
+    prm = m->actuator_dynprm + i*mjNDYN;
+    int j = i-(nu-na);
+
+    // compute act_dot according to dynamics type
+    switch (m->actuator_dyntype[i]) {
+    case mjDYN_INTEGRATOR:          // simple integrator
+      d->act_dot[j] = d->ctrl[i];
+      break;
+
+    case mjDYN_FILTER:              // linear filter: prm = tau
+      tau = mju_max(mjMINVAL, prm[0]);
+      d->act_dot[j] = (d->ctrl[i] - d->act[j]) / tau;
+      break;
+
+    case mjDYN_MUSCLE:              // muscle model: prm = (tau_act, tau_deact)
+      d->act_dot[j] = mju_muscleDynamics(d->ctrl[i], d->act[j], prm);
+      break;
+
+    default:                        // user dynamics
+      if (mjcb_act_dyn) {
+        d->act_dot[j] = mjcb_act_dyn(m, d, i);
+      } else {
+        d->act_dot[j] = 0;
+      }
+    }
+  }
+
+  TM_END(mjTIMER_ACTUATION);
+}
+
+
+
+// add up all non-constraint forces, compute qacc_smooth
+void mj_fwdAcceleration(const mjModel* m, mjData* d) {
+  TM_START;
+  mjMARKSTACK;
+  int nv = m->nv;
+
+  // qforce = sum of all non-constraint forces
+  mju_sub(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, nv);    // qfrc_bias is negative
+  mju_addTo(d->qfrc_smooth, d->qfrc_applied, nv);
+  mju_addTo(d->qfrc_smooth, d->qfrc_actuator, nv);
+  mj_xfrcAccumulate(m, d, d->qfrc_smooth);
+
+  // qacc_smooth = M \ qfr_smooth
+  mj_solveM(m, d, d->qacc_smooth, d->qfrc_smooth, 1);
+
+  mjFREESTACK;
+  TM_END(mjTIMER_ACCELERATION);
+}
+
+
+
+// warmstart/init solver
+static void warmstart(const mjModel* m, mjData* d) {
+  int nv = m->nv, nefc = d->nefc;
+
+  // warmstart with best of (qacc_warmstart, qacc_smooth)
+  if (!mjDISABLED(mjDSBL_WARMSTART)) {
+    mjMARKSTACK;
+    mjtNum* jar = mj_stackAlloc(d, nefc);
+
+    // start with qacc = qacc_warmstart
+    mju_copy(d->qacc, d->qacc_warmstart, nv);
+
+    // compute jar(qacc_warmstart)
+    mj_mulJacVec(m, d, jar, d->qacc_warmstart);
+    mju_subFrom(jar, d->efc_aref, nefc);
+
+    // update constraints, save cost(qacc_warmstart)
+    mjtNum cost_warmstart;
+    mj_constraintUpdate(m, d, jar, &cost_warmstart, 0);
+
+    // PGS
+    if (m->opt.solver==mjSOL_PGS) {
+      // cost(force_warmstart)
+      mjtNum PGS_warmstart = mju_dot(d->efc_force, d->efc_b, nefc);
+      mjtNum* ARf = mj_stackAlloc(d, nefc);
+      if (mj_isSparse(m))
+        mju_mulMatVecSparse(ARf, d->efc_AR, d->efc_force, nefc,
+                            d->efc_AR_rownnz, d->efc_AR_rowadr,
+                            d->efc_AR_colind, NULL);
+      else {
+        mju_mulMatVec(ARf, d->efc_AR, d->efc_force, nefc, nefc);
+      }
+      PGS_warmstart += 0.5*mju_dot(d->efc_force, ARf, nefc);
+
+      // use zero if better
+      if (PGS_warmstart>0) {
+        mju_zero(d->efc_force, nefc);
+        mju_zero(d->qfrc_constraint, nv);
+      }
+    }
+
+    // non-PGS
+    else {
+      // add Gauss to cost(qacc_warmstart)
+      mjtNum* Ma = mj_stackAlloc(d, nv);
+      mj_mulM(m, d, Ma, d->qacc_warmstart);
+      for (int i=0; i<nv; i++) {
+        cost_warmstart += 0.5*(Ma[i]-d->qfrc_smooth[i])*(d->qacc_warmstart[i]-d->qacc_smooth[i]);
+      }
+
+      // cost(qacc_smooth)
+      mjtNum cost_smooth;
+      mj_constraintUpdate(m, d, d->efc_b, &cost_smooth, 0);
+
+      // use qacc_smooth if better
+      if (cost_warmstart>cost_smooth) {
+        mju_copy(d->qacc, d->qacc_smooth, nv);
+      }
+    }
+
+    mjFREESTACK;
+  }
+
+  // coldstart with qacc = qacc_smooth, efc_force = 0
+  else {
+    mju_copy(d->qacc, d->qacc_smooth, nv);
+    mju_zero(d->efc_force, nefc);
+  }
+}
+
+
+
+// compute efc_b, efc_force, qfrc_constraint; update qacc
+void mj_fwdConstraint(const mjModel* m, mjData* d) {
+  TM_START;
+  int nv = m->nv, nefc = d->nefc;
+
+  // no constraints: copy unconstrained acc, clear forces, return
+  if (!nefc) {
+    mju_copy(d->qacc, d->qacc_smooth, nv);
+    mju_copy(d->qacc_warmstart, d->qacc_smooth, nv);
+    mju_zero(d->qfrc_constraint, nv);
+    d->solver_iter = 0;
+    return;
+  }
+
+  // compute efc_b = J*qacc_smooth - aref
+  mj_mulJacVec(m, d, d->efc_b, d->qacc_smooth);
+  mju_subFrom(d->efc_b, d->efc_aref, nefc);
+
+  // warmstart solver
+  warmstart(m, d);
+  d->solver_iter = 0;
+
+  // run main solver
+  switch (m->opt.solver) {
+  case mjSOL_PGS:                     // PGS
+    mj_solPGS(m, d, m->opt.iterations);
+    break;
+
+  case mjSOL_CG:                      // CG
+    mj_solCG(m, d, m->opt.iterations);
+    break;
+
+  case mjSOL_NEWTON:                  // Newton
+    mj_solNewton(m, d, m->opt.iterations);
+    break;
+
+  default:
+    mju_error_i("Unknown solver type %d", m->opt.solver);
+  }
+
+  // save result for next step warmstart
+  mju_copy(d->qacc_warmstart, d->qacc, nv);
+
+  // run noslip solver if enabled
+  if (m->opt.noslip_iterations>0) {
+    mj_solNoSlip(m, d, m->opt.noslip_iterations);
+  }
+
+  TM_END(mjTIMER_CONSTRAINT);
+}
+
+
+
+//-------------------------- integrators  ----------------------------------------------------------
+
+// Euler integrator, semi-implicit in velocity
+void mj_Euler(const mjModel* m, mjData* d) {
+  int i, nv = m->nv, nM = m->nM;
+  mjMARKSTACK;
+  mjtNum* saveM = mj_stackAlloc(d, nM);
+  mjtNum* saveLD = mj_stackAlloc(d, nM);
+  mjtNum* saveLDiagInv = mj_stackAlloc(d, nv);
+  mjtNum* saveLDiagSqrtInv = mj_stackAlloc(d, nv);
+  mjtNum* qfrc = mj_stackAlloc(d, nv);
+  mjtNum* qacc = mj_stackAlloc(d, nv);
+
+  // check for dof damping
+  for (i=0; i<nv; i++) {
+    if (m->dof_damping[i]>0) {
+      break;
+    }
+  }
+
+  // no damping: explicit velocity integration
+  if (i>=nv) {
+    mju_addToScl(d->qvel, d->qacc, m->opt.timestep, nv);
+  }
+
+  // damping: integrate implicitly
+  else {
+    // save M and factorization
+    mju_copy(saveM, d->qM, nM);
+    mju_copy(saveLD, d->qLD, nM);
+    mju_copy(saveLDiagInv, d->qLDiagInv, nv);
+    mju_copy(saveLDiagSqrtInv, d->qLDiagSqrtInv, nv);
+
+    // add hB to diagonal of M
+    for (i=0; i<nv; i++) {
+      d->qM[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i];
+    }
+
+    // factor
+    mj_factorM(m, d);
+
+    // solve
+    mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
+    mj_solveM(m, d, qacc, qfrc, 1);
+
+    // integrate velocity
+    mju_addToScl(d->qvel, qacc, m->opt.timestep, nv);
+
+    // restore M and factorization
+    mju_copy(d->qM, saveM, nM);
+    mju_copy(d->qLD, saveLD, nM);
+    mju_copy(d->qLDiagInv, saveLDiagInv, nv);
+    mju_copy(d->qLDiagSqrtInv, saveLDiagSqrtInv, nv);
+  }
+
+  // update act
+  if (m->na) {
+    mju_addToScl(d->act, d->act_dot, m->opt.timestep, m->na);
+
+    // clamp activations
+    for (i=0; i<m->na; i++) {
+      int iu = i + m->nu - m->na;
+      if (m->actuator_actlimited[iu]) {
+        mjtNum min = m->actuator_actrange[2*iu];
+        mjtNum max = m->actuator_actrange[2*iu+1];
+        if (d->act[i]<min) {
+          d->act[i] = min;
+        } else if (d->act[i]>max) {
+          d->act[i] = max;
+        }
+      }
+    }
+  }
+
+  // update qpos using new qvel
+  mj_integratePos(m, d->qpos, d->qvel, m->opt.timestep);
+
+  // advance time
+  d->time += m->opt.timestep;
+
+  mjFREESTACK;
+}
+
+
+
+// RK4 tableau
+const mjtNum RK4_A[9] = {
+  0.5,    0,      0,
+  0,      0.5,    0,
+  0,      0,      1
+};
+
+const mjtNum RK4_B[4] = {
+  1.0/6.0, 1.0/3.0, 1.0/3.0, 1.0/6.0
+};
+
+
+// Runge Kutta explicit order-N integrator
+//  (A,B) is the tableau, C is set to row_sum(A)
+void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
+  int nv = m->nv, nq = m->nq, na = m->na;
+  mjtNum h = m->opt.timestep, time = d->time;
+  mjtNum C[9], T[9], *X[10], *F[10], *dX;
+  const mjtNum* A = (N==4 ? RK4_A : 0);
+  const mjtNum* B = (N==4 ? RK4_B : 0);
+  mjMARKSTACK;
+
+  // check order
+  if (!A) {
+    mju_error("Supported RK orders: N=4");
+  }
+
+  // allocate space for intermediate solutions
+  dX = mj_stackAlloc(d, 2*nv+na);
+  for (int i=0; i<N; i++) {
+    X[i] = mj_stackAlloc(d, nq+nv+na);
+    F[i] = mj_stackAlloc(d, nv+na);
+  }
+
+  // precompute C and T;  C,T,A have size (N-1)
+  for (int i=1; i<N; i++) {
+    // C(i) = sum_j A(i,j)
+    C[i-1] = 0;
+    for (int j=0; j<i; j++) {
+      C[i-1] += A[(i-1)*(N-1)+j];
+    }
+
+    // compute T
+    T[i-1] = d->time + C[i-1]*h;
+  }
+
+  // init X[0], F[0]; mj_forward() was already called
+  mju_copy(X[0], d->qpos, nq);
+  mju_copy(X[0]+nq, d->qvel, nv);
+  mju_copy(F[0], d->qacc, nv);
+  if (na) {
+    mju_copy(X[0]+nq+nv, d->act, na);
+    mju_copy(F[0]+nv, d->act_dot, na);
+  }
+
+  // compute the remaining X[i], F[i]
+  for (int i=1; i<N; i++) {
+    // compute dX
+    mju_zero(dX, 2*nv+na);
+    for (int j=0; j<i; j++) {
+      mju_addToScl(dX, X[j]+nq, A[(i-1)*(N-1)+j], nv);
+      mju_addToScl(dX+nv, F[j], A[(i-1)*(N-1)+j], nv+na);
+    }
+
+    // compute X[i] = X[0] '+' dX
+    mju_copy(X[i], X[0], nq+nv+na);
+    mj_integratePos(m, X[i], dX, h);
+    mju_addToScl(X[i]+nq, dX+nv, h, nv+na);
+
+    // set X[i], T[i-1] in mjData
+    mju_copy(d->qpos, X[i], nq);
+    mju_copy(d->qvel, X[i]+nq, nv);
+    if (na) {
+      mju_copy(d->act, X[i]+nq+nv, na);
+    }
+    d->time = T[i-1];
+
+    // evaluate F[i]
+    mj_forwardSkip(m, d, mjSTAGE_NONE, 1);  // 1: do not recompute sensors and energy
+    mju_copy(F[i], d->qacc, nv);
+    if (na) {
+      mju_copy(F[i]+nv, d->act_dot, na);
+    }
+  }
+
+  // compute dX for final update (using B instead of A)
+  mju_zero(dX, 2*nv+na);
+  for (int j=0; j<N; j++) {
+    mju_addToScl(dX, X[j]+nq, B[j], nv);
+    mju_addToScl(dX+nv, F[j], B[j], nv+na);
+  }
+
+  // compute Xfinal
+  d->time = time + h;
+  mju_copy(d->qpos, X[0], nq+nv+na);
+  mj_integratePos(m, d->qpos, dX, h);
+  mju_addToScl(d->qvel, dX+nv, h, nv);
+  if (na) {
+    mju_addToScl(d->act, dX+2*nv, h, na);
+
+    // clamp activations
+    for (int i=0; i<m->na; i++) {
+      int iu = i + m->nu - m->na;
+      if (m->actuator_actlimited[iu]) {
+        mjtNum min = m->actuator_actrange[2*iu];
+        mjtNum max = m->actuator_actrange[2*iu+1];
+        if (d->act[i]<min) {
+          d->act[i] = min;
+        } else if (d->act[i]>max) {
+          d->act[i] = max;
+        }
+      }
+    }
+  }
+
+  mjFREESTACK;
+}
+
+
+
+//-------------------------- top-level API ---------------------------------------------------------
+
+// forward dynamics with skip; skipstage is mjtStage
+void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) {
+  TM_START;
+
+  // position-dependent
+  if (skipstage<mjSTAGE_POS) {
+    mj_fwdPosition(m, d);
+    if (!skipsensor) {
+      mj_sensorPos(m, d);
+    }
+    if (mjENABLED(mjENBL_ENERGY)) {
+      mj_energyPos(m, d);
+    }
+  }
+
+  // velocity-dependent
+  if (skipstage<mjSTAGE_VEL) {
+    mj_fwdVelocity(m, d);
+    if (!skipsensor) {
+      mj_sensorVel(m, d);
+    }
+    if (mjENABLED(mjENBL_ENERGY)) {
+      mj_energyVel(m, d);
+    }
+  }
+
+  // acceleration-dependent
+  if (mjcb_control) {
+    mjcb_control(m, d);
+  }
+  mj_fwdActuation(m, d);
+  mj_fwdAcceleration(m, d);
+  mj_fwdConstraint(m, d);
+  if (!skipsensor) {
+    mj_sensorAcc(m, d);
+  }
+
+  TM_END(mjTIMER_FORWARD);
+}
+
+
+
+// forward dynamics
+void mj_forward(const mjModel* m, mjData* d) {
+  mj_forwardSkip(m, d, mjSTAGE_NONE, 0);
+}
+
+
+
+// advance simulation using control callback
+void mj_step(const mjModel* m, mjData* d) {
+  TM_START;
+
+  // common to all integrators
+  mj_checkPos(m, d);
+  mj_checkVel(m, d);
+  mj_forward(m, d);
+  mj_checkAcc(m, d);
+
+  // compare forward and inverse solutions if enabled
+  if (mjENABLED(mjENBL_FWDINV)) {
+    mj_compareFwdInv(m, d);
+  }
+
+  // use selected integrator
+  if (m->opt.integrator==mjINT_RK4) {
+    mj_RungeKutta(m, d, 4);
+  } else {
+    mj_Euler(m, d);
+  }
+
+  TM_END(mjTIMER_STEP);
+}
+
+
+
+// advance simulation in two phases: before input is set by user
+void mj_step1(const mjModel* m, mjData* d) {
+  TM_START;
+  mj_checkPos(m, d);
+  mj_checkVel(m, d);
+  mj_fwdPosition(m, d);
+  mj_sensorPos(m, d);
+  mj_energyPos(m, d);
+  mj_fwdVelocity(m, d);
+  mj_sensorVel(m, d);
+  mj_energyVel(m, d);
+  if (mjcb_control) {
+    mjcb_control(m, d);
+  }
+  TM_END(mjTIMER_STEP);
+}
+
+
+//   >>>>   user can modify ctrl and q/xfrc_applied between step1 and step2   <<<<
+
+
+// advance simulation in two phases: after input is set by user
+void mj_step2(const mjModel* m, mjData* d) {
+  TM_START;
+  mj_fwdActuation(m, d);
+  mj_fwdAcceleration(m, d);
+  mj_fwdConstraint(m, d);
+  mj_sensorAcc(m, d);
+  mj_checkAcc(m, d);
+
+  // compare forward and inverse solutions if enabled
+  if (mjENABLED(mjENBL_FWDINV)) {
+    mj_compareFwdInv(m, d);
+  }
+
+  // integrate with Euler; ignore integrator option
+  mj_Euler(m, d);
+
+  d->timer[mjTIMER_STEP].number--;
+  TM_END(mjTIMER_STEP);
+}

src/engine/engine_forward.h

@@ -0,0 +1,80 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_FORWARD_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_FORWARD_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+// check positions, velocities, accelerations; reset if bad
+MJAPI void mj_checkPos(const mjModel* m, mjData* d);
+MJAPI void mj_checkVel(const mjModel* m, mjData* d);
+MJAPI void mj_checkAcc(const mjModel* m, mjData* d);
+
+
+//-------------------------------- top-level API ---------------------------------------------------
+
+// advance simulation: use control callback, no external force, RK4 available
+MJAPI void mj_step(const mjModel* m, mjData* d);
+
+// advance simulation in two steps: before external force/control is set by user
+MJAPI void mj_step1(const mjModel* m, mjData* d);
+
+// advance simulation in two steps: after external force/control is set by user
+MJAPI void mj_step2(const mjModel* m, mjData* d);
+
+// forward dynamics
+MJAPI void mj_forward(const mjModel* m, mjData* d);
+
+// forward dynamics with skip; skipstage is mjtStage
+MJAPI void mj_forwardSkip(const mjModel* m, mjData* d,
+                          int skipstage, int skipsensor);
+
+
+//-------------------------------- integrators -----------------------------------------------------
+
+// Euler integrator, semi-implicit in velocity
+MJAPI void mj_Euler(const mjModel* m, mjData* d);
+
+// Runge Kutta explicit order-N integrator
+MJAPI void mj_RungeKutta(const mjModel* m, mjData* d, int N);
+
+
+//-------------------------------- solver components -----------------------------------------------
+
+// computations that depend only on qpos
+MJAPI void mj_fwdPosition(const mjModel* m, mjData* d);
+
+// computations that depend only on qpos and qvel
+MJAPI void mj_fwdVelocity(const mjModel* m, mjData* d);
+
+// compute actuator force
+MJAPI void mj_fwdActuation(const mjModel* m, mjData* d);
+
+// add up all non-constraint forces, compute qacc_unc
+MJAPI void mj_fwdAcceleration(const mjModel* m, mjData* d);
+
+// forward constraint
+MJAPI void mj_fwdConstraint(const mjModel* m, mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_FORWARD_H_

src/engine/engine_inverse.c

@@ -0,0 +1,213 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_inverse.h"
+
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_collision_driver.h"
+#include "engine/engine_core_constraint.h"
+#include "engine/engine_core_smooth.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_sensor.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_sparse.h"
+
+// position-dependent computations
+void mj_invPosition(const mjModel* m, mjData* d) {
+  TM_START1;
+  TM_START;
+
+  mj_kinematics(m, d);
+  mj_comPos(m, d);
+  mj_camlight(m, d);
+  mj_tendon(m, d);
+  mj_transmission(m, d);
+  TM_END(mjTIMER_POS_KINEMATICS);
+
+  TM_RESTART;
+  mj_crb(m, d);
+  mj_factorM(m, d);
+  TM_END(mjTIMER_POS_INERTIA);
+
+  TM_RESTART;
+  mj_collision(m, d);
+  TM_END(mjTIMER_POS_COLLISION);
+
+  TM_RESTART;
+  mj_makeConstraint(m, d);
+  TM_END(mjTIMER_POS_MAKE);
+
+  TM_END1(mjTIMER_POSITION);
+}
+
+
+
+// velocity-dependent computations
+void mj_invVelocity(const mjModel* m, mjData* d) {
+  TM_START;
+
+  // tendon velocity: dense or sparse
+  if (mj_isSparse(m)) {
+    mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon,
+                        d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL);
+  } else {
+    mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv);
+  }
+
+  // actuator velocity
+  mju_mulMatVec(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu, m->nv);
+
+  // standard velocity computations
+  mj_comVel(m, d);
+  mj_passive(m, d);
+  mj_referenceConstraint(m, d);
+
+  // compute qfrc_bias with abbreviated RNE (without acceleration)
+  mj_rne(m, d, 0, d->qfrc_bias);
+
+  TM_END(mjTIMER_VELOCITY);
+}
+
+
+
+// inverse constraint solver
+void mj_invConstraint(const mjModel* m, mjData* d) {
+  TM_START;
+  int nefc = d->nefc;
+
+  // no constraints: clear, return
+  if (!nefc) {
+    mju_zero(d->qfrc_constraint, m->nv);
+    TM_END(mjTIMER_CONSTRAINT);
+    return;
+  }
+
+  mjMARKSTACK;
+  mjtNum* jar = mj_stackAlloc(d, nefc);
+
+  // compute jar = Jac*qacc - aref
+  mj_mulJacVec(m, d, jar, d->qacc);
+  mju_subFrom(jar, d->efc_aref, nefc);
+
+  // call update function
+  mj_constraintUpdate(m, d, jar, NULL, 0);
+
+  mjFREESTACK;
+  TM_END(mjTIMER_CONSTRAINT);
+}
+
+
+
+// inverse dynamics with skip; skipstage is mjtStage
+void mj_inverseSkip(const mjModel* m, mjData* d,
+                    int skipstage, int skipsensor) {
+  TM_START;
+  int nv = m->nv;
+
+  // position-dependent
+  if (skipstage<mjSTAGE_POS) {
+    mj_invPosition(m, d);
+    if (!skipsensor) {
+      mj_sensorPos(m, d);
+    }
+    if (mjENABLED(mjENBL_ENERGY)) {
+      mj_energyPos(m, d);
+    }
+  }
+
+  // velocity-dependent
+  if (skipstage<mjSTAGE_VEL) {
+    mj_invVelocity(m, d);
+    if (!skipsensor) {
+      mj_sensorVel(m, d);
+    }
+    if (mjENABLED(mjENBL_ENERGY)) {
+      mj_energyVel(m, d);
+    }
+  }
+
+  // acceleration-dependent
+  mj_invConstraint(m, d);
+  mj_rne(m, d, 1, d->qfrc_inverse);
+  if (!skipsensor) {
+    mj_sensorAcc(m, d);
+  }
+
+  // qfrc_inverse += artmature*qacc - qfrc_passive - qfrc_constraint
+  for (int i=0; i<nv; i++) {
+    d->qfrc_inverse[i] += m->dof_armature[i]*d->qacc[i]
+                          - d->qfrc_passive[i] - d->qfrc_constraint[i];
+  }
+
+  TM_END(mjTIMER_INVERSE);
+}
+
+
+
+// inverse dynamics
+void mj_inverse(const mjModel* m, mjData* d) {
+  mj_inverseSkip(m, d, mjSTAGE_NONE, 0);
+}
+
+
+
+// compare forward and inverse dynamics, without changing results of forward
+//    fwdinv[0] = norm(qfrc_constraint(forward) - qfrc_constraint(inverse))
+//    fwdinv[1] = norm(qfrc_applied(forward) - qfrc_inverse)
+void mj_compareFwdInv(const mjModel* m, mjData* d) {
+  int nv = m->nv, nefc = d->nefc;
+  mjtNum *qforce, *dif, *save_qfrc_constraint, *save_efc_force;
+  mjMARKSTACK;
+
+  // clear result, return if no constraints
+  d->solver_fwdinv[0] = d->solver_fwdinv[1] = 0;
+  if (!nefc) {
+    return;
+  }
+
+  // allocate
+  qforce = mj_stackAlloc(d, nv);
+  dif = mj_stackAlloc(d, nv);
+  save_qfrc_constraint = mj_stackAlloc(d, nv);
+  save_efc_force = mj_stackAlloc(d, nefc);
+
+  // qforce = qfrc_applied + J'*xfrc_applied + qfrc_actuator
+  //  should equal result of inverse dynamics
+  mju_add(qforce, d->qfrc_applied, d->qfrc_actuator, nv);
+  mj_xfrcAccumulate(m, d, qforce);
+
+  // save forward dynamics results that are about to be modified
+  mju_copy(save_qfrc_constraint, d->qfrc_constraint, nv);
+  mju_copy(save_efc_force, d->efc_force, nefc);
+
+  // run inverse dynamics, do not update position and velocity,
+  mj_inverseSkip(m, d, mjSTAGE_VEL, 1);  // 1: do not recompute sensors and energy
+
+  // compute statistics
+  mju_sub(dif, save_qfrc_constraint, d->qfrc_constraint, nv);
+  d->solver_fwdinv[0] = mju_norm(dif, nv);
+  mju_sub(dif, qforce, d->qfrc_inverse, nv);
+  d->solver_fwdinv[1] = mju_norm(dif, nv);
+
+  // restore forward dynamics results
+  mju_copy(d->qfrc_constraint, save_qfrc_constraint, nv);
+  mju_copy(d->efc_force, save_efc_force, nefc);
+
+  mjFREESTACK;
+}

src/engine/engine_inverse.h

@@ -0,0 +1,49 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_INVERSE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_INVERSE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// inverse dynamics
+MJAPI void mj_inverse(const mjModel* m, mjData* d);
+
+// Inverse dynamics with skip; skipstage is mjtStage.
+MJAPI void mj_inverseSkip(const mjModel* m, mjData* d,
+                          int skipstage, int skipsensor);
+
+// position-dependent computations
+MJAPI void mj_invPosition(const mjModel* m, mjData* d);
+
+// velocity-dependent computations
+MJAPI void mj_invVelocity(const mjModel* m, mjData* d);
+
+// inverse constraint solver
+MJAPI void mj_invConstraint(const mjModel* m, mjData* d);
+
+// compare forward and inverse dynamics, without changing results of forward dynamics
+MJAPI void mj_compareFwdInv(const mjModel* m, mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_INVERSE_H_

src/engine/engine_io.h

@@ -0,0 +1,108 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_IO_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_IO_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//------------------------------- initialization ---------------------------------------------------
+
+// Set default options for length range computation.
+MJAPI void mj_defaultLROpt(mjLROpt* opt);
+
+// set default solver paramters
+MJAPI void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
+
+// set options to default values
+MJAPI void mj_defaultOption(mjOption* opt);
+
+// set visual options to default values
+MJAPI void mj_defaultVisual(mjVisual* vis);
+
+// set statistics to default values; compute later in compiler
+void mj_defaultStatistic(mjStatistic* stat);
+
+
+//------------------------------- mjModel ----------------------------------------------------------
+
+// allocate mjModel
+mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int njnt,
+                      int ngeom, int nsite, int ncam, int nlight,
+                      int nmesh, int nmeshvert, int nmeshtexvert, int nmeshface, int nmeshgraph,
+                      int nskin, int nskinvert, int nskintexvert, int nskinface,
+                      int nskinbone, int nskinbonevert, int nhfield, int nhfielddata,
+                      int ntex, int ntexdata, int nmat, int npair, int nexclude,
+                      int neq, int ntendon, int nwrap, int nsensor,
+                      int nnumeric, int nnumericdata, int ntext, int ntextdata,
+                      int ntuple, int ntupledata, int nkey, int nmocap,
+                      int nuser_body, int nuser_jnt, int nuser_geom, int nuser_site, int nuser_cam,
+                      int nuser_tendon, int nuser_actuator, int nuser_sensor, int nnames);
+
+// copy mjModel; allocate new if dest is NULL
+MJAPI mjModel* mj_copyModel(mjModel* dest, const mjModel* src);
+
+// save model to binary file
+MJAPI void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buffer_sz);
+
+// load model from binary MJB file
+//  if vfs is not NULL, look up file in vfs before reading from disk
+MJAPI mjModel* mj_loadModel(const char* filename, const mjVFS* vfs);
+
+// de-allocate model
+MJAPI void mj_deleteModel(mjModel* m);
+
+// size of buffer needed to hold model
+MJAPI int mj_sizeModel(const mjModel* m);
+
+// validate reference fields in a model; return null if valid, error message otherwise
+MJAPI const char* mj_validateReferences(const mjModel* m);
+
+
+//------------------------------- mjData -----------------------------------------------------------
+
+// Allocate mjData correponding to given model.
+// If the model buffer is unallocated the initial configuration will not be set.
+MJAPI mjData* mj_makeData(const mjModel* m);
+
+// Copy mjData.
+// m is only required to contain the size fields from MJMODEL_INTS.
+MJAPI mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
+
+// set data to defaults
+MJAPI void mj_resetData(const mjModel* m, mjData* d);
+
+// set data to defaults, fill everything else with debug_value
+MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
+
+// reset data, set fields from specified keyframe
+MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
+
+// mjData stack allocate
+MJAPI mjtNum* mj_stackAlloc(mjData* d, int size);
+
+// de-allocate data
+MJAPI void mj_deleteData(mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_IO_H_

src/engine/engine_macro.h

@@ -0,0 +1,43 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_
+
+#include "engine/engine_callback.h"
+
+//-------------------------------- utility macros --------------------------------------------------
+
+// mark and free stack
+#define mjMARKSTACK int _mark = d->pstack;
+#define mjFREESTACK d->pstack = _mark;
+
+// check bitflag
+#define mjDISABLED(x) (m->opt.disableflags & (x))
+#define mjENABLED(x)  (m->opt.enableflags & (x))
+
+// max and min macros
+#define mjMAX(a,b)    (((a) > (b)) ? (a) : (b))
+#define mjMIN(a,b)    (((a) < (b)) ? (a) : (b))
+
+
+//-------------------------- timer macros ----------------------------------------------------------
+
+#define TM_START mjtNum _tm = (mjcb_time ? mjcb_time() : 0);
+#define TM_RESTART _tm = (mjcb_time ? mjcb_time() : 0);
+#define TM_END(i) {d->timer[i].duration += ((mjcb_time ? mjcb_time() : 0) - _tm); d->timer[i].number++;}
+#define TM_START1 mjtNum _tm1 = (mjcb_time ? mjcb_time() : 0);
+#define TM_END1(i) {d->timer[i].duration += ((mjcb_time ? mjcb_time() : 0) - _tm1); d->timer[i].number++;}
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_

src/engine/engine_print.c

@@ -0,0 +1,944 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_print.h"
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <string.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjxmacro.h>
+#include "engine/engine_core_constraint.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+
+#define FLOAT_FORMAT "% -9.2g"
+#define FLOAT_FORMAT_MAX_LEN 20
+#define INT_FORMAT " %d"
+#define NAME_FORMAT "%-21s"
+
+
+
+//----------------------------------- static utility functions -------------------------------------
+
+// print 2D array of mjtNum into file
+static void printArray(const char* str, int nr, int nc, const mjtNum* data, FILE* fp,
+                       const char* float_format) {
+  if (nr && nc) {
+    fprintf(fp, "%s\n  ", str);
+    for (int r=0; r<nr; r++) {
+      for (int c=0; c<nc; c++) {
+        fprintf(fp, float_format, data[c + r*nc]);
+        fprintf(fp, " ");
+      }
+      fprintf(fp, "\n  ");
+    }
+    fprintf(fp, "\n");
+  }
+}
+
+
+// print 2D array of int into file
+static void printArrayInt(const char* str, int nr, int nc, const int* data, FILE* fp) {
+  if (nr && nc) {
+    fprintf(fp, "%s\n  ", str);
+    for (int r=0; r<nr; r++) {
+      for (int c=0; c<nc; c++) {
+        fprintf(fp, "%d ", data[c + r*nc]);
+      }
+      fprintf(fp, "\n  ");
+    }
+    fprintf(fp, "\n");
+  }
+}
+
+
+
+// print sparse matrix
+static void printSparse(const char* str, const mjtNum* mat, int nr,
+                        const int* rownnz, const int* rowadr,
+                        const int* colind, FILE* fp, const char* float_format) {
+  fprintf(fp, "%s\n  ", str);
+
+  for (int r=0; r<nr; r++) {
+    for (int adr=rowadr[r]; adr<rowadr[r]+rownnz[r]; adr++) {
+      fprintf(fp, "%d: ", colind[adr]);
+      fprintf(fp, float_format, mat[adr]);
+      fprintf(fp, "  ");
+    }
+    fprintf(fp, "\n  ");
+  }
+  fprintf(fp, "\n");
+}
+
+
+
+// print vector
+static void printVector(const char* str, const mjtNum* data, int n, FILE* fp,
+                        const char* float_format) {
+  // print str
+  fprintf(fp, "%s", str);
+
+  // print data
+  for (int i=0; i<n; i++) {
+    fprintf(fp, float_format, data[i]);
+    fprintf(fp, " ");
+  }
+  fprintf(fp, "\n");
+}
+
+
+
+//------------------------------ printing functions ------------------------------------------------
+
+// return whether float_format is a valid format string for a single float
+static bool validateFloatFormat(const char* float_format) {
+  // check for nullptr;
+  if (!float_format) {
+    return false;
+  }
+
+  // example valid format string: "% -9.2g"
+  if (strnlen(float_format, FLOAT_FORMAT_MAX_LEN + 1) > FLOAT_FORMAT_MAX_LEN) {
+    mju_warning_i("Format string longer than limit of %d.", FLOAT_FORMAT_MAX_LEN);
+    return false;
+  }
+
+  int cur_idx = 0;
+  if (float_format[cur_idx] != '%') {
+    mju_warning("Format string must start with '%'.");
+    return false;
+  }
+  cur_idx++;
+
+  // flag characters. allow at most one of each flag
+  const char flag_characters[] = "-+ #0";
+  int flag_character_counts[sizeof(flag_characters)] = { 0 };
+  char* c;
+  while (c = strchr(flag_characters, float_format[cur_idx]), c != NULL) {
+    int flag_idx = (c - flag_characters)/sizeof(char);
+    flag_character_counts[flag_idx]++;
+    if (flag_character_counts[flag_idx] > 1) {
+      mju_warning("Format string contains repeated flag.");
+      return false;
+    }
+    cur_idx++;
+  }
+
+  // width. disallow *, which requires additional argument
+  while (strchr("0123456789", float_format[cur_idx]) != NULL) {
+    cur_idx++;
+  }
+
+  // precision. disallow *, which requires additional argument
+  if (float_format[cur_idx] == '.') {
+    cur_idx++;
+    while (strchr("0123456789", float_format[cur_idx]) != NULL) {
+      cur_idx++;
+    }
+  }
+
+  // length
+  if (float_format[cur_idx] == 'L') {
+    cur_idx++;
+  }
+
+  // specifier must be a valid float format
+  if (strchr("fgGeE", float_format[cur_idx]) == NULL) {
+    mju_warning("Format string specifier must be one of \"fgGeE\".");
+    return false;
+  }
+  cur_idx++;
+
+  if (float_format[cur_idx] == '\0') {
+    return true;
+  } else {
+    mju_warning_s("Unable to match format string %s with expected pattern for a single float.",
+                  float_format);
+    return false;
+  }
+}
+
+
+// Clang sometimes goes OOM when the -Wuninitialized warning is enabled for this function
+#ifdef __clang__
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wuninitialized"
+#endif
+
+
+// print mjModel to text file, specifying format. float_format must be a
+// valid printf-style format string for a single float value
+void mj_printFormattedModel(const mjModel* m, const char* filename, const char* float_format) {
+  // get file
+  FILE* fp;
+  if (filename) {
+    fp = fopen(filename, "wt");
+  } else {
+    fp = stdout;
+  }
+
+  // check for nullptr
+  if (!fp) {
+    mju_warning_s("Could not open file '%s' for writing mjModel", filename);
+    return;
+  }
+
+  // validate format string
+  if (!validateFloatFormat(float_format)) {
+    mju_warning("WARNING: Received invalid float_format. Using default instead.");
+    float_format = FLOAT_FORMAT;
+  }
+
+  // compute total body mass
+  mjtNum totalmass = 0;
+  for (int i=0; i<m->nbody; i++) {
+    totalmass += m->body_mass[i];
+  }
+
+  // software version and model name
+  fprintf(fp, "MuJoCo version %s\n", mj_versionString());
+  fprintf(fp, "model name     %s\n\n", m->names);
+
+  // sizes
+#define X( name )                                 \
+  if(m->name) {                                   \
+    fprintf(fp, NAME_FORMAT, #name);              \
+    fprintf(fp, INT_FORMAT "\n", m->name);        \
+  }
+
+  MJMODEL_INTS
+#undef X
+  fprintf(fp, "\n");
+
+  // scalar options
+#define X( type, name )                           \
+  fprintf(fp, NAME_FORMAT, #name);                \
+  fprintf(fp, float_format, m->opt.name);         \
+  fprintf(fp, "\n");
+
+  MJOPTION_FLOATS
+#undef X
+
+#define X( type, name )                           \
+  fprintf(fp, NAME_FORMAT, #name);                \
+  fprintf(fp, INT_FORMAT "\n", m->opt.name);
+
+  MJOPTION_INTS
+#undef X
+
+  // vector options
+#define X( name, sz )                             \
+  fprintf(fp, NAME_FORMAT, #name);                \
+  for (int i=0; i < sz; i++) {                    \
+    fprintf(fp, float_format, m->opt.name[i]);    \
+    fprintf(fp, " ");                             \
+  }                                               \
+  fprintf(fp, "\n");
+
+  MJOPTION_VECTORS
+#undef X
+  fprintf(fp, "\n");
+
+  // total mass
+  fprintf(fp, NAME_FORMAT, "totalmass");
+  fprintf(fp, float_format, totalmass);
+  fprintf(fp, "\n\n");
+
+  // statistics
+  fprintf(fp, NAME_FORMAT, "meaninertia");
+  fprintf(fp, float_format, m->stat.meaninertia);
+  fprintf(fp, "\n");
+  fprintf(fp, NAME_FORMAT, "meanmass");
+  fprintf(fp, float_format, m->stat.meanmass);
+  fprintf(fp, "\n");
+  fprintf(fp, NAME_FORMAT, "meansize");
+  fprintf(fp, float_format, m->stat.meansize);
+  fprintf(fp, "\n");
+  fprintf(fp, NAME_FORMAT, "extent");
+  fprintf(fp, float_format, m->stat.extent);
+  fprintf(fp, "\n");
+  fprintf(fp, NAME_FORMAT, "center");
+  fprintf(fp, float_format, m->stat.center[0]);
+  fprintf(fp, float_format, m->stat.center[1]);
+  fprintf(fp, float_format, m->stat.center[2]);
+  fprintf(fp, "\n\n");
+
+  // qpos0
+  fprintf(fp, NAME_FORMAT, "qpos0");
+  for (int i=0; i<m->nq; i++) {
+    fprintf(fp, float_format, m->qpos0[i]);
+    fprintf(fp, " ");
+  }
+  fprintf(fp, "\n\n");
+
+  // qpos_spring
+  fprintf(fp, NAME_FORMAT, "qpos_spring");
+  for (int i=0; i<m->nq; i++) {
+    fprintf(fp, float_format, m->qpos_spring[i]);
+    fprintf(fp, " ");
+  }
+  fprintf(fp, "\n\n");
+
+  // values used by MJMODEL_POINTERS macro
+  MJMODEL_POINTERS_PREAMBLE(m)
+
+  // object_class points to the integer size identifying the class of arrays currently being printed
+  //   used to organise the printout into category groups
+  //   note that comparison is based on the integer address, not its value
+  const int* object_class;
+
+#define X( type, name, num, sz )                                              \
+  if (&m->num == object_class && (strncmp(#name, "name_", 5)!=0) && sz) {     \
+    fprintf(fp, "  ");                                                        \
+    fprintf(fp, NAME_FORMAT, #name);                                          \
+    for (int j=0; j < sz; j++) {                                              \
+       ((strcmp(#type, "mjtNum") == 0) || (strcmp(#type, "float") == 0)) ?    \
+          (fprintf(fp, float_format, (mjtNum)m->name[sz*i+j]),                \
+           fprintf(fp, " ")) :                                                \
+          fprintf(fp, INT_FORMAT " ", (int)m->name[sz*i+j]);                  \
+    }                                                                         \
+    fprintf(fp, "\n");                                                        \
+  }
+
+  // bodies
+  for (int i=0; i<m->nbody; i++) {
+    fprintf(fp, "\nBODY %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_bodyadr[i]);
+    object_class = &m->nbody;
+    MJMODEL_POINTERS
+  }
+  if (m->nbody) fprintf(fp, "\n");
+
+  // joints
+  for (int i=0; i<m->njnt; i++) {
+    fprintf(fp, "\nJOINT %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_jntadr[i]);
+    object_class = &m->njnt;
+    MJMODEL_POINTERS
+  }
+  if (m->njnt) fprintf(fp, "\n");
+
+  // dofs
+  for (int i=0; i<m->nv; i++) {
+    fprintf(fp, "\nDOF %d:\n", i);
+    object_class = &m->nv;
+    MJMODEL_POINTERS
+  }
+  if (m->nv) fprintf(fp, "\n");
+
+  // geoms
+  for (int i=0; i<m->ngeom; i++) {
+    fprintf(fp, "\nGEOM %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_geomadr[i]);
+    object_class = &m->ngeom;
+    MJMODEL_POINTERS
+  }
+  if (m->ngeom) fprintf(fp, "\n");
+
+  // sites
+  for (int i=0; i<m->nsite; i++) {
+    fprintf(fp, "\nSITE %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_siteadr[i]);
+    object_class = &m->nsite;
+    MJMODEL_POINTERS
+  }
+  if (m->nsite) fprintf(fp, "\n");
+
+  // cameras
+  for (int i=0; i<m->ncam; i++) {
+    fprintf(fp, "\nCAMERA %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_camadr[i]);
+    object_class = &m->ncam;
+    MJMODEL_POINTERS
+  }
+  if (m->ncam) fprintf(fp, "\n");
+
+  // lights
+  for (int i=0; i<m->nlight; i++) {
+    fprintf(fp, "\nLIGHT %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_lightadr[i]);
+    object_class = &m->nlight;
+    MJMODEL_POINTERS
+  }
+  if (m->nlight) fprintf(fp, "\n");
+
+  // meshes
+  for (int i=0; i<m->nmesh; i++) {
+    fprintf(fp, "\nMESH %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_meshadr[i]);
+    object_class = &m->nmesh;
+    MJMODEL_POINTERS
+    if (m->mesh_graphadr[i]>=0) {
+      fprintf(fp, "  " NAME_FORMAT, "qhull face");
+      fprintf(fp, " %d\n", m->mesh_graph[m->mesh_graphadr[i]+1]);
+      fprintf(fp, "  " NAME_FORMAT, "qhull vert");
+      fprintf(fp, " %d\n", m->mesh_graph[m->mesh_graphadr[i]]);
+    }
+  }
+  if (m->nmesh) fprintf(fp, "\n");
+
+  // skins
+  for (int i=0; i<m->nskin; i++) {
+    fprintf(fp, "\nSKIN %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_skinadr[i]);
+    object_class = &m->nskin;
+    MJMODEL_POINTERS
+  }
+  if (m->nskin) fprintf(fp, "\n");
+
+  // hfields
+  for (int i=0; i<m->nhfield; i++) {
+    fprintf(fp, "\nHEIGHTFIELD %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, "  %s\n", m->names + m->name_hfieldadr[i]);
+    object_class = &m->nhfield;
+    MJMODEL_POINTERS
+  }
+  if (m->nhfield) fprintf(fp, "\n");
+
+  // textures
+  for (int i=0; i<m->ntex; i++) {
+    fprintf(fp, "\nTEXTURE %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_texadr[i]);
+    object_class = &m->ntex;
+    MJMODEL_POINTERS
+  }
+  if (m->ntex) fprintf(fp, "\n");
+
+  // materials
+  for (int i=0; i<m->nmat; i++) {
+    fprintf(fp, "\nMATERIAL %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_matadr[i]);
+    object_class = &m->nmat;
+    MJMODEL_POINTERS
+  }
+  if (m->nmat) fprintf(fp, "\n");
+
+  // pairs
+  for (int i=0; i<m->npair; i++) {
+    fprintf(fp, "\nPAIR %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_pairadr[i]);
+    object_class = &m->npair;
+    MJMODEL_POINTERS
+  }
+  if (m->npair) fprintf(fp, "\n");
+
+  // excludes
+  for (int i=0; i<m->nexclude; i++) {
+    fprintf(fp, "\nEXCLUDE %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_excludeadr[i]);
+    object_class = &m->nexclude;
+    MJMODEL_POINTERS
+  }
+  if (m->nexclude) fprintf(fp, "\n");
+
+  // equality constraints
+  for (int i=0; i<m->neq; i++) {
+    fprintf(fp, "\nEQUALITY %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_eqadr[i]);
+    object_class = &m->neq;
+    MJMODEL_POINTERS
+  }
+  if (m->neq) fprintf(fp, "\n");
+
+  // tendons
+  for (int i=0; i<m->ntendon; i++) {
+    fprintf(fp, "\nTENDON %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_tendonadr[i]);
+    object_class = &m->ntendon;
+    MJMODEL_POINTERS
+    fprintf(fp, "  path         \n");
+    for (int j=0; j<m->tendon_num[i]; j++) {
+      int k = m->tendon_adr[i]+j;
+      fprintf(fp, "    %d %d ", m->wrap_type[k], m->wrap_objid[k]);
+      fprintf(fp, float_format, m->wrap_prm[k]);
+      fprintf(fp, "\n");
+    }
+    fprintf(fp, "\n");
+  }
+  if (m->ntendon) fprintf(fp, "\n");
+
+  // actuators
+  for (int i=0; i<m->nu; i++) {
+    fprintf(fp, "\nACTUATOR %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_actuatoradr[i]);
+    object_class = &m->nu;
+    MJMODEL_POINTERS
+  }
+  if (m->nu) fprintf(fp, "\n");
+
+  // sensors
+  for (int i=0; i<m->nsensor; i++) {
+    fprintf(fp, "\nSENSOR %d:\n", i);
+    fprintf(fp, "  " NAME_FORMAT, "name");
+    fprintf(fp, " %s\n", m->names + m->name_sensoradr[i]);
+    object_class = &m->nsensor;
+    MJMODEL_POINTERS
+  }
+  if (m->nsensor) fprintf(fp, "\n");
+
+  // custom numeric parameters
+  for (int i=0; i<m->nnumeric; i++) {
+    fprintf(fp, "\nNUMERIC %d:\n", i);
+    fprintf(fp, "  name         %s\n", m->names + m->name_numericadr[i]);
+    fprintf(fp, "  size         %d\n", m->numeric_size[i]);
+    fprintf(fp, "  value       ");
+    for (int j=0; j<m->numeric_size[i]; j++) {
+      fprintf(fp, float_format, m->numeric_data[m->numeric_adr[i]+j]);
+    }
+    fprintf(fp, "\n");
+  }
+  if (m->nnumeric) fprintf(fp, "\n");
+
+  // custom text parameters
+  for (int i=0; i<m->ntext; i++) {
+    fprintf(fp, "\nTEXT %d:\n", i);
+    fprintf(fp, "  name         %s\n", m->names + m->name_textadr[i]);
+    fprintf(fp, "  size         %d\n", m->text_size[i]);
+    fprintf(fp, "  value        %s\n", m->text_data + m->text_adr[i]);
+  }
+  if (m->ntext) fprintf(fp, "\n");
+
+  // custom tuple parameters
+  for (int i=0; i<m->ntuple; i++) {
+    fprintf(fp, "\nTUPLE %d:\n", i);
+    fprintf(fp, "  name         %s\n", m->names + m->name_tupleadr[i]);
+    fprintf(fp, "  size         %d\n", m->tuple_size[i]);
+    fprintf(fp, "  elements\n");
+    for (int j=m->tuple_adr[i]; j<m->tuple_adr[i]+m->tuple_size[i]; j++) {
+      fprintf(fp, "       %s %d, prm = ",
+              mju_type2Str(m->tuple_objtype[j]), m->tuple_objid[j]);
+      fprintf(fp, float_format, m->tuple_objprm[j]);
+      fprintf(fp, "\n");
+    }
+  }
+  if (m->ntuple) fprintf(fp, "\n");
+
+  // keyframes (only if different from default)
+  for (int i=0; i<m->nkey; i++) {
+    // print name
+    if (m->names[m->name_keyadr[i]]) {
+      fprintf(fp, "key_name%d    %s\n", i, m->names + m->name_keyadr[i]);
+    }
+
+    // print time if non-0
+    if (m->key_time[i]!=0) {
+      fprintf(fp, "key_time%d    %.4f\n", i, m->key_time[i]);
+    }
+
+    // check qpos for difference
+    int k = 0;
+    for (int j=0; j<m->nq; j++)
+      if (m->qpos0[j] != m->key_qpos[i*m->nq + j]) {
+        k = 1;
+      }
+
+    // print if different
+    if (k==1) {
+      fprintf(fp, "key_qpos%d   ", i);
+      for (int j=0; j<m->nq; j++) {
+        fprintf(fp, float_format, m->key_qpos[i*m->nq + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+    // check qvel for nonzero
+    for (int j=0; j<m->nv; j++)
+      if (m->key_qvel[i*m->nv + j]) {
+        k = 2;
+      }
+
+    // print if nozero
+    if (k==2) {
+      fprintf(fp, "key_qvel%d   ", i);
+      for (int j=0; j<m->nv; j++) {
+        fprintf(fp, float_format, m->key_qvel[i*m->nv + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+    // check act for nonzero
+    for (int j=0; j<m->na; j++)
+      if (m->key_act[i*m->na + j]) {
+        k = 3;
+      }
+
+    // print if nonzero
+    if (k==3) {
+      fprintf(fp, "key_act%d   ", i);
+      for (int j=0; j<m->na; j++) {
+        fprintf(fp, float_format, m->key_act[i*m->na + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+    // check mpos for difference
+    if (m->nmocap) {
+      for (int j=0; j<m->nbody; j++) {
+        if (m->body_mocapid[j]>=0) {
+          int id = m->body_mocapid[j];
+          if (m->body_pos[3*j]   != m->key_mpos[i*3*m->nmocap + 3*id] ||
+              m->body_pos[3*j+1] != m->key_mpos[i*3*m->nmocap + 3*id+1] ||
+              m->body_pos[3*j+2] != m->key_mpos[i*3*m->nmocap + 3*id+2]) {
+            k = 4;
+            break;
+          }
+        }
+      }
+    }
+
+    // print if nonzero
+    if (k==4) {
+      fprintf(fp, "key_mpos%d   ", i);
+      for (int j=0; j<3*m->nmocap; j++) {
+        fprintf(fp, float_format, m->key_mpos[i*3*m->nmocap + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+    // check mquat for difference
+    if (m->nmocap) {
+      for (int j=0; j<m->nbody; j++) {
+        if (m->body_mocapid[j]>=0) {
+          int id = m->body_mocapid[j];
+          if (m->body_quat[4*j]   != m->key_mquat[i*4*m->nmocap + 4*id] ||
+              m->body_quat[4*j+1] != m->key_mquat[i*4*m->nmocap + 4*id+1] ||
+              m->body_quat[4*j+2] != m->key_mquat[i*4*m->nmocap + 4*id+2] ||
+              m->body_quat[4*j+3] != m->key_mquat[i*4*m->nmocap + 4*id+3]) {
+            k = 5;
+            break;
+          }
+        }
+      }
+    }
+
+    // print if nonzero
+    if (k==5) {
+      fprintf(fp, "key_mquat%d   ", i);
+      for (int j=0; j<4*m->nmocap; j++) {
+        fprintf(fp, float_format, m->key_mquat[i*4*m->nmocap + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+    // new line if any data was written
+    if (k) {
+      fprintf(fp, "\n");
+    }
+  }
+
+#undef X
+
+  if (filename) {
+    fclose(fp);
+  }
+}
+
+
+// print mjModel to text file
+void mj_printModel(const mjModel* m, const char* filename) {
+  mj_printFormattedModel(m, filename, FLOAT_FORMAT);
+}
+
+
+// print mjModel to text file, specifying format. float_format must be a
+// valid printf-style format string for a single float value
+void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
+                           const char* float_format) {
+  mjtNum *M;
+  mjMARKSTACK;
+
+  // check format string
+  if (!validateFloatFormat(float_format)) {
+    mju_warning("WARNING: Received invalid float_format. Using default instead.");
+    float_format = FLOAT_FORMAT;
+  }
+
+  // stack in use, SHOULD NOT OCCUR
+  if (d->pstack) {
+    mju_error("Attempting to print mjData when stack is in use");
+  }
+
+  // get file
+  FILE* fp;
+  if (filename) {
+    fp = fopen(filename, "wt");
+  } else {
+    fp = stdout;
+  }
+
+  // check for nullptr
+  if (!fp) {
+    mju_warning_s("Could not open file '%s' for writing mjModel", filename);
+    mjFREESTACK;
+    return;
+  }
+
+  // allocate full inertia
+  M = mj_stackAlloc(d, m->nv*m->nv);
+
+  // ---------------------------------- print mjData fields
+
+  fprintf(fp, "SIZES\n");
+#define X( type, name )                          \
+  if(strcmp(#type, "int")==0) {                  \
+    fprintf(fp, "  ");                           \
+    fprintf(fp, NAME_FORMAT, #name);             \
+    fprintf(fp, INT_FORMAT "\n", (int)d->name);  \
+  }
+
+  MJDATA_SCALAR
+#undef X
+  fprintf(fp, "\n");
+
+  // WARNING
+  int active_warnings = 0;
+  for (int i=0; i<mjNWARNING; i++) {
+    active_warnings += d->warning[i].number;
+  }
+  if (active_warnings) {
+    fprintf(fp, "WARNING\n");
+    for (int i=0; i<mjNWARNING; i++)
+      if (d->warning[i].number)
+        fprintf(fp, "    %d:  lastinfo = %d   number = %d\n",
+                i, d->warning[i].lastinfo, d->warning[i].number);
+    fprintf(fp, "\n");
+  }
+
+  // TIMER
+  mjtNum active_timers = 0;
+  for (int i=0; i<mjNTIMER; i++) {
+    active_timers += d->timer[i].duration;
+  }
+  if (active_timers) {
+    fprintf(fp, "TIMER\n");
+    for (int i=0; i<mjNTIMER; i++) {
+      fprintf(fp, "    %d:  duration = ",i);
+      fprintf(fp, float_format, d->timer[i].duration);
+      fprintf(fp, "   number = %d\n", d->timer[i].number);
+    }
+    fprintf(fp, "\n");
+  }
+
+  // SOLVER STAT
+  if (d->solver_iter) {
+    fprintf(fp, "SOLVER STAT\n");
+    fprintf(fp, "  solver_iter = %d\n", d->solver_iter);
+    fprintf(fp, "  solver_nnz = %d\n",  d->solver_nnz);
+    for (int i=0; i<mjMIN(mjNSOLVER, d->solver_iter); i++) {
+      fprintf(fp, "    %d:  improvement = ", i);
+      fprintf(fp, float_format, d->solver[i].improvement);
+      fprintf(fp, "  gradient = ");
+      fprintf(fp, float_format, d->solver[i].gradient);
+      fprintf(fp, "  lineslope = ");
+      fprintf(fp, float_format, d->solver[i].lineslope);
+      fprintf(fp, "\n");
+      fprintf(fp, "        nactive = %d   nchange = %d   neval = %d   nupdate = %d\n",
+              d->solver[i].nactive, d->solver[i].nchange,
+              d->solver[i].neval, d->solver[i].nupdate);
+    }
+    printVector("solver_fwdinv = ", d->solver_fwdinv, 2, fp, float_format);
+    fprintf(fp, "\n");
+  }
+
+  printVector("ENERGY = ", d->energy, 2, fp, float_format);
+  fprintf(fp, "\n");
+
+  fprintf(fp, "TIME = ");
+  fprintf(fp, float_format, d->time);
+  fprintf(fp, "\n\n");
+
+  printArray("QPOS", m->nq, 1, d->qpos, fp, float_format);
+  printArray("QVEL", m->nv, 1, d->qvel, fp, float_format);
+  printArray("ACT", m->na, 1, d->act, fp, float_format);
+  printArray("QACC_WARMSTART", m->nv, 1, d->qacc_warmstart, fp, float_format);
+  printArray("CTRL", m->nu, 1, d->ctrl, fp, float_format);
+  printArray("QFRC_APPLIED", m->nq, 1, d->qfrc_applied, fp, float_format);
+  printArray("XFRC_APPLIED", m->nbody, 6, d->xfrc_applied, fp, float_format);
+  printArray("MOCAP_POS", m->nmocap, 3, d->mocap_pos, fp, float_format);
+  printArray("MOCAP_QUAT", m->nmocap, 4, d->mocap_quat, fp, float_format);
+  printArray("QACC", m->nv, 1, d->qacc, fp, float_format);
+  printArray("ACT_DOT", m->na, 1, d->act_dot, fp, float_format);
+  printArray("USERDATA", m->nuserdata, 1, d->userdata, fp, float_format);
+  printArray("SENSOR", m->nsensordata, 1, d->sensordata, fp, float_format);
+
+  printArray("XPOS", m->nbody, 3, d->xpos, fp, float_format);
+  printArray("XQUAT", m->nbody, 4, d->xquat, fp, float_format);
+  printArray("XMAT", m->nbody, 9, d->xmat, fp, float_format);
+  printArray("XIPOS", m->nbody, 3, d->xipos, fp, float_format);
+  printArray("XIMAT", m->nbody, 9, d->ximat, fp, float_format);
+  printArray("XANCHOR", m->njnt, 3, d->xanchor, fp, float_format);
+  printArray("XAXIS", m->njnt, 3, d->xaxis, fp, float_format);
+  printArray("GEOM_XPOS", m->ngeom, 3, d->geom_xpos, fp, float_format);
+  printArray("GEOM_XMAT", m->ngeom, 9, d->geom_xmat, fp, float_format);
+  printArray("SITE_XPOS", m->nsite, 3, d->site_xpos, fp, float_format);
+  printArray("SITE_XMAT", m->nsite, 9, d->site_xmat, fp, float_format);
+  printArray("CAM_XPOS", m->ncam, 3, d->cam_xpos, fp, float_format);
+  printArray("CAM_XMAT", m->ncam, 9, d->cam_xmat, fp, float_format);
+  printArray("LIGHT_XPOS", m->nlight, 3, d->light_xpos, fp, float_format);
+  printArray("LIGHT_XDIR", m->nlight, 3, d->light_xdir, fp, float_format);
+
+  printArray("SUBTREE_COM", m->nbody, 3, d->subtree_com, fp, float_format);
+  printArray("CDOF", m->nv, 6, d->cdof, fp, float_format);
+  printArray("CINERT", m->nbody, 10, d->cinert, fp, float_format);
+
+  printArray("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format);
+  if (!mj_isSparse(m)) {
+    printArray("TEN_MOMENT", m->ntendon, m->nv, d->ten_J, fp, float_format);
+  } else {
+    printArrayInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp);
+    printArrayInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp);
+    printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz,
+                         d->ten_J_rowadr, d->ten_J_colind, fp, float_format);
+  }
+  for (int i=0; i<m->ntendon; i++) {
+    fprintf(fp, "TENDON %d: %d wrap points\n", i, d->ten_wrapnum[i]);
+    for (int j=0; j<d->ten_wrapnum[i]; j++) {
+      fprintf(fp, "    %d:  ", d->wrap_obj[d->ten_wrapadr[i]+j]);
+      printVector("", d->wrap_xpos+3*(d->ten_wrapadr[i]+j), 3, fp, float_format);
+    }
+    fprintf(fp, "\n");
+  }
+
+  printArray("ACTUATOR_LENGTH", m->nu, 1, d->actuator_length, fp, float_format);
+  printArray("ACTUATOR_MOMENT", m->nu, m->nv, d->actuator_moment, fp, float_format);
+  printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
+
+  // construct and print full M matrix
+  mj_fullM(m, M, d->qM);
+  printArray("QM", m->nv, m->nv, M, fp, float_format);
+
+  // construct and print full LD matrix
+  mj_fullM(m, M, d->qLD);
+  printArray("QLD", m->nv, m->nv, M, fp, float_format);
+
+  printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
+  printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
+
+  // contact
+  fprintf(fp, "CONTACT\n");
+  for (int i=0; i<d->ncon; i++) {
+    fprintf(fp, "  %d:\n     dim           %d\n     geom          %d %d\n",
+            i, d->contact[i].dim, d->contact[i].geom1, d->contact[i].geom2);
+    fprintf(fp, "     exclude       %d\n     efc_address   %d\n",
+            d->contact[i].exclude, d->contact[i].efc_address);
+    printVector("     solref       ", d->contact[i].solref, mjNREF, fp, float_format);
+    printVector("     solimp       ", d->contact[i].solimp, mjNIMP, fp, float_format);
+    printVector("     dist         ", &d->contact[i].dist, 1, fp, float_format);
+    printVector("     includemargin", &d->contact[i].includemargin, 1, fp, float_format);
+    printVector("     pos          ", d->contact[i].pos, 3, fp, float_format);
+    printVector("     frame        ", d->contact[i].frame, 9, fp, float_format);
+    printVector("     friction     ", d->contact[i].friction, 5, fp, float_format);
+    printVector("     mu           ", &d->contact[i].mu, 1, fp, float_format);
+  }
+  if (d->ncon) fprintf(fp, "\n");
+
+  printArrayInt("EFC_TYPE", d->nefc, 1, d->efc_type, fp);
+  printArrayInt("EFC_ID", d->nefc, 1, d->efc_id, fp);
+
+  if (!mj_isSparse(m)) {
+    printArray("EFC_J", d->nefc, m->nv, d->efc_J, fp, float_format);
+    printArray("EFC_AR", d->nefc, d->nefc, d->efc_AR, fp, float_format);
+  } else {
+    printArrayInt("EFC_J_ROWNNZ", d->nefc, 1, d->efc_J_rownnz, fp);
+    printArrayInt("EFC_J_ROWADR", d->nefc, 1, d->efc_J_rowadr, fp);
+    printSparse("EFC_J", d->efc_J, d->nefc, d->efc_J_rownnz,
+                d->efc_J_rowadr, d->efc_J_colind, fp, float_format);
+
+    printArrayInt("EFC_AR_ROWNNZ", d->nefc, 1, d->efc_AR_rownnz, fp);
+    printArrayInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp);
+    printSparse("EFC_AR", d->efc_AR, d->nefc, d->efc_AR_rownnz,
+                d->efc_AR_rowadr, d->efc_AR_colind, fp, float_format);
+  }
+
+  printArray("EFC_POS", d->nefc, 1, d->efc_pos, fp, float_format);
+  printArray("EFC_MARGIN", d->nefc, 1, d->efc_margin, fp, float_format);
+  printArray("EFC_FRICTIONLOSS", d->nefc, 1, d->efc_frictionloss, fp, float_format);
+  printArray("EFC_DIAGAPPROX", d->nefc, 1, d->efc_diagApprox, fp, float_format);
+  printArray("EFC_KBIP", d->nefc, 4, d->efc_KBIP, fp, float_format);
+  printArray("EFC_D", d->nefc, 1, d->efc_D, fp, float_format);
+  printArray("EFC_R", d->nefc, 1, d->efc_R, fp, float_format);
+
+  printArray("TEN_VELOCITY", m->ntendon, 1, d->ten_velocity, fp, float_format);
+  printArray("ACTUATOR_VELOCITY", m->nu, 1, d->actuator_velocity, fp, float_format);
+
+  printArray("CVEL", m->nbody, 6, d->cvel, fp, float_format);
+  printArray("CDOF_DOT", m->nv, 6, d->cdof_dot, fp, float_format);
+
+  printArray("QFRC_BIAS", m->nv, 1, d->qfrc_bias, fp, float_format);
+
+  printArray("QFRC_PASSIVE", m->nv, 1, d->qfrc_passive, fp, float_format);
+
+  printArray("EFC_VEL", d->nefc, 1, d->efc_vel, fp, float_format);
+  printArray("EFC_AREF", d->nefc, 1, d->efc_aref, fp, float_format);
+
+  printArray("SUBTREE_LINVEL", m->nbody, 3, d->subtree_linvel, fp, float_format);
+  printArray("SUBTREE_ANGMOM", m->nbody, 3, d->subtree_angmom, fp, float_format);
+
+  printArray("ACTUATOR_FORCE", m->nu, 1, d->actuator_force, fp, float_format);
+  printArray("QFRC_ACTUATOR", m->nv, 1, d->qfrc_actuator, fp, float_format);
+
+  printArray("QFRC_SMOOTH", m->nv, 1, d->qfrc_smooth, fp, float_format);
+  printArray("QACC_SMOOTH", m->nv, 1, d->qacc_smooth, fp, float_format);
+
+  printArray("EFC_B", d->nefc, 1, d->efc_b, fp, float_format);
+  printArray("EFC_FORCE", d->nefc, 1, d->efc_force, fp, float_format);
+  printArrayInt("EFC_STATE", d->nefc, 1, d->efc_state, fp);
+  printArray("QFRC_CONSTRAINT", m->nv, 1, d->qfrc_constraint, fp, float_format);
+
+  printArray("QFRC_INVERSE", m->nv, 1, d->qfrc_inverse, fp, float_format);
+
+  printArray("CACC", m->nbody, 6, d->cacc, fp, float_format);
+  printArray("CFRC_INT", m->nbody, 6, d->cfrc_int, fp, float_format);
+  printArray("CFRC_EXT", m->nbody, 6, d->cfrc_ext, fp, float_format);
+
+  if (filename) {
+    fclose(fp);
+  }
+
+  mjFREESTACK;
+}
+
+
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif
+
+
+// print mjData to text file
+void mj_printData(const mjModel* m, mjData* d, const char* filename) {
+  mj_printFormattedData(m, d, filename, FLOAT_FORMAT);
+}

src/engine/engine_print.h

@@ -0,0 +1,47 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_PRINT_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_PRINT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// print mjModel to text file, specifying format
+// float_format must be a valid printf-style format string for a single float value
+MJAPI void mj_printFormattedModel(const mjModel* m, const char* filename,
+                                  const char* float_format);
+
+// print model and option to text file
+MJAPI void mj_printModel(const mjModel* m, const char* filename);
+
+
+// print mjData to text file, specifying format
+// float_format must be a valid printf-style format string for a single float value
+MJAPI void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
+                                 const char* float_format);
+
+// print data to text file
+MJAPI void mj_printData(const mjModel* m, mjData* d, const char* filename);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_PRINT_H_

src/engine/engine_ray.c

@@ -0,0 +1,821 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//---------------------------------//
+
+#include "engine/engine_ray.h"
+
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+#include "engine/engine_macro.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_spatial.h"
+
+//---------------------------- utility functions ---------------------------------------------------
+
+// map ray to local geom frame
+static void ray_map(const mjtNum* pos, const mjtNum* mat, const mjtNum* pnt, const mjtNum* vec,
+                    mjtNum* lpnt, mjtNum* lvec) {
+  const mjtNum dif[3] = {pnt[0]-pos[0], pnt[1]-pos[1], pnt[2]-pos[2]};
+
+  // lpnt = mat' * dif
+  lpnt[0] = mat[0]*dif[0] + mat[3]*dif[1] + mat[6]*dif[2];
+  lpnt[1] = mat[1]*dif[0] + mat[4]*dif[1] + mat[7]*dif[2];
+  lpnt[2] = mat[2]*dif[0] + mat[5]*dif[1] + mat[8]*dif[2];
+
+  // lvec = mat' * vec
+  lvec[0] = mat[0]*vec[0] + mat[3]*vec[1] + mat[6]*vec[2];
+  lvec[1] = mat[1]*vec[0] + mat[4]*vec[1] + mat[7]*vec[2];
+  lvec[2] = mat[2]*vec[0] + mat[5]*vec[1] + mat[8]*vec[2];
+}
+
+
+
+// eliminate geom
+static int ray_eliminate(const mjModel* m, const mjData* d, int geomid,
+                         const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude) {
+  // body exclusion
+  if (m->geom_bodyid[geomid]==bodyexclude) {
+    return 1;
+  }
+
+  // invisible geom exclusion
+  if (m->geom_matid[geomid]<0 && m->geom_rgba[4*geomid+3]==0) {
+    return 1;
+  }
+
+  // invisible material exclusion
+  if (m->geom_matid[geomid]>=0 && m->mat_rgba[4*m->geom_matid[geomid]+3]==0) {
+    return 1;
+  }
+
+  // static exclusion
+  if (!flg_static && m->geom_bodyid[geomid]==0) {
+    return 1;
+  }
+
+  // plane and hfield inclusion
+  if (m->geom_type[geomid]==mjGEOM_PLANE || m->geom_type[geomid]==mjGEOM_HFIELD) {
+    return 0;
+  }
+
+  // no geomgroup inclusion
+  if (!geomgroup) {
+    return 0;
+  }
+
+  // group inclusion/exclusion
+  int groupid = mjMIN(mjNGROUP-1, mjMAX(0, m->geom_group[geomid]));
+
+  return (geomgroup[groupid]==0);
+}
+
+
+
+// compute solution from quadratic:  a*x^2 + 2*b*x + c = 0
+static mjtNum ray_quad(mjtNum a, mjtNum b, mjtNum c, mjtNum* x) {
+  // compute determinant and check
+  mjtNum det = b*b - a*c;
+  if (det<mjMINVAL) {
+    x[0] = -1;
+    x[1] = -1;
+    return -1;
+  }
+  det = mju_sqrt(det);
+
+  // compute the two solutions
+  x[0] = (-b-det)/a;
+  x[1] = (-b+det)/a;
+
+  // finalize result
+  if (x[0]>=0) {
+    return x[0];
+  } else if (x[1]>=0) {
+    return x[1];
+  } else {
+    return -1;
+  }
+}
+
+
+
+// intersect ray with triangle
+static mjtNum ray_triangle(mjtNum v[][3], const mjtNum* lpnt, const mjtNum* lvec,
+                           const mjtNum* b0, const mjtNum* b1) {
+  // dif = v[i] - lpnt
+  mjtNum dif[3][3];
+  for (int i=0; i<3; i++) {
+    for (int j=0; j<3; j++) {
+      dif[i][j] = v[i][j] - lpnt[j];
+    }
+  }
+
+  // project difference vectors in normal plane
+  mjtNum planar[3][2];
+  for (int i=0; i<3; i++) {
+    planar[i][0] = mju_dot3(b0, dif[i]);
+    planar[i][1] = mju_dot3(b1, dif[i]);
+  }
+
+  // reject if on the same side of any coordinate axis
+  if ((planar[0][0]>0 && planar[1][0]>0 && planar[2][0]>0) ||
+      (planar[0][0]<0 && planar[1][0]<0 && planar[2][0]<0) ||
+      (planar[0][1]>0 && planar[1][1]>0 && planar[2][1]>0) ||
+      (planar[0][1]<0 && planar[1][1]<0 && planar[2][1]<0)) {
+    return -1;
+  }
+
+  // determine if origin is inside planar projection of triangle
+  // A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  mjtNum A[4] = {planar[0][0]-planar[2][0], planar[1][0]-planar[2][0],
+                 planar[0][1]-planar[2][1], planar[1][1]-planar[2][1]};
+  mjtNum b[2] = {-planar[2][0], -planar[2][1]};
+  mjtNum det = A[0]*A[3] - A[1]*A[2];
+  if (mju_abs(det)<mjMINVAL) {
+    return -1;
+  }
+  mjtNum t0 = (A[3]*b[0] - A[1]*b[1]) / det;
+  mjtNum t1 = (-A[2]*b[0] + A[0]*b[1]) / det;
+
+  // check if outside
+  if (t0<0 || t1<0|| t0+t1>1) {
+    return -1;
+  }
+
+  // intersect ray with plane of triangle
+  mju_sub3(dif[0], v[0], v[2]);       // v0-v2
+  mju_sub3(dif[1], v[1], v[2]);       // v1-v2
+  mju_sub3(dif[2], lpnt, v[2]);       // lp-v2
+  mjtNum nrm[3];
+  mju_cross(nrm, dif[0], dif[1]);     // normal to triangle plane
+  mjtNum denom = mju_dot3(lvec, nrm);
+  if (mju_abs(denom)<mjMINVAL) {
+    return -1;
+  }
+
+  return (-mju_dot3(dif[2], nrm) / denom);
+}
+
+
+
+//---------------------------- geom-specific intersection functions --------------------------------
+
+// plane
+static mjtNum ray_plane(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                        const mjtNum* pnt, const mjtNum* vec) {
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(pos, mat, pnt, vec, lpnt, lvec);
+
+  // z-vec not pointing towards front face: reject
+  if (lvec[2]>-mjMINVAL) {
+    return -1;
+  }
+
+  // intersection with plane
+  const mjtNum x = -lpnt[2]/lvec[2];
+  if (x<0) {
+    return -1;
+  }
+  mjtNum p0 = lpnt[0] + x*lvec[0];
+  mjtNum p1 = lpnt[1] + x*lvec[1];
+
+  // accept only within rendered rectangle
+  if ((size[0]<=0 || mju_abs(p0)<=size[0]) &&
+      (size[1]<=0 || mju_abs(p1)<=size[1])) {
+    return x;
+  } else {
+    return -1;
+  }
+}
+
+
+
+// sphere
+static mjtNum ray_sphere(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                         const mjtNum* pnt, const mjtNum* vec) {
+  // (x*vec+pnt-pos)'*(x*vec+pnt-pos) = size[0]*size[0]
+  mjtNum dif[3] = {pnt[0]-pos[0], pnt[1]-pos[1], pnt[2]-pos[2]};
+  mjtNum a = vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2];
+  mjtNum b = vec[0]*dif[0] + vec[1]*dif[1] + vec[2]*dif[2];
+  mjtNum c = dif[0]*dif[0] + dif[1]*dif[1] + dif[2]*dif[2] - size[0]*size[0];
+
+  // solve a*x^2 + 2*b*x + c = 0
+  mjtNum xx[2];
+  return ray_quad(a, b, c, xx);
+}
+
+
+
+// capsule
+static mjtNum ray_capsule(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                          const mjtNum* pnt, const mjtNum* vec) {
+  // bounding sphere test
+  mjtNum ssz = size[0] + size[1];
+  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+    return -1;
+  }
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(pos, mat, pnt, vec, lpnt, lvec);
+
+  // init solution
+  mjtNum x = -1, sol, xx[2];
+
+  // cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  mjtNum a = lvec[0]*lvec[0] + lvec[1]*lvec[1];
+  mjtNum b = lvec[0]*lpnt[0] + lvec[1]*lpnt[1];
+  mjtNum c = lpnt[0]*lpnt[0] + lpnt[1]*lpnt[1] - size[0]*size[0];
+
+  // solve a*x^2 + 2*b*x + c = 0
+  sol = ray_quad(a, b, c, xx);
+
+  // make sure round solution is between flat sides
+  if (sol>=0 && mju_abs(lpnt[2]+sol*lvec[2])<=size[1]) {
+    if (x<0 || sol<x) {
+      x = sol;
+    }
+  }
+
+  // top cap
+  mjtNum ldif[3] = {lpnt[0], lpnt[1], lpnt[2]-size[1]};
+  a = lvec[0]*lvec[0] + lvec[1]*lvec[1] + lvec[2]*lvec[2];
+  b = lvec[0]*ldif[0] + lvec[1]*ldif[1] + lvec[2]*ldif[2];
+  c = ldif[0]*ldif[0] + ldif[1]*ldif[1] + ldif[2]*ldif[2] - size[0]*size[0];
+  ray_quad(a, b, c, xx);
+
+  // accept only top half of sphere
+  for (int i=0; i<2; i++) {
+    if (xx[i]>=0 && lpnt[2]+xx[i]*lvec[2]>=size[1]) {
+      if (x<0 || xx[i]<x) {
+        x = xx[i];
+      }
+    }
+  }
+
+  // bottom cap
+  ldif[2] = lpnt[2]+size[1];
+  b = lvec[0]*ldif[0] + lvec[1]*ldif[1] + lvec[2]*ldif[2];
+  c = ldif[0]*ldif[0] + ldif[1]*ldif[1] + ldif[2]*ldif[2] - size[0]*size[0];
+  ray_quad(a, b, c, xx);
+
+  // accept only bottom half of sphere
+  for (int i=0; i<2; i++) {
+    if (xx[i]>=0 && lpnt[2]+xx[i]*lvec[2]<=-size[1]) {
+      if (x<0 || xx[i]<x) {
+        x = xx[i];
+      }
+    }
+  }
+
+  return x;
+}
+
+
+
+// ellipsoid
+static mjtNum ray_ellipsoid(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                            const mjtNum* pnt, const mjtNum* vec) {
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(pos, mat, pnt, vec, lpnt, lvec);
+
+  // invert size^2
+  mjtNum s[3] = {1/(size[0]*size[0]), 1/(size[1]*size[1]), 1/(size[2]*size[2])};
+
+  // (x*lvec+lpnt)' * diag(1./size^2) * (x*lvec+lpnt) = 1
+  mjtNum a = s[0]*lvec[0]*lvec[0] + s[1]*lvec[1]*lvec[1] + s[2]*lvec[2]*lvec[2];
+  mjtNum b = s[0]*lvec[0]*lpnt[0] + s[1]*lvec[1]*lpnt[1] + s[2]*lvec[2]*lpnt[2];
+  mjtNum c = s[0]*lpnt[0]*lpnt[0] + s[1]*lpnt[1]*lpnt[1] + s[2]*lpnt[2]*lpnt[2] - 1;
+
+  // solve a*x^2 + 2*b*x + c = 0
+  mjtNum xx[2];
+  return ray_quad(a, b, c, xx);
+}
+
+
+
+// cylinder
+static mjtNum ray_cylinder(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                           const mjtNum* pnt, const mjtNum* vec) {
+  // bounding sphere test
+  mjtNum ssz = mju_sqrt(size[0]*size[0] + size[1]*size[1]);
+  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+    return -1;
+  }
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(pos, mat, pnt, vec, lpnt, lvec);
+
+  // init solution
+  mjtNum x = -1, sol;
+
+  // flat sides
+  int side;
+  if (mju_abs(lvec[2])>mjMINVAL) {
+    for (side=-1; side<=1; side+=2) {
+      // soludion of: lpnt[2] + x*lvec[2] = side*height_size
+      sol = (side*size[1]-lpnt[2])/lvec[2];
+
+      // process if non-negative
+      if (sol>=0) {
+        // intersection with horizontal face
+        mjtNum p0 = lpnt[0] + sol*lvec[0];
+        mjtNum p1 = lpnt[1] + sol*lvec[1];
+
+        // accept within radius
+        if (p0*p0 + p1*p1 <= size[0]*size[0]) {
+          if (x<0 || sol<x) {
+            x = sol;
+          }
+        }
+      }
+    }
+  }
+
+  // (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  mjtNum a = lvec[0]*lvec[0] + lvec[1]*lvec[1];
+  mjtNum b = lvec[0]*lpnt[0] + lvec[1]*lpnt[1];
+  mjtNum c = lpnt[0]*lpnt[0] + lpnt[1]*lpnt[1] - size[0]*size[0];
+
+  // solve a*x^2 + 2*b*x + c = 0
+  mjtNum xx[2];
+  sol = ray_quad(a, b, c, xx);
+
+  // make sure round solution is between flat sides
+  if (sol>=0 && mju_abs(lpnt[2]+sol*lvec[2])<=size[1]) {
+    if (x<0 || sol<x) {
+      x = sol;
+    }
+  }
+
+  return x;
+}
+
+
+
+// box
+static mjtNum ray_box(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                      const mjtNum* pnt, const mjtNum* vec, mjtNum* all) {
+  // clear all
+  if (all) {
+    for (int i=0; i<6; i++) {
+      all[i] = -1;
+    }
+  }
+
+  // bounding sphere test
+  mjtNum ssz = mju_sqrt(size[0]*size[0] + size[1]*size[1] + size[2]*size[2]);
+  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+    return -1;
+  }
+
+  // faces
+  const int iface[3][2] = {
+    {1, 2},
+    {0, 2},
+    {0, 1}
+  };
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(pos, mat, pnt, vec, lpnt, lvec);
+
+  // init solution
+  mjtNum x = -1, sol;
+
+  // loop over axes with non-zero vec
+  for (int i=0; i<3; i++) {
+    if (mju_abs(lvec[i])>mjMINVAL) {
+      for (int side=-1; side<=1; side+=2) {
+        // soludion of: lpnt[i] + x*lvec[i] = side*size[i]
+        sol = (side*size[i]-lpnt[i])/lvec[i];
+
+        // process if non-negative
+        if (sol>=0) {
+          // intersection with face
+          mjtNum p0 = lpnt[iface[i][0]] + sol*lvec[iface[i][0]];
+          mjtNum p1 = lpnt[iface[i][1]] + sol*lvec[iface[i][1]];
+
+          // accept within rectangle
+          if (mju_abs(p0)<=size[iface[i][0]] &&
+              mju_abs(p1)<=size[iface[i][1]]) {
+            // update
+            if (x<0 || sol<x) {
+              x = sol;
+            }
+
+            // save in all
+            if (all) {
+              all[2*i+(side+1)/2] = sol;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  return x;
+}
+
+
+
+// interect ray with hfield
+mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int id,
+                    const mjtNum* pnt, const mjtNum* vec) {
+  // check geom type
+  if (m->geom_type[id]!=mjGEOM_HFIELD) {
+    mju_error("mj_rayHfield: geom with hfield type expected");
+  }
+
+  // hfield id and dimensions
+  int hid = m->geom_dataid[id];
+  int nrow = m->hfield_nrow[hid];
+  int ncol = m->hfield_ncol[hid];
+  const mjtNum* size = m->hfield_size + 4*hid;
+  const float* data = m->hfield_data + m->hfield_adr[hid];
+
+  // compute size and pos of base box
+  mjtNum base_size[3] = {size[0], size[1], size[3]*0.5};
+  mjtNum base_pos[3] = {
+    d->geom_xpos[3*id]   - d->geom_xmat[9*id+2]*size[3]*0.5,
+    d->geom_xpos[3*id+1] - d->geom_xmat[9*id+5]*size[3]*0.5,
+    d->geom_xpos[3*id+2] - d->geom_xmat[9*id+8]*size[3]*0.5
+  };
+
+  // compute size and pos of top box
+  mjtNum top_size[3] = {size[0], size[1], size[2]*0.5};
+  mjtNum top_pos[3] = {
+    d->geom_xpos[3*id]   + d->geom_xmat[9*id+2]*size[2]*0.5,
+    d->geom_xpos[3*id+1] + d->geom_xmat[9*id+5]*size[2]*0.5,
+    d->geom_xpos[3*id+2] + d->geom_xmat[9*id+8]*size[2]*0.5
+  };
+
+  // init: intersection with base box
+  mjtNum x = ray_box(base_pos, d->geom_xmat+9*id, base_size, pnt, vec, NULL);
+
+  // check top box: done if no intersection
+  mjtNum all[6];
+  mjtNum top_intersect = ray_box(top_pos, d->geom_xmat+9*id, top_size, pnt, vec, all);
+  if (top_intersect<0) {
+    return x;
+  }
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(d->geom_xpos+3*id, d->geom_xmat+9*id, pnt, vec, lpnt, lvec);
+
+  // construct basis vectors of normal plane
+  mjtNum b0[3] = {1, 1, 1}, b1[3];
+  if (mju_abs(lvec[0])>=mju_abs(lvec[1]) && mju_abs(lvec[0])>=mju_abs(lvec[2])) {
+    b0[0] = 0;
+  } else if (mju_abs(lvec[1])>=mju_abs(lvec[2])) {
+    b0[1] = 0;
+  } else {
+    b0[2] = 0;
+  }
+  mju_addScl3(b1, b0, lvec, -mju_dot3(lvec, b0)/mju_dot3(lvec, lvec));
+  mju_normalize3(b1);
+  mju_cross(b0, b1, lvec);
+  mju_normalize3(b0);
+
+  // find ray segment intersecting top box
+  mjtNum seg[2] = {0, top_intersect};
+  for (int i=0; i<6; i++) {
+    if (all[i]>seg[1]) {
+      seg[0] = top_intersect;
+      seg[1] = all[i];
+    }
+  }
+
+  // project segment endpoints in horizontal plane, discretize
+  mjtNum dx = (2.0*size[0]) / (ncol-1);
+  mjtNum dy = (2.0*size[1]) / (nrow-1);
+  mjtNum SX[2], SY[2];
+  for (int i=0; i<2; i++) {
+    SX[i] = (lpnt[0] + seg[i]*lvec[0] + size[0]) / dx;
+    SY[i] = (lpnt[1] + seg[i]*lvec[1] + size[1]) / dy;
+  }
+
+  // compute ranges, with +1 padding
+  int cmin = mjMAX(0, (int)mju_floor(mjMIN(SX[0], SX[1]))-1);
+  int cmax = mjMIN(ncol-1, (int)mju_ceil(mjMAX(SX[0], SX[1]))+1);
+  int rmin = mjMAX(0, (int)mju_floor(mjMIN(SY[0], SY[1]))-1);
+  int rmax = mjMIN(nrow-1, (int)mju_ceil(mjMAX(SY[0], SY[1]))+1);
+
+  // check triangles within bounds
+  for (int r=rmin; r<rmax; r++) {
+    for (int c=cmin; c<cmax; c++) {
+      // first triangle
+      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-size[1], data[r*ncol+(c+1)]*size[2]}
+      };
+      mjtNum sol = ray_triangle(va, lpnt, lvec, b0, b1);
+      if (sol>=0 && (x<0 || sol<x)) {
+        x = sol;
+      }
+
+      // second triangle
+      mjtNum vb[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-size[0], dy*(r+1)-size[1], data[(r+1)*ncol+c]*size[2]}
+      };
+      sol = ray_triangle(vb, lpnt, lvec, b0, b1);
+      if (sol>=0 && (x<0 || sol<x)) {
+        x = sol;
+      }
+    }
+  }
+
+  // check viable sides of top box
+  for (int i=0; i<4; i++) {
+    if (all[i]>=0 && (all[i]<x || x<0)) {
+      // normalized height of intersection point
+      mjtNum z = (lpnt[2] + all[i]*lvec[2]) / size[2];
+
+      // rectangle points
+      mjtNum y, y0, z0, z1;
+
+      // side normal to x-axis
+      if (i<2) {
+        y = (lpnt[1] + all[i]*lvec[1] + size[1]) / dy;
+        y0 = mjMAX(0, mjMIN(nrow-2, mju_floor(y)));
+        z0 = (mjtNum)data[mju_round(y0)*nrow + (i==1 ? ncol-1 : 0)];
+        z1 = (mjtNum)data[mju_round(y0+1)*nrow + (i==1 ? ncol-1 : 0)];
+      }
+
+      // side normal to y-axis
+      else {
+        y = (lpnt[0] + all[i]*lvec[0] + size[0]) / dx;
+        y0 = mjMAX(0, mjMIN(ncol-2, mju_floor(y)));
+        z0 = (mjtNum)data[mju_round(y0) + (i==3 ? (nrow-1)*ncol : 0)];
+        z1 = (mjtNum)data[mju_round(y0+1) + (i==3 ? (nrow-1)*ncol : 0)];
+      }
+
+      // check if point is below line segment
+      if (z < z0*(y0+1-y) + z1*(y-y0)) {
+        x = all[i];
+      }
+    }
+  }
+
+  return x;
+}
+
+
+
+// interect ray with mesh
+mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int id,
+                  const mjtNum* pnt, const mjtNum* vec) {
+  // check geom type
+  if (m->geom_type[id]!=mjGEOM_MESH) {
+    mju_error("mj_rayMesh: geom with mesh type expected");
+  }
+
+  // bounding box test
+  if (ray_box(d->geom_xpos+3*id, d->geom_xmat+9*id, m->geom_size+3*id, pnt, vec, NULL)<0) {
+    return -1;
+  }
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  ray_map(d->geom_xpos+3*id, d->geom_xmat+9*id, pnt, vec, lpnt, lvec);
+
+  // construct basis vectors of normal plane
+  mjtNum b0[3] = {1, 1, 1}, b1[3];
+  if (mju_abs(lvec[0])>=mju_abs(lvec[1]) && mju_abs(lvec[0])>=mju_abs(lvec[2])) {
+    b0[0] = 0;
+  } else if (mju_abs(lvec[1])>=mju_abs(lvec[2])) {
+    b0[1] = 0;
+  } else {
+    b0[2] = 0;
+  }
+  mju_addScl3(b1, b0, lvec, -mju_dot3(lvec, b0)/mju_dot3(lvec, lvec));
+  mju_normalize3(b1);
+  mju_cross(b0, b1, lvec);
+  mju_normalize3(b0);
+
+  // init solution
+  mjtNum x = -1, sol;
+
+  // process all triangles
+  int face, meshid = m->geom_dataid[id];
+  for (face = m->mesh_faceadr[meshid];
+       face < m->mesh_faceadr[meshid] + m->mesh_facenum[meshid];
+       face++) {
+    // get float vertices
+    float* vf[3];
+    vf[0] = m->mesh_vert + 3*(m->mesh_face[3*face]   + m->mesh_vertadr[meshid]);
+    vf[1] = m->mesh_vert + 3*(m->mesh_face[3*face+1] + m->mesh_vertadr[meshid]);
+    vf[2] = m->mesh_vert + 3*(m->mesh_face[3*face+2] + m->mesh_vertadr[meshid]);
+
+    // convert to mjtNum
+    mjtNum v[3][3];
+    for (int i=0; i<3; i++) {
+      for (int j=0; j<3; j++) {
+        v[i][j] = (mjtNum)vf[i][j];
+      }
+    }
+
+    // solve
+    sol = ray_triangle(v, lpnt, lvec, b0, b1);
+
+    // update
+    if (sol>=0 && (x<0 || sol<x)) {
+      x = sol;
+    }
+  }
+
+  return x;
+}
+
+
+
+// interect ray with pure geom, no meshes or hfields
+mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                   const mjtNum* pnt, const mjtNum* vec, int geomtype) {
+  switch (geomtype) {
+  case mjGEOM_PLANE:
+    return ray_plane(pos, mat, size, pnt, vec);
+
+  case mjGEOM_SPHERE:
+    return ray_sphere(pos, mat, size, pnt, vec);
+
+  case mjGEOM_CAPSULE:
+    return ray_capsule(pos, mat, size, pnt, vec);
+
+  case mjGEOM_ELLIPSOID:
+    return ray_ellipsoid(pos, mat, size, pnt, vec);
+
+  case mjGEOM_CYLINDER:
+    return ray_cylinder(pos, mat, size, pnt, vec);
+
+  case mjGEOM_BOX:
+    return ray_box(pos, mat, size, pnt, vec, NULL);
+
+  default:
+    mju_error_i("mju_rayGeom: unexpected geom type %d", geomtype);
+    return -1;
+  }
+}
+
+
+
+// interect ray with skin, return nearest vertex id
+mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
+                   const mjtNum* pnt, const mjtNum* vec, int vertid[1]) {
+  // compute bounding box
+  mjtNum box[3][2] = {{0, 0}, {0, 0}, {0, 0}};
+  for (int i=0; i<nvert; i++) {
+    for (int j=0; j<3; j++) {
+      // update minimum along side j
+      if (box[j][0]>vert[3*i+j] || i==0) {
+        box[j][0] = vert[3*i+j];
+      }
+
+      // update maximum along side j
+      if (box[j][1]<vert[3*i+j] || i==0) {
+        box[j][1] = vert[3*i+j];
+      }
+    }
+  }
+
+  // construct box geom
+  mjtNum pos[3], size[3], mat[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
+  for (int j=0; j<3; j++) {
+    pos[j] = 0.5*(box[j][0]+box[j][1]);
+    size[j] = 0.5*(box[j][1]-box[j][0]);
+  }
+
+  // apply bounding-box filter
+  if (ray_box(pos, mat, size, pnt, vec, NULL)<0) {
+    return -1;
+  }
+
+  // construct basis vectors of normal plane
+  mjtNum b0[3] = {1, 1, 1}, b1[3];
+  if (mju_abs(vec[0])>=mju_abs(vec[1]) && mju_abs(vec[0])>=mju_abs(vec[2])) {
+    b0[0] = 0;
+  } else if (mju_abs(vec[1])>=mju_abs(vec[2])) {
+    b0[1] = 0;
+  } else {
+    b0[2] = 0;
+  }
+  mju_addScl3(b1, b0, vec, -mju_dot3(vec, b0)/mju_dot3(vec, vec));
+  mju_normalize3(b1);
+  mju_cross(b0, b1, vec);
+  mju_normalize3(b0);
+
+  // init solution
+  mjtNum x = -1, sol;
+
+  // process all faces
+  for (int i=0; i<nface; i++) {
+    // get float vertices
+    const float* vf[3];
+    vf[0] = vert + 3*(face[3*i]);
+    vf[1] = vert + 3*(face[3*i+1]);
+    vf[2] = vert + 3*(face[3*i+2]);
+
+    // convert to mjtNum
+    mjtNum v[3][3];
+    for (int j=0; j<3; j++) {
+      for (int k=0; k<3; k++) {
+        v[j][k] = (mjtNum)vf[j][k];
+      }
+    }
+
+    // solve
+    sol = ray_triangle(v, pnt, vec, b0, b1);
+
+    // update
+    if (sol>=0 && (x<0 || sol<x)) {
+      x = sol;
+
+      // construct intersection point
+      mjtNum intersect[3];
+      mju_addScl3(intersect, pnt, vec, sol);
+
+      // find nearest vertex
+      mjtNum dist = mju_dist3(intersect, v[0]);
+      *vertid = face[3*i];
+      for (int j=1; j<3; j++) {
+        mjtNum newdist = mju_dist3(intersect, v[j]);
+        if (newdist<dist) {
+          dist = newdist;
+          *vertid = face[3*i+j];
+        }
+      }
+    }
+  }
+
+  return x;
+}
+
+
+
+//---------------------------- main entry point ---------------------------------------------------
+
+// intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
+//  return geomid and distance (x) to nearest surface, or -1 if no intersection
+//  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion
+mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum* vec,
+              const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+              int geomid[1]) {
+  mjtNum dist, newdist;
+
+  // check vector length
+  if (mju_norm3(vec)<mjMINVAL) {
+    mju_error("mj_ray: vector length is too small");
+  }
+
+  // clear result
+  dist = -1;
+  *geomid = -1;
+
+  // loop over geoms not eliminated by mask and bodyexclude
+  for (int i=0; i<m->ngeom; i++) {
+    if (!ray_eliminate(m, d, i, geomgroup, flg_static, bodyexclude)) {
+      // handle mesh and hfield separately
+      if (m->geom_type[i]==mjGEOM_MESH) {
+        newdist = mj_rayMesh(m, d, i, pnt, vec);
+      } else if (m->geom_type[i]==mjGEOM_HFIELD) {
+        newdist = mj_rayHfield(m, d, i, pnt, vec);
+      }
+
+      // otherwise general dispatch
+      else {
+        newdist = mju_rayGeom(d->geom_xpos+3*i, d->geom_xmat+9*i,
+                              m->geom_size+3*i, pnt, vec, m->geom_type[i]);
+      }
+
+      // update if closer intersection found
+      if (newdist>=0 && (newdist<dist || dist<0)) {
+        dist = newdist;
+        *geomid = i;
+      }
+    }
+  }
+
+  return dist;
+}

src/engine/engine_ray.h

@@ -0,0 +1,53 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_RAY_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_RAY_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
+//  return geomid and distance (x) to nearest surface, or -1 if no intersection
+//  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion
+MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum* vec,
+                    const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                    int geomid[1]);
+
+// interect ray with hfield
+MJAPI mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
+                          const mjtNum* pnt, const mjtNum* vec);
+
+// interect ray with mesh
+MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
+                        const mjtNum* pnt, const mjtNum* vec);
+
+// interect ray with pure geom, no meshes or hfields
+MJAPI mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                         const mjtNum* pnt, const mjtNum* vec, int geomtype);
+
+// interect ray with skin, return nearest vertex id
+MJAPI mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
+                         const mjtNum* pnt, const mjtNum* vec, int vertid[1]);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_RAY_H_

src/engine/engine_sensor.c

@@ -0,0 +1,747 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_sensor.h"
+
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_callback.h"
+#include "engine/engine_core_smooth.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_ray.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_spatial.h"
+
+
+//-------------------------------- utility ---------------------------------------------------------
+
+// add sensor noise after each stage
+static void add_noise(const mjModel* m, mjData* d, mjtStage stage) {
+  int adr, dim;
+  mjtNum rnd[4], noise, quat[4], res[4];
+
+  // process sensors matching stage and having positive noise
+  for (int i=0; i<m->nsensor; i++) {
+    if (m->sensor_needstage[i]==stage && m->sensor_noise[i]>0) {
+      // get sensor info
+      adr = m->sensor_adr[i];
+      dim = m->sensor_dim[i];
+      noise = m->sensor_noise[i];
+
+      // real or positive: add noise directly, with clamp for positive
+      if (m->sensor_datatype[i]==mjDATATYPE_REAL ||
+          m->sensor_datatype[i]==mjDATATYPE_POSITIVE)
+        for (int j=0; j<dim; j++) {
+          // get random numbers; use only the first one
+          rnd[0] = mju_standardNormal(rnd+1);
+
+          // positive
+          if (m->sensor_datatype[i]==mjDATATYPE_POSITIVE) {
+            // add noise only if positive, keep it positive
+            if (d->sensordata[adr+j]>0) {
+              d->sensordata[adr+j] = mjMAX(0, d->sensordata[adr+j]+rnd[0]*noise);
+            }
+          }
+
+          // real
+          else {
+            d->sensordata[adr+j] += rnd[0]*noise;
+          }
+        }
+
+      // axis or quat: rotate around random axis by random angle
+      else {
+        // get four random numbers
+        rnd[0] = mju_standardNormal(rnd+1);
+        rnd[2] = mju_standardNormal(rnd+3);
+
+        // scale angle, normalize axis, make quaterion
+        rnd[0] *= noise;
+        mju_normalize3(rnd+1);
+        mju_axisAngle2Quat(quat, rnd+1, rnd[0]);
+
+        // axis
+        if (m->sensor_datatype[i]==mjDATATYPE_AXIS) {
+          // apply quaternion rotation to axis, assign
+          mju_rotVecQuat(res, d->sensordata+adr, quat);
+          mju_copy3(d->sensordata+adr, res);
+        }
+
+        // quaternion
+        else if (m->sensor_datatype[i]==mjDATATYPE_QUATERNION) {
+          // apply quaternion rotation to quaternion, assign
+          mju_mulQuat(res, d->sensordata+adr, quat);
+          mju_copy4(d->sensordata+adr, res);
+        }
+
+        // unknown datatype
+        else {
+          mju_error_i("Unknown datatype in sensor %d", i);
+        }
+      }
+    }
+  }
+}
+
+
+
+// apply cutoff after each stage
+static void apply_cutoff(const mjModel* m, mjData* d, mjtStage stage) {
+  // process sensors matching stage and having positive cutoff
+  for (int i=0; i<m->nsensor; i++) {
+    if (m->sensor_needstage[i]==stage && m->sensor_cutoff[i]>0) {
+      // get sensor info
+      int adr = m->sensor_adr[i];
+      int dim = m->sensor_dim[i];
+      mjtNum cutoff = m->sensor_cutoff[i];
+
+      // process all dimensions
+      for (int j=0; j<dim; j++)
+
+        // real: apply on both sides
+        if (m->sensor_datatype[i]==mjDATATYPE_REAL)
+          d->sensordata[adr+j] =
+            mju_min(cutoff, mju_max(-cutoff, d->sensordata[adr+j]));
+
+        // positive: apply on positive side only
+        else if (m->sensor_datatype[i]==mjDATATYPE_POSITIVE)
+          d->sensordata[adr+j] =
+            mju_min(cutoff, d->sensordata[adr+j]);
+    }
+  }
+}
+
+
+
+// get xpos and xmat pointers to an object in mjData
+static void get_xpos_xmat(const mjData* d, int type, int id, int sensor_id,
+                          mjtNum **xpos, mjtNum **xmat) {
+  switch (type) {
+  case mjOBJ_XBODY:
+    *xpos = d->xpos + 3*id;
+    *xmat = d->xmat + 9*id;
+    break;
+  case mjOBJ_BODY:
+    *xpos = d->xipos + 3*id;
+    *xmat = d->ximat + 9*id;
+    break;
+  case mjOBJ_GEOM:
+    *xpos = d->geom_xpos + 3*id;
+    *xmat = d->geom_xmat + 9*id;
+    break;
+  case mjOBJ_SITE:
+    *xpos = d->site_xpos + 3*id;
+    *xmat = d->site_xmat + 9*id;
+    break;
+  case mjOBJ_CAMERA:
+    *xpos = d->cam_xpos + 3*id;
+    *xmat = d->cam_xmat + 9*id;
+    break;
+  default:
+    mju_error_i("Invalid object type in sensor %d", sensor_id);
+  }
+}
+
+// get global quaternion of an object in mjData
+static void get_xquat(const mjModel* m, const mjData* d, int type, int id, int sensor_id,
+                      mjtNum *quat) {
+  switch (type) {
+  case mjOBJ_XBODY:
+    mju_copy4(quat, d->xquat+4*id);
+    break;
+  case mjOBJ_BODY:
+    mju_mulQuat(quat, d->xquat+4*id, m->body_iquat+4*id);
+    break;
+  case mjOBJ_GEOM:
+    mju_mulQuat(quat, d->xquat+4*m->geom_bodyid[id], m->geom_quat+4*id);
+    break;
+  case mjOBJ_SITE:
+    mju_mulQuat(quat, d->xquat+4*m->site_bodyid[id], m->site_quat+4*id);
+    break;
+  case mjOBJ_CAMERA:
+    mju_mulQuat(quat, d->xquat+4*m->cam_bodyid[id], m->cam_quat+4*id);
+    break;
+  default:
+    mju_error_i("Invalid object type in sensor %d", sensor_id);
+  }
+}
+
+
+//-------------------------------- sensor ----------------------------------------------------------
+
+// position-dependent sensors
+void mj_sensorPos(const mjModel* m, mjData* d) {
+  int rgeomid, objtype, objid, reftype, refid, adr, offset, nusersensor = 0;
+  int ne = d->ne, nf = d->nf, nefc = d->nefc;
+  mjtNum rvec[3], *xpos, *xmat, *xpos_ref, *xmat_ref;
+
+  // process sensors matching stage
+  for (int i=0; i<m->nsensor; i++) {
+    if (m->sensor_needstage[i]==mjSTAGE_POS) {
+      // get sensor info
+      objtype = m->sensor_objtype[i];
+      objid = m->sensor_objid[i];
+      refid = m->sensor_refid[i];
+      reftype = m->sensor_reftype[i];
+      adr = m->sensor_adr[i];
+
+      // process according to type
+      switch (m->sensor_type[i]) {
+      case mjSENS_MAGNETOMETER:                           // magnetometer
+        mju_mulMatTVec(d->sensordata+adr, d->site_xmat+9*objid, m->opt.magnetic, 3, 3);
+        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], &rgeomid);
+        break;
+
+      case mjSENS_JOINTPOS:                               // jointpos
+        d->sensordata[adr] = d->qpos[m->jnt_qposadr[objid]];
+        break;
+
+      case mjSENS_TENDONPOS:                              // tendonpos
+        d->sensordata[adr] = d->ten_length[objid];
+        break;
+
+      case mjSENS_ACTUATORPOS:                            // actuatorpos
+        d->sensordata[adr] = d->actuator_length[objid];
+        break;
+
+      case mjSENS_BALLQUAT:                               // ballquat
+        mju_copy4(d->sensordata+adr, d->qpos+m->jnt_qposadr[objid]);
+        break;
+
+      case mjSENS_JOINTLIMITPOS:                          // jointlimitpos
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_JOINT && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_pos[j] - d->efc_margin[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_TENDONLIMITPOS:                         // tendonlimitpos
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_TENDON && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_pos[j] - d->efc_margin[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_FRAMEPOS:                               // framepos
+      case mjSENS_FRAMEXAXIS:                             // framexaxis
+      case mjSENS_FRAMEYAXIS:                             // frameyaxis
+      case mjSENS_FRAMEZAXIS:                             // framezaxis
+        // get xpos and xmat pointers for object frame
+        get_xpos_xmat(d, objtype, objid, i, &xpos, &xmat);
+
+        // reference frame unspecified: global frame
+        if (refid == -1) {
+          if (m->sensor_type[i]==mjSENS_FRAMEPOS) {
+            mju_copy3(d->sensordata+adr, xpos);
+          } else {
+            // offset = (0 or 1 or 2) for (x or y or z)-axis sensors, respectively
+            offset = m->sensor_type[i] - mjSENS_FRAMEXAXIS;
+            d->sensordata[adr] = xmat[offset];
+            d->sensordata[adr+1] = xmat[offset+3];
+            d->sensordata[adr+2] = xmat[offset+6];
+          }
+        }
+
+        // reference frame specified
+        else {
+          get_xpos_xmat(d, reftype, refid, i, &xpos_ref, &xmat_ref);
+          if (m->sensor_type[i]==mjSENS_FRAMEPOS) {
+            mju_sub3(rvec, xpos, xpos_ref);
+            mju_rotVecMatT(d->sensordata+adr, rvec, xmat_ref);
+          } else {
+            // offset = (0 or 1 or 2) for (x or y or z)-axis sensors, respectively
+            offset = m->sensor_type[i] - mjSENS_FRAMEXAXIS;
+            mjtNum axis[3] = {xmat[offset], xmat[offset+3], xmat[offset+6]};
+            mju_rotVecMatT(d->sensordata+adr, axis, xmat_ref);
+          }
+        }
+        break;
+
+      case mjSENS_FRAMEQUAT:                              // framequat
+        {
+          // get global object quaternion
+          mjtNum objquat[4];
+          get_xquat(m, d, objtype, objid, i, objquat);
+
+          // reference frame unspecified: copy object quaternion
+          if (refid == -1) {
+            mju_copy4(d->sensordata+adr, objquat);
+          } else {
+            // reference frame specified, get global reference quaternion
+            mjtNum refquat[4];
+            get_xquat(m, d, reftype, refid, i, refquat);
+
+            // relative quaternion
+            mju_negQuat(refquat, refquat);
+            mju_mulQuat(d->sensordata+adr, refquat, objquat);
+          }
+        }
+        break;
+
+      case mjSENS_SUBTREECOM:                             // subtreecom
+        mju_copy3(d->sensordata+adr, d->subtree_com+3*objid);
+        break;
+
+      case mjSENS_USER:                                   // user
+        nusersensor++;
+        break;
+
+      default:
+        mju_error_i("Invalid sensor type in POS stage, sensor %d", i);
+      }
+    }
+  }
+
+  // fill in user sensors if detected
+  if (nusersensor && mjcb_sensor) {
+    mjcb_sensor(m, d, mjSTAGE_POS);
+  }
+
+  // add noise if enabled
+  if (mjENABLED(mjENBL_SENSORNOISE)) {
+    add_noise(m, d, mjSTAGE_POS);
+  }
+
+  // cutoff
+  apply_cutoff(m, d, mjSTAGE_POS);
+}
+
+
+
+// velocity-dependent sensors
+void mj_sensorVel(const mjModel* m, mjData* d) {
+  int type, objtype, objid, reftype, refid, adr, nusersensor = 0;
+  int ne = d->ne, nf = d->nf, nefc = d->nefc;
+  mjtNum xvel[6];
+
+  // process sensors matching stage
+  int subtreeVel = 0;
+  for (int i=0; i<m->nsensor; i++) {
+    if (m->sensor_needstage[i]==mjSTAGE_VEL) {
+      // get sensor info
+      type = m->sensor_type[i];
+      objtype = m->sensor_objtype[i];
+      objid = m->sensor_objid[i];
+      refid = m->sensor_refid[i];
+      reftype = m->sensor_reftype[i];
+      adr = m->sensor_adr[i];
+
+      // call mj_subtreeVel when first relevant sensor is encountered
+      if (subtreeVel==0 &&
+          (type==mjSENS_SUBTREELINVEL ||
+           type==mjSENS_SUBTREEANGMOM ||
+           type==mjSENS_USER)) {
+        // compute subtree_linvel, subtree_angmom
+        mj_subtreeVel(m, d);
+
+        // mark computed
+        subtreeVel = 1;
+      }
+
+      // process according to type
+      switch (type) {
+      case mjSENS_VELOCIMETER:                            // velocimeter
+        // xvel = site velocity, in site frame
+        mj_objectVelocity(m, d, mjOBJ_SITE, objid, xvel, 1);
+
+        // assign linear velocity
+        mju_copy3(d->sensordata+adr, xvel+3);
+        break;
+
+      case mjSENS_GYRO:                                   // gyro
+        // xvel = site velocity, in site frame
+        mj_objectVelocity(m, d, mjOBJ_SITE, objid, xvel, 1);
+
+        // assign angular velocity
+        mju_copy3(d->sensordata+adr, xvel);
+        break;
+
+      case mjSENS_JOINTVEL:                               // jointvel
+        d->sensordata[adr] = d->qvel[m->jnt_dofadr[objid]];
+        break;
+
+      case mjSENS_TENDONVEL:                              // tendonvel
+        d->sensordata[adr] = d->ten_velocity[objid];
+        break;
+
+      case mjSENS_ACTUATORVEL:                            // actuatorvel
+        d->sensordata[adr] = d->actuator_velocity[objid];
+        break;
+
+      case mjSENS_BALLANGVEL:                             // ballangvel
+        mju_copy3(d->sensordata+adr, d->qvel+m->jnt_dofadr[objid]);
+        break;
+
+      case mjSENS_JOINTLIMITVEL:                          // jointlimitvel
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_JOINT && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_vel[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_TENDONLIMITVEL:                         // tendonlimitvel
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_TENDON && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_vel[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_FRAMELINVEL:                            // framelinvel
+      case mjSENS_FRAMEANGVEL:                            // frameangvel
+        // xvel = 6D object velocity, in global frame
+        mj_objectVelocity(m, d, objtype, objid, xvel, 0);
+
+        if (refid > -1) {  // reference frame specified
+          mjtNum *xpos, *xmat, *xpos_ref, *xmat_ref, xvel_ref[6], rel_vel[6], cross[3], rvec[3];
+
+          // in global frame: object and reference position, reference orientation and velocity
+          get_xpos_xmat(d, objtype, objid, i, &xpos, &xmat);
+          get_xpos_xmat(d, reftype, refid, i, &xpos_ref, &xmat_ref);
+          mj_objectVelocity(m, d, reftype, refid, xvel_ref, 0);
+
+          // subtract velocities
+          mju_sub(rel_vel, xvel, xvel_ref, 6);
+
+          // linear velocity: add correction due to rotating reference frame
+          mju_sub3(rvec, xpos, xpos_ref);
+          mju_cross(cross, rvec, xvel_ref);
+          mju_addTo3(rel_vel+3, cross);
+
+          // project into reference frame
+          mju_rotVecMatT(xvel, rel_vel, xmat_ref);
+          mju_rotVecMatT(xvel+3, rel_vel+3, xmat_ref);
+        }
+
+        // copy linear or angular component
+        if (m->sensor_type[i]==mjSENS_FRAMELINVEL) {
+          mju_copy3(d->sensordata+adr, xvel+3);
+        } else {
+          mju_copy3(d->sensordata+adr, xvel);
+        }
+        break;
+
+      case mjSENS_SUBTREELINVEL:                          // subtreelinvel
+        mju_copy3(d->sensordata+adr, d->subtree_linvel+3*objid);
+        break;
+
+      case mjSENS_SUBTREEANGMOM:                          // subtreeangmom
+        mju_copy3(d->sensordata+adr, d->subtree_angmom+3*objid);
+        break;
+
+      case mjSENS_USER:                                   // user
+        nusersensor++;
+        break;
+
+      default:
+        mju_error_i("Invalid type in VEL stage, sensor %d", i);
+      }
+    }
+  }
+
+  // fill in user sensors if detected
+  if (nusersensor && mjcb_sensor) {
+    mjcb_sensor(m, d, mjSTAGE_VEL);
+  }
+
+  // add noise if enabled
+  if (mjENABLED(mjENBL_SENSORNOISE)) {
+    add_noise(m, d, mjSTAGE_VEL);
+  }
+
+  // cutoff
+  apply_cutoff(m, d, mjSTAGE_VEL);
+}
+
+
+
+// acceleration/force-dependent sensors
+void mj_sensorAcc(const mjModel* m, mjData* d) {
+  int rootid, bodyid, type, objtype, objid, body1, body2, adr, nusersensor = 0;
+  int ne = d->ne, nf = d->nf, nefc = d->nefc;
+  mjtNum tmp[6], conforce[6], conray[3];
+  mjContact* con;
+
+  // process sensors matching stage
+  int rnePost = 0;
+  for (int i=0; i<m->nsensor; i++) {
+    if (m->sensor_needstage[i]==mjSTAGE_ACC) {
+      // get sensor info
+      type =  m->sensor_type[i];
+      objtype = m->sensor_objtype[i];
+      objid = m->sensor_objid[i];
+      adr = m->sensor_adr[i];
+
+      // call mj_rnePostConstraint when first relevant sensor is encountered
+      if (rnePost==0                  &&
+          type!=mjSENS_TOUCH          &&
+          type!=mjSENS_ACTUATORFRC    &&
+          type!=mjSENS_JOINTLIMITFRC  &&
+          type!=mjSENS_TENDONLIMITFRC) {
+        // compute cacc, cfrc_int, cfrc_ext
+        mj_rnePostConstraint(m, d);
+
+        // mark computed
+        rnePost = 1;
+      }
+
+      // process according to type
+      switch (type) {
+      case mjSENS_TOUCH:                                  // touch
+        // extract body data
+        bodyid = m->site_bodyid[objid];
+        rootid = m->body_rootid[bodyid];
+
+        // clear result
+        d->sensordata[adr] = 0;
+
+        // find contacts in sensor zone, add normal forces
+        for (int j=0; j<d->ncon; j++) {
+          // contact pointer, contacting bodies
+          con = d->contact + j;
+          body1 = m->geom_bodyid[con->geom1];
+          body2 = m->geom_bodyid[con->geom2];
+
+          // select contacts involving sensorized body
+          if (con->efc_address>=0 && (bodyid==body1 || bodyid==body2)) {
+            // get contact force:torque in contact frame
+            mj_contactForce(m, d, j, conforce);
+
+            // nothing to do if normal is zero
+            if (conforce[0]<=0) {
+              continue;
+            }
+
+            // convert contact normal force to global frame, normalize
+            mju_scl3(conray, con->frame, conforce[0]);
+            mju_normalize3(conray);
+
+            // flip ray direction if sensor is on body2
+            if (bodyid==body2) {
+              mju_scl3(conray, conray, -1);
+            }
+
+            // add if ray-zone intersection (always true when con->pos inside zone)
+            if (mju_rayGeom(d->site_xpos+3*objid, d->site_xmat+9*objid,
+                            m->site_size+3*objid, con->pos, conray,
+                            m->site_type[objid]) >= 0) {
+              d->sensordata[adr] += conforce[0];
+            }
+          }
+        }
+        break;
+
+      case mjSENS_ACCELEROMETER:                          // accelerometer
+        // tmp = site acceleration, in site frame
+        mj_objectAcceleration(m, d, mjOBJ_SITE, objid, tmp, 1);
+
+        // assign linear acceleration
+        mju_copy3(d->sensordata+adr, tmp+3);
+        break;
+
+      case mjSENS_FORCE:                                  // force
+        // extract body data
+        bodyid = m->site_bodyid[objid];
+        rootid = m->body_rootid[bodyid];
+
+        // tmp = interaction force between body and parent, in site frame
+        mju_transformSpatial(tmp, d->cfrc_int+6*bodyid, 1,
+                             d->site_xpos+3*objid, d->subtree_com+3*rootid, d->site_xmat+9*objid);
+
+        // assign force
+        mju_copy3(d->sensordata+adr, tmp+3);
+        break;
+
+      case mjSENS_TORQUE:                                 // torque
+        // extract body data
+        bodyid = m->site_bodyid[objid];
+        rootid = m->body_rootid[bodyid];
+
+        // tmp = interaction force between body and parent, in site frame
+        mju_transformSpatial(tmp, d->cfrc_int+6*bodyid, 1,
+                             d->site_xpos+3*objid, d->subtree_com+3*rootid, d->site_xmat+9*objid);
+
+        // assign torque
+        mju_copy3(d->sensordata+adr, tmp);
+        break;
+
+      case mjSENS_ACTUATORFRC:                            // actuatorfrc
+        d->sensordata[adr] = d->actuator_force[objid];
+        break;
+
+      case mjSENS_JOINTLIMITFRC:                          // jointlimitfrc
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_JOINT && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_force[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_TENDONLIMITFRC:                         // tendonlimitfrc
+        d->sensordata[adr] = 0;
+        for (int j=ne+nf; j<nefc; j++) {
+          if (d->efc_type[j]==mjCNSTR_LIMIT_TENDON && d->efc_id[j]==objid) {
+            d->sensordata[adr] = d->efc_force[j];
+            break;
+          }
+        }
+        break;
+
+      case mjSENS_FRAMELINACC:                            // framelinacc
+      case mjSENS_FRAMEANGACC:                            // frameangacc
+        // get 6D object acceleration, in global frame
+        mj_objectAcceleration(m, d, objtype, objid, tmp, 0);
+
+        // copy linear or angular component
+        if (m->sensor_type[i]==mjSENS_FRAMELINACC) {
+          mju_copy3(d->sensordata+adr, tmp+3);
+        } else {
+          mju_copy3(d->sensordata+adr, tmp);
+        }
+        break;
+
+      case mjSENS_USER:                                   // user
+        nusersensor++;
+        break;
+
+      default:
+        mju_error_i("Invalid type in ACC stage, sensor %d", i);
+      }
+    }
+  }
+
+  // fill in user sensors if detected
+  if (nusersensor && mjcb_sensor) {
+    mjcb_sensor(m, d, mjSTAGE_ACC);
+  }
+
+  // add noise if enabled
+  if (mjENABLED(mjENBL_SENSORNOISE)) {
+    add_noise(m, d, mjSTAGE_ACC);
+  }
+
+  // cutoff
+  apply_cutoff(m, d, mjSTAGE_ACC);
+}
+
+
+
+//-------------------------------- energy ----------------------------------------------------------
+
+// position-dependent energy (potential)
+void mj_energyPos(const mjModel* m, mjData* d) {
+  int padr;
+  mjtNum dif[3], stiffness;
+
+  // disabled: clear and return
+  if (!mjENABLED(mjENBL_ENERGY)) {
+    d->energy[0] = d->energy[1] = 0;
+    return;
+  }
+
+  // init potential energy:  -sum_i body(i).mass * mju_dot(body(i).pos, gravity)
+  d->energy[0] = 0;
+  if (!mjDISABLED(mjDSBL_GRAVITY)) {
+    for (int i=1; i<m->nbody; i++) {
+      d->energy[0] -= m->body_mass[i] * mju_dot3(m->opt.gravity, d->xipos+3*i);
+    }
+  }
+
+  // add joint-level springs
+  if (!mjDISABLED(mjDSBL_PASSIVE)) {
+    for (int i=0; i<m->njnt; i++) {
+      stiffness = m->jnt_stiffness[i];
+      padr = m->jnt_qposadr[i];
+
+      switch (m->jnt_type[i]) {
+      case mjJNT_FREE:
+        mju_sub3(dif, d->qpos+padr, m->qpos_spring+padr);
+        d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif);
+
+        // continue with rotations
+        padr += 3;
+
+      case mjJNT_BALL:
+        // covert quatertion difference into angular "velocity"
+        mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr);
+        d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif);
+        break;
+
+      case mjJNT_SLIDE:
+      case mjJNT_HINGE:
+        d->energy[0] += 0.5*stiffness*
+                        (d->qpos[padr] - m->qpos_spring[padr])*
+                        (d->qpos[padr] - m->qpos_spring[padr]);
+        break;
+      }
+    }
+  }
+
+  // add tendon-level springs
+  if (!mjDISABLED(mjDSBL_PASSIVE)) {
+    for (int i=0; i<m->ntendon; i++) {
+      stiffness = m->tendon_stiffness[i];
+
+      d->energy[0] += 0.5*stiffness*(d->ten_length[i] - m->tendon_lengthspring[i])*
+                      (d->ten_length[i] - m->tendon_lengthspring[i]);
+    }
+  }
+}
+
+
+
+// velocity-dependent energy (kinetic)
+void mj_energyVel(const mjModel* m, mjData* d) {
+  mjtNum *vec;
+  mjMARKSTACK;
+
+  // return if disabled (already cleared in potential)
+  if (!mjENABLED(mjENBL_ENERGY)) {
+    return;
+  }
+
+  vec = mj_stackAlloc(d, m->nv);
+
+  // kinetic energy:  0.5 * qvel' * M * qvel
+  mj_mulM(m, d, vec, d->qvel);
+  d->energy[1] = 0.5*mju_dot(vec, d->qvel, m->nv);
+
+  mjFREESTACK;
+}

src/engine/engine_sensor.h

@@ -0,0 +1,50 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_SENSOR_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_SENSOR_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//-------------------------------- sensors ---------------------------------------------------------
+
+// position-dependent sensors
+MJAPI void mj_sensorPos(const mjModel* m, mjData* d);
+
+// velocity-dependent sensors
+MJAPI void mj_sensorVel(const mjModel* m, mjData* d);
+
+// acceleration/force-dependent sensors
+MJAPI void mj_sensorAcc(const mjModel* m, mjData* d);
+
+
+//-------------------------------- energy ----------------------------------------------------------
+
+// position-dependent energy (potential)
+MJAPI void mj_energyPos(const mjModel* m, mjData* d);
+
+// velocity-dependent energy (kinetic)
+MJAPI void mj_energyVel(const mjModel* m, mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_SENSOR_H_

src/engine/engine_setconst.c

@@ -0,0 +1,556 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_setconst.h"
+
+#include <stdio.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_core_constraint.h"
+#include "engine/engine_core_smooth.h"
+#include "engine/engine_forward.h"
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_spatial.h"
+
+// set quantities that depend on qpos0
+static void set0(mjModel* m, mjData* d) {
+  int id, id1, id2, dnum, nv = m->nv;
+  mjtNum A[36] = {0}, pos[3], quat[4];
+  mjMARKSTACK;
+  mjtNum* jac = mj_stackAlloc(d, 6*nv);
+  mjtNum* tmp = mj_stackAlloc(d, 6*nv);
+  int* cammode = 0;
+  int* lightmode = 0;
+
+  // save camera and light mode, set to fixed
+  if (m->ncam) {
+    cammode = (int*) mj_stackAlloc(d, m->ncam);
+    for (int i=0; i<m->ncam; i++) {
+      cammode[i] = m->cam_mode[i];
+      m->cam_mode[i] = mjCAMLIGHT_FIXED;
+    }
+  }
+  if (m->nlight) {
+    lightmode = (int*) mj_stackAlloc(d, m->nlight);
+    for (int i=0; i<m->nlight; i++) {
+      lightmode[i] = m->light_mode[i];
+      m->light_mode[i] = mjCAMLIGHT_FIXED;
+    }
+  }
+
+  // run computations in qpos0
+  mju_copy(d->qpos, m->qpos0, m->nq);
+  mj_kinematics(m, d);
+  mj_comPos(m, d);
+  mj_camlight(m, d);
+  mj_crbSkip(m, d, 0);
+
+  // save dof_M0
+  for (int i=0; i<nv; i++) {
+    m->dof_M0[i] = d->qM[m->dof_Madr[i]];
+  }
+
+  // run remaining computations (factorM needs dof_M0)
+  mj_factorM(m, d);
+  mj_tendon(m, d);
+  mj_transmission(m, d);
+
+  // restore camera and light mode
+  for (int i=0; i<m->ncam; i++) {
+    m->cam_mode[i] = cammode[i];
+  }
+  for (int i=0; i<m->nlight; i++) {
+    m->light_mode[i] = lightmode[i];
+  }
+
+  // set tendon_length0, actuator_length0
+  mju_copy(m->tendon_length0, d->ten_length, m->ntendon);
+  mju_copy(m->actuator_length0, d->actuator_length, m->nu);
+
+  // compute body_invweight0
+  m->body_invweight0[0] = m->body_invweight0[1] = 0.0;
+  for (int i=1; i<m->nbody; i++) {
+    if (nv) {
+      // inverse spatial inertia:  A = J*inv(M)*J'
+      mj_jacBodyCom(m, d, jac, jac+3*nv, i);
+      mj_solveM(m, d, tmp, jac, 6);
+      mju_mulMatMatT(A, jac, tmp, 6, nv, 6);
+    }
+
+    // average diagonal and assign
+    m->body_invweight0[2*i] = (A[0] + A[7] + A[14])/3;
+    m->body_invweight0[2*i+1] = (A[21] + A[28] + A[35])/3;
+  }
+
+  // compute dof_invweight0
+  for (int i=0; i<m->njnt; i++) {
+    id = m->jnt_dofadr[i];
+
+    // get number of components
+    if (m->jnt_type[i]==mjJNT_FREE) {
+      dnum = 6;
+    } else if (m->jnt_type[i]==mjJNT_BALL) {
+      dnum = 3;
+    } else {
+      dnum = 1;
+    }
+
+    // inverse joint inertia:  A = J*inv(M)*J'
+    if (nv) {
+      mju_zero(jac, dnum*nv);
+      for (int j=0; j<dnum; j++) {
+        jac[j*(nv+1) + id] = 1;
+      }
+      mj_solveM(m, d, tmp, jac, dnum);
+      mju_mulMatMatT(A, jac, tmp, dnum, nv, dnum);
+    }
+
+    // average diagonal and assign
+    if (dnum==6) {
+      m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
+                                                          (A[0] + A[7] + A[14])/3;
+      m->dof_invweight0[id+3] = m->dof_invweight0[id+4] = m->dof_invweight0[id+5] =
+                                                            (A[21] + A[28] + A[35])/3;
+    } else if (dnum==3)
+      m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
+                                                          (A[0] + A[4] + A[8])/3;
+    else {
+      m->dof_invweight0[id] = A[0];
+    }
+  }
+
+  // compute tendon_invweight0
+  if (nv) {
+    for (int i=0; i<m->ntendon; i++) {
+      // make dense vector into tmp
+      if (mj_isSparse(m)) {
+        mju_zero(tmp, nv);
+        int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i];
+        for (int j=d->ten_J_rowadr[i]; j<end; j++) {
+          tmp[d->ten_J_colind[j]] = d->ten_J[j];
+        }
+      } else {
+        mju_copy(tmp, d->ten_J+i*nv, nv);
+      }
+
+      // solve into tmp+nv
+      mj_solveM(m, d, tmp+nv, tmp, 1);
+      m->tendon_invweight0[i] = mju_dot(tmp, tmp+nv, nv);
+    }
+
+    // compute actuator_acc0
+    for (int i=0; i<m->nu; i++) {
+      mj_solveM(m, d, tmp, d->actuator_moment+i*nv, 1);
+      m->actuator_acc0[i] = mju_norm(tmp, nv);
+    }
+  } else {
+    for (int i=0; i<m->nu; i++) {
+      m->actuator_acc0[i] = 0;
+    }
+  }
+
+  // compute missing eq_data for body constraints
+  for (int i=0; i<m->neq; i++) {
+    // get ids
+    id1 = m->eq_obj1id[i];
+    id2 = m->eq_obj2id[i];
+
+    // connect constraint
+    if (m->eq_type[i]==mjEQ_CONNECT) {
+      // pos = anchor position in global frame
+      mj_local2Global(d, pos, 0, m->eq_data+mjNEQDATA*i, 0, id1, 0);
+
+      // data[3-5] = anchor position in body2 local frame
+      mju_subFrom3(pos, d->xpos+3*id2);
+      mju_rotVecMatT(m->eq_data+mjNEQDATA*i+3, pos, d->xmat+9*id2);
+    }
+
+    // weld constraint
+    else if (m->eq_type[i]==mjEQ_WELD) {
+      // skip if user has set any quaternion data
+      if (m->eq_data[mjNEQDATA*i+3] ||
+          m->eq_data[mjNEQDATA*i+4] ||
+          m->eq_data[mjNEQDATA*i+5] ||
+          m->eq_data[mjNEQDATA*i+6]) {
+        // normalize quaternion just in case
+        mju_normalize4(m->eq_data+mjNEQDATA*i+3);
+        continue;
+      }
+
+      // data[0-2] = xpos2-xpos1 in body1 local frame
+      mju_sub3(pos, d->xpos+3*id2, d->xpos+3*id1);
+      mju_rotVecMatT(m->eq_data+mjNEQDATA*i, pos, d->xmat+9*id1);
+
+      // data[3-6] = neg(xquat1)*xquat2 = "xquat2-xquat1" in body1 local frame
+      mju_negQuat(quat, d->xquat+4*id1);
+      mju_mulQuat(m->eq_data+mjNEQDATA*i+3, quat, d->xquat+4*id2);
+    }
+  }
+
+  // camera compos0, pos0, mat0
+  for (int i=0; i<m->ncam; i++) {
+    // get body ids
+    id = m->cam_bodyid[i];              // camera body
+    id1 = m->cam_targetbodyid[i];       // target body
+
+    // compute positional offsets
+    mju_sub3(m->cam_pos0+3*i, d->cam_xpos+3*i, d->xpos+3*id);
+    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);
+  }
+
+  // light compos0, pos0, dir0
+  for (int i=0; i<m->nlight; i++) {
+    // get body ids
+    id = m->light_bodyid[i];            // light body
+    id1 = m->light_targetbodyid[i];     // target body
+
+    // compute positional offsets
+    mju_sub3(m->light_pos0+3*i, d->light_xpos+3*i, d->xpos+3*id);
+    mju_sub3(m->light_poscom0+3*i, d->light_xpos+3*i, d->subtree_com+ (id1>=0 ? 3*id1 : 3*id));
+
+    // copy dir
+    mju_copy3(m->light_dir0+3*i, d->light_xdir+3*i);
+  }
+
+  mjFREESTACK;
+}
+
+
+
+// accumulate bounding box
+static void updateBox(mjtNum* xmin, mjtNum* xmax, mjtNum* pos, mjtNum radius) {
+  for (int i=0; i<3; i++) {
+    xmin[i] = mjMIN(xmin[i], pos[i] - radius);
+    xmax[i] = mjMAX(xmax[i], pos[i] + radius);
+  }
+}
+
+
+// compute stat; assume computations already executed in qpos0
+static void setStat(mjModel* m, mjData* d) {
+  mjtNum xmin[3] = {1E+10, 1E+10, 1E+10};
+  mjtNum xmax[3] = {-1E+10, -1E+10, -1E+10};
+  mjtNum rbound;
+  mjMARKSTACK;
+  mjtNum* body = mj_stackAlloc(d, m->nbody);
+
+  // compute bounding box of bodies, joint centers, geoms and sites
+  for (int i=1; i<m->nbody; i++) {
+    updateBox(xmin, xmax, d->xpos+3*i, 0);
+    updateBox(xmin, xmax, d->xipos+3*i, 0);
+  }
+  for (int i=0; i<m->njnt; i++) {
+    updateBox(xmin, xmax, d->xanchor+3*i, 0);
+  }
+  for (int i=0; i<m->nsite; i++) {
+    updateBox(xmin, xmax, d->site_xpos+3*i, 0);
+  }
+  for (int i=0; i<m->ngeom; i++) {
+    // set rbound: regular geom rbound, or 0.1 of plane or hfield max size
+    rbound = 0;
+    if (m->geom_rbound[i] > 0) {
+      rbound = m->geom_rbound[i];
+    } else if (m->geom_type[i]==mjGEOM_PLANE) {
+      // finite in at least one direction
+      if (m->geom_size[3*i] || m->geom_size[3*i+1]) {
+        rbound = mjMAX(m->geom_size[3*i], m->geom_size[3*i+1]) * 0.1;
+      }
+
+      // infinite in both directions
+      else {
+        rbound = 1;
+      }
+    } else if (m->geom_type[i]==mjGEOM_HFIELD) {
+      int j = m->geom_dataid[i];
+      rbound = mjMAX(m->hfield_size[4*j],
+                     mjMAX(m->hfield_size[4*j+1],
+                           mjMAX(m->hfield_size[4*j+2], m->hfield_size[4*j+3]))) * 0.1;
+    }
+
+    updateBox(xmin, xmax, d->geom_xpos+3*i, rbound);
+  }
+
+  // compute center
+  mju_add3(m->stat.center, xmin, xmax);
+  mju_scl3(m->stat.center, m->stat.center, 0.5);
+
+  // compute bounding box size
+  if (xmax[0]>xmin[0])
+    m->stat.extent = mju_max(1E-5,
+                             mju_max(xmax[0]-xmin[0], mju_max(xmax[1]-xmin[1], xmax[2]-xmin[2])));
+
+  // set body size to max com-joint distance
+  mju_zero(body, m->nbody);
+  for (int i=0; i<m->njnt; i++) {
+    // handle this body
+    int id = m->jnt_bodyid[i];
+    body[id] = mju_max(body[id], mju_dist3(d->xipos+3*id, d->xanchor+3*i));
+
+    // handle parent body
+    id = m->body_parentid[id];
+    body[id] = mju_max(body[id], mju_dist3(d->xipos+3*id, d->xanchor+3*i));
+  }
+  body[0] = 0;
+
+  // set body size to max of old value, and geom rbound + com-geom dist
+  for (int i=1; i<m->nbody; i++) {
+    for (int id=m->body_geomadr[i]; id<m->body_geomadr[i]+m->body_geomnum[i]; id++) {
+      if (m->geom_rbound[id]>0) {
+        body[i] = mju_max(body[i], m->geom_rbound[id] + mju_dist3(d->xipos+3*i, d->geom_xpos+3*id));
+      }
+    }
+  }
+
+  // compute meansize, make sure all sizes are above min
+  if (m->nbody>1) {
+    m->stat.meansize = 0;
+    for (int i=1; i<m->nbody; i++) {
+      body[i] = mju_max(body[i], 1E-5);
+      m->stat.meansize += body[i]/(m->nbody-1);
+    }
+  }
+
+  // fix extent if too small compared to meanbody
+  m->stat.extent = mju_max(m->stat.extent, 2 * m->stat.meansize);
+
+  // compute meanmass
+  if (m->nbody>1) {
+    m->stat.meanmass = 0;
+    for (int i=1; i<m->nbody; i++) {
+      m->stat.meanmass += m->body_mass[i];
+    }
+    m->stat.meanmass /= (m->nbody-1);
+  }
+
+  // compute meaninertia
+  if (m->nv) {
+    m->stat.meaninertia = 0;
+    for (int i=0; i<m->nv; i++) {
+      m->stat.meaninertia += d->qM[m->dof_Madr[i]];
+    }
+    m->stat.meaninertia /= m->nv;
+  }
+
+  mjFREESTACK;
+}
+
+
+
+// set quantities that depend on qpos_spring
+static void setSpring(mjModel* m, mjData* d) {
+  // run computations in qpos_spring
+  mju_copy(d->qpos, m->qpos_spring, m->nq);
+  mj_kinematics(m, d);
+  mj_comPos(m, d);
+  mj_tendon(m, d);
+  mj_transmission(m, d);
+
+  // copy if model spring length is negative
+  for (int i=0; i<m->ntendon; i++) {
+    if (m->tendon_lengthspring[i]<0) {
+      m->tendon_lengthspring[i] = d->ten_length[i];
+    }
+  }
+}
+
+
+
+// entry point: set all constant fields of mjModel, except for lengthrange
+void mj_setConst(mjModel* m, mjData* d) {
+  // compute subtreemass
+  for (int i=0; i<m->nbody; i++) {
+    m->body_subtreemass[i] = m->body_mass[i];
+  }
+  for (int i=m->nbody-1; i>0; i--) {
+    m->body_subtreemass[m->body_parentid[i]] += m->body_subtreemass[i];
+  }
+
+  // call functions
+  set0(m, d);
+  setStat(m, d);
+  setSpring(m, d);
+}
+
+
+
+//----------------------------- actuator length range computation ----------------------------------
+
+// evaluate actuator length, advance special dynamics
+static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
+                      const mjLROpt* opt) {
+  int nv = m->nv;
+
+  // reduce velocity
+  mju_scl(d->qvel, d->qvel, mju_exp(-m->opt.timestep/mjMAX(0.01, opt->timeconst)), nv);
+
+  // step1: compute inertia and actuator moments
+  mj_step1(m, d);
+
+  // set force to generate desired acceleration
+  mj_solveM(m, d, d->qfrc_applied, d->actuator_moment+index*nv, 1);
+  mjtNum nrm = mju_norm(d->qfrc_applied, nv);
+  mju_scl(d->qfrc_applied, d->actuator_moment+index*nv,
+          (2*side-1)*opt->accel/mjMAX(mjMINVAL, nrm), nv);
+
+  // impose maxforce
+  nrm = mju_norm(d->qfrc_applied, nv);
+  if (opt->maxforce>0 && nrm>opt->maxforce) {
+    mju_scl(d->qfrc_applied, d->qfrc_applied, opt->maxforce/mjMAX(mjMINVAL, nrm), nv);
+  }
+
+  // step2: apply force
+  mj_step2(m, d);
+
+  // return actuator length
+  return d->actuator_length[index];
+}
+
+
+
+// Set length range for specified actuator, return 1 if ok, 0 if error.
+int mj_setLengthRange(mjModel* m, mjData* d, int index,
+                      const mjLROpt* opt, char* error, int error_sz) {
+  // check index
+  if (index<0 || index>=m->nu) {
+    mju_error("Invalid actuator index in mj_setLengthRange");
+  }
+
+  // skip depending on mode and type
+  int ismuscle = (m->actuator_gaintype[index]==mjGAIN_MUSCLE ||
+                  m->actuator_biastype[index]==mjBIAS_MUSCLE);
+  int isuser = (m->actuator_gaintype[index]==mjGAIN_USER ||
+                m->actuator_biastype[index]==mjBIAS_USER);
+  if ((opt->mode==mjLRMODE_NONE) ||
+      (opt->mode==mjLRMODE_MUSCLE && !ismuscle) ||
+      (opt->mode==mjLRMODE_MUSCLEUSER && !ismuscle && !isuser)) {
+    return 1;
+  }
+
+  // use existing length range if available
+  if (opt->useexisting && (m->actuator_lengthrange[2*index] < m->actuator_lengthrange[2*index+1])) {
+    return 1;
+  }
+
+  // get transmission id
+  int threadid = m->actuator_trnid[index];
+
+  // use joint and tendon limits if available
+  if (opt->uselimit) {
+    // joint or jointinparent
+    if (m->actuator_trntype[index]==mjTRN_JOINT ||
+        m->actuator_trntype[index]==mjTRN_JOINTINPARENT) {
+      // make sure joint is limited
+      if (m->jnt_limited[threadid]) {
+        // copy range
+        m->actuator_lengthrange[2*index] = m->jnt_range[2*threadid];
+        m->actuator_lengthrange[2*index+1] = m->jnt_range[2*threadid+1];
+
+        // skip optimization
+        return 1;
+      }
+    }
+
+    // tendon
+    if (m->actuator_trntype[index]==mjTRN_TENDON) {
+      // make sure tendon is limited
+      if (m->tendon_limited[threadid]) {
+        // copy range
+        m->actuator_lengthrange[2*index] = m->tendon_range[2*threadid];
+        m->actuator_lengthrange[2*index+1] = m->tendon_range[2*threadid+1];
+
+        // skip optimization
+        return 1;
+      }
+    }
+  }
+
+  // optimize in both directions
+  mjtNum lmin[2] = {0, 0}, lmax[2] = {0, 0};
+  int side;
+  for (side=0; side<2; side++) {
+    // init at qpos0
+    mj_resetData(m, d);
+
+    // simulate
+    int updated = 0;
+    while (d->time < opt->inttotal) {
+      // advance and get length
+      mjtNum len = evalAct(m, d, index, side, opt);
+
+      // reset: cannot proceed
+      if (d->time==0) {
+        snprintf(error, error_sz, "Unstable lengthrange simulation in actuator %d", index);
+        return 0;
+      }
+
+      // update limits
+      if (d->time > opt->inttotal-opt->inteval) {
+        if (len<lmin[side] || !updated) {
+          lmin[side] = len;
+        }
+        if (len>lmax[side] || !updated) {
+          lmax[side] = len;
+        }
+
+        updated = 1;
+      }
+    }
+
+    // assign
+    m->actuator_lengthrange[2*index+side] = (side==0 ? lmin[side] : lmax[side]);
+  }
+
+  // check range
+  mjtNum dif = m->actuator_lengthrange[2*index+1] - m->actuator_lengthrange[2*index];
+  if (dif<=0) {
+    snprintf(error, error_sz,
+             "Invalid lengthrange (%g, %g) in actuator %d",
+             m->actuator_lengthrange[2*index],
+             m->actuator_lengthrange[2*index+1], index);
+    return 0;
+  }
+
+  // check convergence, side 0
+  if (lmax[0]-lmin[0]>opt->tolrange*dif) {
+    snprintf(error, error_sz,
+             "Lengthrange computation did not converge in actuator %d:\n"
+             "  eval (%g, %g)\n  range (%g, %g)",
+             index, lmin[0], lmax[0],
+             m->actuator_lengthrange[2*index],
+             m->actuator_lengthrange[2*index+1]);
+    return 0;
+  }
+
+  // check convergence, side 1
+  if (lmax[1]-lmin[1]>opt->tolrange*dif) {
+    snprintf(error, error_sz,
+             "Lengthrange computation did not converge in actuator %d:\n"
+             "  eval (%g, %g)\n range (%g, %g)",
+             index, lmin[1], lmax[1],
+             m->actuator_lengthrange[2*index],
+             m->actuator_lengthrange[2*index+1]);
+    return 0;
+  }
+
+  return 1;
+}

src/engine/engine_setconst.h

@@ -0,0 +1,37 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_SETCONST_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_SETCONST_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Set constant fields of mjModel, corresponding to qpos0 configuration.
+MJAPI void mj_setConst(mjModel* m, mjData* d);
+
+// Set actuator_lengthrange for specified actuator; return 1 if ok, 0 if error.
+MJAPI int mj_setLengthRange(mjModel* m, mjData* d, int index,
+                            const mjLROpt* opt, char* error, int error_sz);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_SETCONST_H_

src/engine/engine_solver.h

@@ -0,0 +1,34 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+
+
+// PGS solver
+void mj_solPGS(const mjModel* m, mjData* d, int maxiter);
+
+// No Slip solver (modified PGS)
+void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter);
+
+// CG solver
+void mj_solCG(const mjModel* m, mjData* d, int maxiter);
+
+// Newton solver
+void mj_solNewton(const mjModel* m, mjData* d, int maxiter);
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_

src/engine/engine_support.h

@@ -0,0 +1,162 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_SUPPORT_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_SUPPORT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// strings
+MJAPI extern const char* mjDISABLESTRING[mjNDISABLE];
+MJAPI extern const char* mjENABLESTRING[mjNENABLE];
+MJAPI extern const char* mjTIMERSTRING[mjNTIMER];
+
+
+//-------------------------- Jacobians -------------------------------------------------------------
+
+// compute 3/6-by-nv Jacobian of global point attached to given body
+MJAPI void mj_jac(const mjModel* m, const mjData* d,
+                  mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
+
+// compute body frame Jacobian
+MJAPI void mj_jacBody(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int body);
+
+// compute body center-of-mass Jacobian
+MJAPI void mj_jacBodyCom(const mjModel* m, const mjData* d,
+                         mjtNum* jacp, mjtNum* jacr, int body);
+
+// compute geom Jacobian
+MJAPI void mj_jacGeom(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int geom);
+
+// compute site Jacobian
+MJAPI void mj_jacSite(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, int site);
+
+// compute translation Jacobian of point, and rotation Jacobian of axis
+MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
+                           mjtNum* jacPoint, mjtNum* jacAxis,
+                           const mjtNum point[3], const mjtNum axis[3], int body);
+
+// compute 3/6-by-nv sparse Jacobian of global point attached to given body
+void mj_jacSparse(const mjModel* m, const mjData* d,
+                  mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                  int NV, int* chain);
+
+// sparse Jacobian difference for simple body contacts
+void mj_jacSparseSimple(const mjModel* m, const mjData* d,
+                        mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
+                        int body, int flg_second, int NV, int start);
+
+// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
+int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
+                  int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
+                  mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
+                  mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
+
+
+//-------------------------- name functions --------------------------------------------------------
+
+// get id of object with specified name; -1: not found; type is mjtObj
+MJAPI int mj_name2id(const mjModel* m, int type, const char* name);
+
+// get name of object with specified id; 0: invalid type or id; type is mjtObj
+MJAPI const char* mj_id2name(const mjModel* m, int type, int id);
+
+
+//-------------------------- inertia functions -----------------------------------------------------
+
+// convert sparse inertia matrix M into full matrix
+MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
+
+// multiply vector by inertia matrix
+MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
+
+// multiply vector by (inertia matrix)^(1/2)
+MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
+
+// add inertia matrix to destination matrix
+//  destination can be sparse uncompressed, or dense when all int* are NULL
+MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
+                   int* rownnz, int* rowadr, int* colind);
+
+
+//-------------------------- perturbations ---------------------------------------------------------
+
+// apply cartesian force and torque
+MJAPI void mj_applyFT(const mjModel* m, mjData* d,
+                      const mjtNum force[3], const mjtNum torque[3],
+                      const mjtNum point[3], int body, mjtNum* qfrc_target);
+
+// accumulate xfrc_applied in qfrc
+void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc);
+
+
+//-------------------------- coordinate transformation ---------------------------------------------
+
+// compute object 6D velocity in object-centered frame, world/local orientation
+MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d,
+                             int objtype, int objid, mjtNum res[6], int flg_local);
+
+// compute object 6D acceleration in object-centered frame, world/local orientation
+MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d,
+                                 int objtype, int objid, mjtNum res[6], int flg_local);
+
+
+//-------------------------- miscellaneous ---------------------------------------------------------
+
+// extract 6D force:torque for one contact, in contact frame
+MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
+
+// compute velocity by finite-differencing two positions
+MJAPI void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
+                               const mjtNum* qpos1, const mjtNum* qpos2);
+
+// integrate position with given velocity
+MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt);
+
+// normalize all quaterions in qpos-type vector
+MJAPI void mj_normalizeQuat(const mjModel* m, mjtNum* qpos);
+
+// map from body local to global Cartesian coordinates
+MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
+                           const mjtNum pos[3], const mjtNum quat[4],
+                           int body, mjtByte sameframe);
+
+// sum all body masses
+MJAPI mjtNum mj_getTotalmass(const mjModel* m);
+
+// scale body masses and inertias to achieve specified total mass
+MJAPI void mj_setTotalmass(mjModel* m, mjtNum newmass);
+
+// high-level warning function: count warnings in mjData, print only the first time
+MJAPI void mj_warning(mjData* d, int warning, int info);
+
+// version number
+MJAPI int mj_version(void);
+
+// current version of MuJoCo as a null-terminated string
+MJAPI const char* mj_versionString();
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_SUPPORT_H_

src/engine/engine_util_blas.c

@@ -0,0 +1,795 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_util_blas.h"
+
+#include <string.h>
+
+#include <mujoco/mjmodel.h>
+
+#ifdef mjUSEPLATFORMSIMD
+  #if defined(__AVX__) && defined(mjUSEDOUBLE)
+    #define mjUSEAVX
+    #include "immintrin.h"
+  #endif
+#endif
+
+
+
+//------------------------------ 3D vector and matrix-vector operations ----------------------------
+
+// res = 0
+void mju_zero3(mjtNum res[3]) {
+  res[0] = 0;
+  res[1] = 0;
+  res[2] = 0;
+}
+
+
+
+// res = vec
+void mju_copy3(mjtNum res[3], const mjtNum data[3]) {
+  res[0] = data[0];
+  res[1] = data[1];
+  res[2] = data[2];
+}
+
+
+
+// res = vec*scl
+void mju_scl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl) {
+  res[0] = vec[0] * scl;
+  res[1] = vec[1] * scl;
+  res[2] = vec[2] * scl;
+}
+
+
+
+// res = vec1 + vec2
+void mju_add3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]) {
+  res[0] = vec1[0] + vec2[0];
+  res[1] = vec1[1] + vec2[1];
+  res[2] = vec1[2] + vec2[2];
+}
+
+
+
+// res = vec1 - vec2
+void mju_sub3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]) {
+  res[0] = vec1[0] - vec2[0];
+  res[1] = vec1[1] - vec2[1];
+  res[2] = vec1[2] - vec2[2];
+}
+
+
+
+// res += vec
+void mju_addTo3(mjtNum res[3], const mjtNum vec[3]) {
+  res[0] += vec[0];
+  res[1] += vec[1];
+  res[2] += vec[2];
+}
+
+
+
+// res -= vec
+void mju_subFrom3(mjtNum res[3], const mjtNum vec[3]) {
+  res[0] -= vec[0];
+  res[1] -= vec[1];
+  res[2] -= vec[2];
+}
+
+
+
+// res += vec*scl
+void mju_addToScl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl) {
+  res[0] += vec[0] * scl;
+  res[1] += vec[1] * scl;
+  res[2] += vec[2] * scl;
+}
+
+
+
+// res = vec1 + vec2*scl
+void mju_addScl3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3], mjtNum scl) {
+  res[0] = vec1[0] + scl*vec2[0];
+  res[1] = vec1[1] + scl*vec2[1];
+  res[2] = vec1[2] + scl*vec2[2];
+}
+
+
+
+// normalize vector, return length before normalization
+mjtNum mju_normalize3(mjtNum vec[3]) {
+  mjtNum norm = mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2]);
+
+  if (norm<mjMINVAL) {
+    vec[0] = 1;
+    vec[1] = 0;
+    vec[2] = 0;
+  } else {
+    mjtNum normInv = 1/norm;
+    vec[0] *= normInv;
+    vec[1] *= normInv;
+    vec[2] *= normInv;
+  }
+
+  return norm;
+}
+
+
+
+// compute vector length (without normalizing)
+mjtNum mju_norm3(const mjtNum vec[3]) {
+  return mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2]);
+}
+
+
+
+// vector dot-product
+mjtNum mju_dot3(const mjtNum vec1[3], const mjtNum vec2[3]) {
+  return vec1[0]*vec2[0] + vec1[1]*vec2[1] + vec1[2]*vec2[2];
+}
+
+
+
+// Cartesian distance between 3D vectors
+mjtNum mju_dist3(const mjtNum pos1[3], const mjtNum pos2[3]) {
+  mjtNum dif[3] = {pos1[0]-pos2[0], pos1[1]-pos2[1], pos1[2]-pos2[2]};
+  return mju_sqrt(dif[0]*dif[0] + dif[1]*dif[1] + dif[2]*dif[2]);
+}
+
+
+
+// multiply vector by 3D rotation matrix
+void mju_rotVecMat(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]) {
+  res[0] = mat[0]*vec[0] + mat[1]*vec[1] + mat[2]*vec[2];
+  res[1] = mat[3]*vec[0] + mat[4]*vec[1] + mat[5]*vec[2];
+  res[2] = mat[6]*vec[0] + mat[7]*vec[1] + mat[8]*vec[2];
+}
+
+
+
+// multiply vector by transposed 3D rotation matrix
+void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]) {
+  res[0] = mat[0]*vec[0] + mat[3]*vec[1] + mat[6]*vec[2];
+  res[1] = mat[1]*vec[0] + mat[4]*vec[1] + mat[7]*vec[2];
+  res[2] = mat[2]*vec[0] + mat[5]*vec[1] + mat[8]*vec[2];
+}
+
+
+
+//------------------------------ 4D vector and matrix-vector operations ----------------------------
+
+// res = 0
+void mju_zero4(mjtNum res[4]) {
+  res[0] = 0;
+  res[1] = 0;
+  res[2] = 0;
+  res[3] = 0;
+}
+
+
+
+// res = (1,0,0,0)
+void mju_unit4(mjtNum res[4]) {
+  res[0] = 1;
+  res[1] = 0;
+  res[2] = 0;
+  res[3] = 0;
+}
+
+
+// res = vec
+void mju_copy4(mjtNum res[4], const mjtNum data[4]) {
+  res[0] = data[0];
+  res[1] = data[1];
+  res[2] = data[2];
+  res[3] = data[3];
+}
+
+
+
+// normalize vector, return length before normalization
+mjtNum mju_normalize4(mjtNum vec[4]) {
+  mjtNum norm = mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2] + vec[3]*vec[3]);
+
+  if (norm<mjMINVAL) {
+    vec[0] = 1;
+    vec[1] = 0;
+    vec[2] = 0;
+    vec[3] = 0;
+  } else {
+    mjtNum normInv = 1/norm;
+    vec[0] *= normInv;
+    vec[1] *= normInv;
+    vec[2] *= normInv;
+    vec[3] *= normInv;
+  }
+
+  return norm;
+}
+
+
+
+//------------------------------ vector operations -------------------------------------------------
+
+// res = 0
+void mju_zero(mjtNum* res, int n) {
+  if (n>0) {
+    memset(res, 0, n*sizeof(mjtNum));
+  }
+}
+
+
+
+// res = vec
+void mju_copy(mjtNum* res, const mjtNum* vec, int n) {
+  if (n>0) {
+    memcpy(res, vec, n*sizeof(mjtNum));
+  }
+}
+
+
+
+// sum(vec)
+mjtNum mju_sum(const mjtNum* vec, int n) {
+  mjtNum res = 0;
+
+  for (int i=0; i<n; i++) {
+    res += vec[i];
+  }
+
+  return res;
+}
+
+
+
+// sum(abs(vec))
+mjtNum mju_L1(const mjtNum* vec, int n) {
+  mjtNum res = 0;
+
+  for (int i=0; i<n; i++) {
+    res += mju_abs(vec[i]);
+  }
+
+  return res;
+}
+
+
+
+// res = vec*scl
+void mju_scl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sclpar, val1, val1scl;
+
+    // init
+    sclpar = _mm256_set1_pd(scl);
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(vec+i);
+      val1scl = _mm256_mul_pd(val1, sclpar);
+      _mm256_storeu_pd(res+i, val1scl);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] = vec[i]*scl;
+    res[i+1] = vec[i+1]*scl;
+    res[i+2] = vec[i+2]*scl;
+  } else if (n_i==2) {
+    res[i] = vec[i]*scl;
+    res[i+1] = vec[i+1]*scl;
+  } else if (n_i==1) {
+    res[i] = vec[i]*scl;
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] = vec[i]*scl;
+  }
+#endif
+}
+
+
+
+// res = vec1 + vec2
+void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sum, val1, val2;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(vec1+i);
+      val2 = _mm256_loadu_pd(vec2+i);
+      sum = _mm256_add_pd(val1, val2);
+      _mm256_storeu_pd(res+i, sum);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] = vec1[i] + vec2[i];
+    res[i+1] = vec1[i+1] + vec2[i+1];
+    res[i+2] = vec1[i+2] + vec2[i+2];
+  } else if (n_i==2) {
+    res[i] = vec1[i] + vec2[i];
+    res[i+1] = vec1[i+1] + vec2[i+1];
+  } else if (n_i==1) {
+    res[i] = vec1[i] + vec2[i];
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] = vec1[i] + vec2[i];
+  }
+#endif
+}
+
+
+
+// res = vec1 - vec2
+void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d dif, val1, val2;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(vec1+i);
+      val2 = _mm256_loadu_pd(vec2+i);
+      dif = _mm256_sub_pd(val1, val2);
+      _mm256_storeu_pd(res+i, dif);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] = vec1[i] - vec2[i];
+    res[i+1] = vec1[i+1] - vec2[i+1];
+    res[i+2] = vec1[i+2] - vec2[i+2];
+  } else if (n_i==2) {
+    res[i] = vec1[i] - vec2[i];
+    res[i+1] = vec1[i+1] - vec2[i+1];
+  } else if (n_i==1) {
+    res[i] = vec1[i] - vec2[i];
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] = vec1[i] - vec2[i];
+  }
+#endif
+}
+
+
+
+// res += vec
+void mju_addTo(mjtNum* res, const mjtNum* vec, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sum, val1, val2;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(res+i);
+      val2 = _mm256_loadu_pd(vec+i);
+      sum = _mm256_add_pd(val1, val2);
+      _mm256_storeu_pd(res+i, sum);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] += vec[i];
+    res[i+1] += vec[i+1];
+    res[i+2] += vec[i+2];
+  } else if (n_i==2) {
+    res[i] += vec[i];
+    res[i+1] += vec[i+1];
+  } else if (n_i==1) {
+    res[i] += vec[i];
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] += vec[i];
+  }
+#endif
+}
+
+
+
+// res -= vec
+void mju_subFrom(mjtNum* res, const mjtNum* vec, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d dif, val1, val2;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(res+i);
+      val2 = _mm256_loadu_pd(vec+i);
+      dif = _mm256_sub_pd(val1, val2);
+      _mm256_storeu_pd(res+i, dif);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] -= vec[i];
+    res[i+1] -= vec[i+1];
+    res[i+2] -= vec[i+2];
+  } else if (n_i==2) {
+    res[i] -= vec[i];
+    res[i+1] -= vec[i+1];
+  } else if (n_i==1) {
+    res[i] -= vec[i];
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] -= vec[i];
+  }
+#endif
+}
+
+
+
+// res += vec*scl
+void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sclpar, sum, val1, val2, val2scl;
+
+    // init
+    sclpar = _mm256_set1_pd(scl);
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(res+i);
+      val2 = _mm256_loadu_pd(vec+i);
+      val2scl = _mm256_mul_pd(val2, sclpar);
+      sum = _mm256_add_pd(val1, val2scl);
+      _mm256_storeu_pd(res+i, sum);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] += vec[i]*scl;
+    res[i+1] += vec[i+1]*scl;
+    res[i+2] += vec[i+2]*scl;
+  } else if (n_i==2) {
+    res[i] += vec[i]*scl;
+    res[i+1] += vec[i+1]*scl;
+  } else if (n_i==1) {
+    res[i] += vec[i]*scl;
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] += vec[i]*scl;
+  }
+#endif
+}
+
+// res = vec1 + vec2*scl
+void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n) {
+  int i = 0;
+
+#if defined(__AVX__) && defined(mjUSEAVX)  && defined(mjUSEDOUBLE)
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sclpar, sum, val1, val2, val2scl;
+
+    // init
+    sclpar = _mm256_set1_pd(scl);
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(vec1+i);
+      val2 = _mm256_loadu_pd(vec2+i);
+      val2scl = _mm256_mul_pd(val2, sclpar);
+      sum = _mm256_add_pd(val1, val2scl);
+      _mm256_storeu_pd(res+i, sum);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] = vec1[i] + vec2[i]*scl;
+    res[i+1] = vec1[i+1] + vec2[i+1]*scl;
+    res[i+2] = vec1[i+2] + vec2[i+2]*scl;
+  } else if (n_i==2) {
+    res[i] = vec1[i] + vec2[i]*scl;
+    res[i+1] = vec1[i+1] + vec2[i+1]*scl;
+  } else if (n_i==1) {
+    res[i] = vec1[i] + vec2[i]*scl;
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] = vec1[i] + vec2[i]*scl;
+  }
+#endif
+}
+
+
+
+// normalize vector, return length before normalization
+mjtNum mju_normalize(mjtNum* res, int n) {
+  mjtNum norm = (mjtNum)mju_sqrt(mju_dot(res, res, n));
+  mjtNum normInv;
+
+  if (norm<mjMINVAL) {
+    res[0] = 1;
+    for (int i=1; i<n; i++) {
+      res[i] = 0;
+    }
+  } else {
+    normInv = 1/norm;
+    for (int i=0; i<n; i++) {
+      res[i] *= normInv;
+    }
+  }
+
+  return norm;
+}
+
+
+
+// compute vector length (without normalizing)
+mjtNum mju_norm(const mjtNum* res, int n) {
+  return mju_sqrt(mju_dot(res, res, n));
+}
+
+
+
+// vector dot-product
+mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, const int n) {
+  mjtNum res = 0;
+  int i = 0;
+  int n_4 = n - 4;
+#ifdef mjUSEAVX
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sum, prod, val1, val2;
+    __m128d vlow, vhigh, high64;
+
+    // init
+    val1 = _mm256_loadu_pd(vec1);
+    val2 = _mm256_loadu_pd(vec2);
+    sum = _mm256_mul_pd(val1, val2);
+    i = 4;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(vec1+i);
+      val2 = _mm256_loadu_pd(vec2+i);
+      prod = _mm256_mul_pd(val1, val2);
+      sum = _mm256_add_pd(sum, prod);
+      i += 4;
+    }
+
+    // reduce
+    vlow = _mm256_castpd256_pd128(sum);
+    vhigh = _mm256_extractf128_pd(sum, 1);
+    vlow = _mm_add_pd(vlow, vhigh);
+    high64 = _mm_unpackhi_pd(vlow, vlow);
+    res = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
+  }
+
+#else
+  // do the same order of additions as the AVX intrinsics implementation.
+  // this is faster than the simple for loop you'd expect for a dot product,
+  // and produces exactly the same results.
+  mjtNum res0 = 0;
+  mjtNum res1 = 0;
+  mjtNum res2 = 0;
+  mjtNum res3 = 0;
+
+  for (; i<=n_4; i+=4) {
+    res0 += vec1[i] * vec2[i];
+    res1 += vec1[i+1] * vec2[i+1];
+    res2 += vec1[i+2] * vec2[i+2];
+    res3 += vec1[i+3] * vec2[i+3];
+  }
+  res = (res0 + res2) + (res1 + res3);
+#endif
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res += vec1[i]*vec2[i] + vec1[i+1]*vec2[i+1] + vec1[i+2]*vec2[i+2];
+  } else if (n_i==2) {
+    res += vec1[i]*vec2[i] + vec1[i+1]*vec2[i+1];
+  } else if (n_i==1) {
+    res += vec1[i]*vec2[i];
+  }
+  return res;
+}
+
+//------------------------------ matrix-vector operations ------------------------------------------
+
+// multiply matrix and vector
+void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                   int nr, int nc) {
+  for (int r=0; r<nr; r++) {
+    res[r] = mju_dot(mat + r*nc, vec, nc);
+  }
+}
+
+
+
+// multiply transposed matrix and vector
+void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                    int nr, int nc) {
+  mjtNum tmp;
+  mju_zero(res, nc);
+
+  for (int r=0; r<nr; r++) {
+    if ((tmp = vec[r])) {
+      mju_addToScl(res, mat+r*nc, tmp, nc);
+    }
+  }
+}
+
+
+
+//------------------------------ matrix-matrix operations ------------------------------------------
+
+// transpose matrix
+void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) {
+  for (int i=0; i<nr; i++) {
+    for (int j=0; j<nc; j++) {
+      res[j*nr+i] = mat[i*nc+j];
+    }
+  }
+}
+
+
+
+// multiply matrices, exploit sparsity of mat1
+void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                   int r1, int c1, int c2) {
+  mjtNum tmp;
+
+  mju_zero(res, r1*c2);
+
+  for (int i=0; i<r1; i++) {
+    for (int k=0; k<c1; k++) {
+      if ((tmp = mat1[i*c1+k])) {
+        mju_addToScl(res+i*c2, mat2+k*c2, tmp, c2);
+      }
+    }
+  }
+}
+
+
+
+// multiply matrices, second argument transposed
+void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                    int r1, int c1, int r2) {
+  for (int i=0; i<r1; i++) {
+    for (int j=0; j<r2; j++) {
+      res[i*r2+j] = mju_dot(mat1+i*c1, mat2+j*c1, c1);
+    }
+  }
+}
+
+
+
+// compute M'*diag*M (diag=NULL: compute M'*M)
+void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
+  mjtNum tmp;
+
+  // half of MatMat routine: only lower triangle
+  mju_zero(res, nc*nc);
+  if (diag) {
+    for (int j=0; j<nr; j++) {
+      if (diag[j]) {
+        for (int i=0; i<nc; i++) {
+          if ((tmp = mat[j*nc+i])) {
+            mju_addToScl(res+i*nc, mat+j*nc, tmp*diag[j], i+1);
+          }
+        }
+      }
+    }
+  } else {
+    for (int i=0; i<nc; i++) {
+      for (int j=0; j<nr; j++) {
+        if ((tmp = mat[j*nc+i])) {
+          mju_addToScl(res+i*nc, mat+j*nc, tmp, i+1);
+        }
+      }
+    }
+  }
+
+  // make symmetric
+  for (int i=0; i<nc; i++) {
+    for (int j=i+1; j<nc; j++) {
+      res[i*nc+j] = res[j*nc+i];
+    }
+  }
+}
+
+
+
+// multiply matrices, first argument transposed
+void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                    int r1, int c1, int c2) {
+  mjtNum tmp;
+
+  mju_zero(res, c1*c2);
+
+  for (int i=0; i<r1; i++) {
+    for (int j=0; j<c1; j++) {
+      if ((tmp = mat1[i*c1+j])) {
+        mju_addToScl(res+j*c2, mat2+i*c2, tmp, c2);
+      }
+    }
+  }
+}

src/engine/engine_util_blas.h

@@ -0,0 +1,207 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_BLAS_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_BLAS_H_
+
+#include <math.h>
+
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//------------------------------ standard library fuctions -----------------------------------------
+
+#ifdef mjUSEDOUBLE
+  #define mju_sqrt    sqrt
+  #define mju_exp     exp
+  #define mju_sin     sin
+  #define mju_cos     cos
+  #define mju_tan     tan
+  #define mju_asin    asin
+  #define mju_acos    acos
+  #define mju_atan2   atan2
+  #define mju_tanh    tanh
+  #define mju_pow     pow
+  #define mju_abs     fabs
+  #define mju_log     log
+  #define mju_log10   log10
+  #define mju_floor   floor
+  #define mju_ceil    ceil
+
+#else
+  #define mju_sqrt    sqrtf
+  #define mju_exp     expf
+  #define mju_sin     sinf
+  #define mju_cos     cosf
+  #define mju_tan     tanf
+  #define mju_asin    asinf
+  #define mju_acos    acosf
+  #define mju_atan2   atan2f
+  #define mju_tanh    tanhf
+  #define mju_pow     powf
+  #define mju_abs     fabsf
+  #define mju_log     logf
+  #define mju_log10   log10f
+  #define mju_floor   floorf
+  #define mju_ceil    ceilf
+#endif
+
+
+//------------------------------ 3D vector and matrix-vector operations ----------------------------
+
+// res = 0
+MJAPI void mju_zero3(mjtNum res[3]);
+
+// res = vec
+MJAPI void mju_copy3(mjtNum res[3], const mjtNum data[3]);
+
+// res = vec*scl
+MJAPI void mju_scl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl);
+
+// res = vec1 + vec2
+MJAPI void mju_add3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]);
+
+// res = vec1 - vec2
+MJAPI void mju_sub3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]);
+
+// res += vec
+MJAPI void mju_addTo3(mjtNum res[3], const mjtNum vec[3]);
+
+// res -= vec
+MJAPI void mju_subFrom3(mjtNum res[3], const mjtNum vec[3]);
+
+// res += vec*scl
+MJAPI void mju_addToScl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl);
+
+// res = vec1 + vec2*scl
+MJAPI void mju_addScl3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3], mjtNum scl);
+
+// normalize vector, return length before normalization
+MJAPI mjtNum mju_normalize3(mjtNum vec[3]);
+
+// compute vector length (without normalizing)
+MJAPI mjtNum mju_norm3(const mjtNum vec[3]);
+
+// vector dot-product
+MJAPI mjtNum mju_dot3(const mjtNum vec1[3], const mjtNum vec2[3]);
+
+// Cartesian distance between 3D vectors
+MJAPI mjtNum mju_dist3(const mjtNum pos1[3], const mjtNum pos2[3]);
+
+// multiply vector by 3D rotation matrix
+MJAPI void mju_rotVecMat(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]);
+
+// multiply vector by transposed 3D rotation matrix
+MJAPI void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]);
+
+
+//------------------------------ 4D/quaternion operations ------------------------------------------
+
+// res = 0
+MJAPI void mju_zero4(mjtNum res[4]);
+
+// res = (1,0,0,0)
+MJAPI void mju_unit4(mjtNum res[4]);
+
+// res = vec
+MJAPI void mju_copy4(mjtNum res[4], const mjtNum data[4]);
+
+// normalize vector, return length before normalization
+MJAPI mjtNum mju_normalize4(mjtNum vec[4]);
+
+
+//------------------------------ general vector operations -----------------------------------------
+
+// res = 0
+MJAPI void mju_zero(mjtNum* res, int n);
+
+// res = vec
+MJAPI void mju_copy(mjtNum* res, const mjtNum* vec, int n);
+
+// sum(vec)
+MJAPI mjtNum mju_sum(const mjtNum* vec, int n);
+
+// sum(abs(vec))
+MJAPI mjtNum mju_L1(const mjtNum* vec, int n);
+
+// res = vec*scl
+MJAPI void mju_scl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n);
+
+// res = vec1 + vec2
+MJAPI void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n);
+
+// res = vec1 - vec2
+MJAPI void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n);
+
+// res += vec
+MJAPI void mju_addTo(mjtNum* res, const mjtNum* vec, int n);
+
+// res -= vec
+MJAPI void mju_subFrom(mjtNum* res, const mjtNum* vec, int n);
+
+// res += vec*scl
+MJAPI void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n);
+
+// res = vec1 + vec2*scl
+MJAPI void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n);
+
+// normalize vector, return length before normalization
+MJAPI mjtNum mju_normalize(mjtNum* res, int n);
+
+// compute vector length (without normalizing)
+MJAPI mjtNum mju_norm(const mjtNum* res, int n);
+
+// vector dot-product
+MJAPI mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, const int n);
+
+
+//------------------------------ matrix-vector operations ------------------------------------------
+
+// multiply matrix and vector
+MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                         int nr, int nc);
+
+// multiply transposed matrix and vector
+MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                          int nr, int nc);
+
+
+//------------------------------ matrix-matrix operations ------------------------------------------
+
+// transpose matrix
+MJAPI void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc);
+
+// multiply matrices
+MJAPI void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                         int r1, int c1, int c2);
+
+// multiply matrices, second argument transposed
+MJAPI void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                          int r1, int c1, int r2);
+
+// multiply matrices, first argument transposed
+MJAPI void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                          int r1, int c1, int c2);
+
+// compute M'*diag*M (diag=NULL: compute M'*M)
+MJAPI void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_BLAS_H_

src/engine/engine_util_errmem.c

@@ -0,0 +1,232 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_util_errmem.h"
+
+#include <errno.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+
+#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
+#include <unistd.h>
+#endif
+
+//------------------------- cross-platform aligned malloc/free -------------------------------------
+
+static inline void* mju_alignedMalloc(size_t size, size_t align) {
+#ifdef _WIN32
+  return _aligned_malloc(size, align);
+#elif defined(_POSIX_VERSION) && _POSIX_VERSION >= 200112L
+  // Prefer posix_memalign since C11 aligned_alloc isn't available on macOS < 10.15.
+  void* ptr;
+  const int err = posix_memalign(&ptr, align, size);
+  if (err) {
+    ptr = NULL;
+    errno = err;
+  }
+  return ptr;
+#elif defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L
+  return aligned_alloc(align, size);
+#endif
+}
+
+static inline void mju_alignedFree(void* ptr) {
+#ifdef _WIN32
+  _aligned_free(ptr);
+#else
+  free(ptr);
+#endif
+}
+
+
+//------------------------- default user handlers --------------------------------------------------
+
+// define and clear handlers
+void (*mju_user_error) (const char*) = 0;
+void (*mju_user_warning) (const char*) = 0;
+void* (*mju_user_malloc) (size_t) = 0;
+void (*mju_user_free) (void*) = 0;
+
+
+// restore default processing
+void mju_clearHandlers(void) {
+  mju_user_error = 0;
+  mju_user_warning = 0;
+  mju_user_malloc = 0;
+  mju_user_free = 0;
+}
+
+//------------------------- internal-only handlers -------------------------------------------------
+
+typedef void (*callback_fn)(const char*);
+
+#ifdef _MSC_VER
+#define mjTHREADLOCAL __declspec(thread)
+#else
+#define mjTHREADLOCAL _Thread_local
+#endif
+static mjTHREADLOCAL callback_fn _mjPRIVATE_tls_error_fn = NULL;
+static mjTHREADLOCAL callback_fn _mjPRIVATE_tls_warning_fn = NULL;
+
+callback_fn _mjPRIVATE__get_tls_error_fn() {
+  return _mjPRIVATE_tls_error_fn;
+}
+
+void _mjPRIVATE__set_tls_error_fn(callback_fn h) {
+  _mjPRIVATE_tls_error_fn = h;
+}
+
+callback_fn _mjPRIVATE__get_tls_warning_fn() {
+  return _mjPRIVATE_tls_warning_fn;
+}
+
+void _mjPRIVATE__set_tls_warning_fn(callback_fn h) {
+  _mjPRIVATE_tls_warning_fn = h;
+}
+
+//------------------------------ error hadling -----------------------------------------------------
+
+// write datetime, type: message to MUJOCO_LOG.TXT
+void mju_writeLog(const char* type, const char* msg) {
+  time_t rawtime;
+  struct tm *timeinfo;
+  FILE* fp = fopen("MUJOCO_LOG.TXT", "a+t");
+  if (fp) {
+    // get time
+    time(&rawtime);
+    timeinfo = localtime(&rawtime);
+
+    // write to log file
+    fprintf(fp, "%s%s: %s\n\n", asctime(timeinfo), type, msg);
+    fclose(fp);
+  }
+}
+
+
+// write message to logfile and console, pause and exit
+void mju_error(const char* msg) {
+  if (_mjPRIVATE_tls_error_fn) {
+    _mjPRIVATE_tls_error_fn(msg);
+  } else if (mju_user_error) {
+    mju_user_error(msg);
+  } else {
+     // write to log and console
+    mju_writeLog("ERROR", msg);
+    printf("ERROR: %s\n\nPress Enter to exit ...", msg);
+
+    // pause, exit
+    getchar();
+    exit(1);
+  }
+}
+
+
+// write message to logfile and console
+void mju_warning(const char* msg) {
+  if (_mjPRIVATE_tls_warning_fn) {
+    _mjPRIVATE_tls_warning_fn(msg);
+  } else if (mju_user_warning) {
+    mju_user_warning(msg);
+  } else {
+    // write to log file and console
+    mju_writeLog("WARNING", msg);
+    printf("WARNING: %s\n\n", msg);
+  }
+}
+
+
+// error with int argument
+void mju_error_i(const char* msg, int i) {
+  char errmsg[1000];
+  snprintf(errmsg, 1000, msg, i);
+  errmsg[999] = '\0';
+  mju_error(errmsg);
+}
+
+
+// warning with int argument
+void mju_warning_i(const char* msg, int i) {
+  char wrnmsg[1000];
+  snprintf(wrnmsg, 1000, msg, i);
+  wrnmsg[999] = '\0';
+  mju_warning(wrnmsg);
+}
+
+
+// error string argument
+void mju_error_s(const char* msg, const char* text) {
+  char errmsg[1000];
+  snprintf(errmsg, 1000, msg, text);
+  errmsg[999] = '\0';
+  mju_error(errmsg);
+}
+
+
+// warning string argument
+void mju_warning_s(const char* msg, const char* text) {
+  char wrnmsg[1000];
+  snprintf(wrnmsg, 1000, msg, text);
+  wrnmsg[999] = '\0';
+  mju_warning(wrnmsg);
+}
+
+
+
+//------------------------------ malloc and free ---------------------------------------------------
+
+// allocate memory; byte-align on 8; pad size to multiple of 8
+void* mju_malloc(size_t size) {
+  void* ptr = 0;
+
+  // user allocator
+  if (mju_user_malloc) {
+    ptr = mju_user_malloc(size);
+  }
+
+  // default allocator
+  else {
+    // pad size to multiple of 8
+    if ((size%8)) {
+      size += 8 - (size%8);
+    }
+
+    // allocate
+    ptr = mju_alignedMalloc(size, 8);
+  }
+
+  // error if null pointer
+  if (!ptr) {
+    mju_error("Could not allocate memory");
+  }
+
+  return ptr;
+}
+
+
+// free memory
+void mju_free(void* ptr) {
+  // return if null
+  if (!ptr) {
+    return;
+  }
+
+  // free with user or built-in function
+  if (mju_user_free) {
+    mju_user_free(ptr);
+  } else {
+    mju_alignedFree(ptr);
+  }
+}

src/engine/engine_util_errmem.h

@@ -0,0 +1,69 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_ERRMEM_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_ERRMEM_H_
+
+#include <stddef.h>
+
+#include <mujoco/mjexport.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//------------------------------ user handlers -----------------------------------------------------
+
+MJAPI extern void (*mju_user_error)(const char*);
+MJAPI extern void (*mju_user_warning)(const char*);
+MJAPI extern void* (*mju_user_malloc)(size_t);
+MJAPI extern void (*mju_user_free)(void*);
+
+// clear user handlers; restore default processing
+MJAPI void mju_clearHandlers(void);
+
+// gets/sets thread-local error/warning handlers for internal use
+MJAPI void (*_mjPRIVATE__get_tls_error_fn(void))(const char*);
+MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
+MJAPI void (*_mjPRIVATE__get_tls_warning_fn(void))(const char*);
+MJAPI void _mjPRIVATE__set_tls_warning_fn(void (*h)(const char*));
+
+//------------------------------ errors and warnings -----------------------------------------------
+
+// errors
+MJAPI void mju_error(const char* msg);
+MJAPI void mju_error_i(const char* msg, int i);
+MJAPI void mju_error_s(const char* msg, const char* text);
+
+// warnings
+MJAPI void mju_warning(const char* msg);
+MJAPI void mju_warning_i(const char* msg, int i);
+MJAPI void mju_warning_s(const char* msg, const char* text);
+
+// write [datetime, type: message] to MUJOCO_LOG.TXT
+MJAPI void mju_writeLog(const char* type, const char* msg);
+
+
+//------------------------------ malloc and free ---------------------------------------------------
+
+// allocate memory; byte-align on 8; pad size to multiple of 8
+MJAPI void* mju_malloc(size_t size);
+
+// free memory with free() by default
+MJAPI void mju_free(void* ptr);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_ERRMEM_H_

src/engine/engine_util_misc.h

@@ -0,0 +1,126 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_MISC_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_MISC_H_
+
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//------------------------------ tendons and actuators ---------------------------------------------
+
+// wrap tendons around spheres and cylinders
+mjtNum mju_wrap(mjtNum* wpnt, const mjtNum* x0, const mjtNum* x1,
+                const mjtNum* xpos, const mjtNum* xmat, const mjtNum* size,
+                int type, const mjtNum* side);
+
+// muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax)
+MJAPI mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2],
+                            mjtNum acc0, const mjtNum prm[9]);
+
+// muscle passive force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax)
+MJAPI mjtNum mju_muscleBias(mjtNum len, const mjtNum lengthrange[2],
+                            mjtNum acc0, const mjtNum prm[9]);
+
+// muscle activation dynamics, prm = (tau_act, tau_deact)
+MJAPI mjtNum mju_muscleDynamics(mjtNum ctrl, mjtNum act, const mjtNum prm[2]);
+
+
+
+//------------------------------ misclellaneous ----------------------------------------------------
+
+// convert contact force to pyramid representation
+MJAPI void mju_encodePyramid(mjtNum* pyramid, const mjtNum* force,
+                             const mjtNum* mu, int dim);
+
+// convert pyramid representation to contact force
+MJAPI void mju_decodePyramid(mjtNum* force, const mjtNum* pyramid,
+                             const mjtNum* mu, int dim);
+
+// integrate spring-damper analytically, return pos(dt)
+MJAPI mjtNum mju_springDamper(mjtNum pos0, mjtNum vel0, mjtNum Kp, mjtNum Kv, mjtNum dt);
+
+// print matrix
+MJAPI void mju_printMat(const mjtNum* mat, int nr, int nc);
+
+// print sparse matrix to screen
+MJAPI void mju_printMatSparse(const mjtNum* mat, int nr,
+                              const int* rownnz, const int* rowadr,
+                              const int* colind);
+
+// min function, single evaluation of a and b
+MJAPI mjtNum mju_min(mjtNum a, mjtNum b);
+
+// max function, single evaluation of a and b
+MJAPI mjtNum mju_max(mjtNum a, mjtNum b);
+
+// sign function
+MJAPI mjtNum mju_sign(mjtNum x);
+
+// round to nearest integer
+MJAPI int mju_round(mjtNum x);
+
+// convert type id (mjtObj) to type name
+MJAPI const char* mju_type2Str(int type);
+
+// convert type name to type id (mjtObj)
+MJAPI int mju_str2Type(const char* str);
+
+// warning text
+MJAPI const char* mju_warningText(int warning, int info);
+
+// return 1 if nan or abs(x)>mjMAXVAL, 0 otherwise
+MJAPI int mju_isBad(mjtNum x);
+
+// return 1 if all elements are 0
+MJAPI int mju_isZero(mjtNum* vec, int n);
+
+// standard normal random number generator (optional second number)
+MJAPI mjtNum mju_standardNormal(mjtNum* num2);
+
+// convert from float to mjtNum
+MJAPI void mju_f2n(mjtNum* res, const float* vec, int n);
+
+// convert from mjtNum to float
+MJAPI void mju_n2f(float* res, const mjtNum* vec, int n);
+
+// convert from double to mjtNum
+MJAPI void mju_d2n(mjtNum* res, const double* vec, int n);
+
+// convert from mjtNum to double
+MJAPI void mju_n2d(double* res, const mjtNum* vec, int n);
+
+// insertion sort, increasing order
+MJAPI void mju_insertionSort(mjtNum* list, int n);
+
+// integer insertion sort, increasing order
+MJAPI void mju_insertionSortInt(int* list, int n);
+
+// Halton sequence
+MJAPI mjtNum mju_Halton(int index, int base);
+
+// Call strncpy, then set dst[n-1] = 0.
+MJAPI char* mju_strncpy(char *dst, const char *src, int n);
+
+// Sigmoid function over 0<=x<=1 constructed from half-quadratics.
+MJAPI mjtNum mju_sigmoid(mjtNum x);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_MISC_H_

src/engine/engine_util_solve.c

@@ -0,0 +1,632 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_util_solve.h"
+
+#include <math.h>
+#include <string.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_sparse.h"
+#include "engine/engine_util_spatial.h"
+
+//---------------------------- dense Cholesky ------------------------------------------------------
+
+// Cholesky decomposition: mat = L*L'; return 'rank'
+int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag) {
+  int rank = n;
+  mjtNum tmp;
+
+  // in-place Cholesky factorization
+  for (int j=0; j<n; j++) {
+    // compute new diagonal
+    tmp = mat[j*(n+1)];
+    if (j) {
+      tmp -= mju_dot(mat+j*n, mat+j*n, j);
+    }
+
+    // correct diagonal values below threshold
+    if (tmp<mindiag) {
+      tmp = mindiag;
+      rank--;
+    }
+
+    // save diagonal
+    mat[j*(n+1)] = mju_sqrt(tmp);
+
+    // process off-diagonal entries
+    tmp = 1/mat[j*(n+1)];
+    for (int i=j+1; i<n; i++) {
+      mat[i*n+j] = (mat[i*n+j] - mju_dot(mat+i*n, mat+j*n, j)) * tmp;
+    }
+  }
+
+  return rank;
+}
+
+
+
+// Cholesky solve
+void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n) {
+  // copy if source and destination are different
+  if (res!=vec) {
+    mju_copy(res, vec, n);
+  }
+
+  // forward substitution: solve L*res = vec
+  for (int i=0; i<n; i++) {
+    if (i) {
+      res[i] -= mju_dot(mat+i*n, res, i);
+    }
+
+    // diagonal
+    res[i] /= mat[i*(n+1)];
+  }
+
+  // backward substitution: solve L'*res = res
+  for (int i=n-1; i>=0; i--) {
+    if (i<n-1) {
+      for (int j=i+1; j<n; j++) {
+        res[i] -= mat[j*n+i] * res[j];
+      }
+    }
+    // diagonal
+    res[i] /= mat[i*(n+1)];
+  }
+}
+
+
+
+// Cholesky rank-one update: L*L' +/- x*x'; return rank
+int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus) {
+  int rank = n;
+  mjtNum r, c, cinv, s, Lkk, tmp;
+
+  for (int k=0; k<n; k++) {
+    if (x[k]) {
+      // prepare constants
+      Lkk = mat[k*(n+1)];
+      tmp = Lkk*Lkk + (flg_plus ? x[k]*x[k] : -x[k]*x[k]);
+      if (tmp<mjMINVAL) {
+        tmp = mjMINVAL;
+        rank--;
+      }
+      r = mju_sqrt(tmp);
+      c = r / Lkk;
+      cinv = 1 / c;
+      s = x[k] / Lkk;
+
+      // update diagonal
+      mat[k*(n+1)] = r;
+
+      // update mat
+      if (flg_plus)
+        for (int i=k+1; i<n; i++) {
+          mat[i*n+k] = (mat[i*n+k] + s*x[i])*cinv;
+        } else
+        for (int i=k+1; i<n; i++) {
+          mat[i*n+k] = (mat[i*n+k] - s*x[i])*cinv;
+        }
+
+      // update x
+      for (int i=k+1; i<n; i++) {
+        x[i] = c*x[i] - s*mat[i*n+k];
+      }
+    }
+  }
+
+  return rank;
+}
+
+
+
+//---------------------------- sparse Cholesky -----------------------------------------------------
+
+// sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank'
+//  mat must have uncompressed layout; rownnz is modified to end at diagonal
+int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
+                         int* rownnz, int* rowadr, int* colind,
+                         mjData* d) {
+  int rank = n;
+
+  mjMARKSTACK;
+  int* buf_ind = (int*) mj_stackAlloc(d, n);
+  mjtNum* sparse_buf = mj_stackAlloc(d, n);
+
+  // shrink rows so that rownnz ends at diagonal
+  for (int r=0; r<n; r++) {
+    // shrink
+    while (rownnz[r]>0 && colind[rowadr[r]+rownnz[r]-1]>r) {
+      rownnz[r]--;
+    }
+
+    // check
+    if (rownnz[r]==0 || colind[rowadr[r]+rownnz[r]-1]!=r) {
+      mju_error("Matrix must have non-zero diagonal in mju_cholFactorSparse");
+    }
+  }
+
+  // backpass over rows
+  for (int r=n-1; r>=0; r--) {
+    // get rownnz and rowadr for row r
+    int nnz = rownnz[r], adr = rowadr[r];
+
+    // update row r diagonal
+    mjtNum tmp = mat[adr+nnz-1];
+    if (tmp<mindiag) {
+      tmp = mindiag;
+      rank--;
+    }
+    mat[adr+nnz-1] = mju_sqrt(tmp);
+    tmp = 1/mat[adr+nnz-1];
+
+    // update row r before diagonal
+    for (int i=0; i<nnz-1; i++) {
+      mat[adr+i] *= tmp;
+    }
+
+    // update row c<r where mat(r,c)!=0
+    for (int i=0; i<nnz-1; i++) {
+      // get column index
+      int c = colind[adr+i];
+
+      // mat(c,0:c) = mat(c,0:c) - mat(r,c) * mat(r,0:c)
+      int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], c + 1, 1, -mat[adr+i],
+                                    rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r],
+                                    sparse_buf, buf_ind);
+
+      // assign new nnz to row c
+      rownnz[c] = nnz_c;
+    }
+  }
+
+  mjFREESTACK;
+  return rank;
+}
+
+
+
+// sparse reverse-order Cholesky solve
+void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
+                         const int* rownnz, const int* rowadr, const int* colind) {
+  // copy input into result
+  mju_copy(res, vec, n);
+
+  // vec <- L^-T vec
+  for (int i=n-1; i>=0; i--) {
+    if (res[i]) {
+      // get rowadr[i], rownnz[i]
+      const int adr = rowadr[i], nnz = rownnz[i];
+
+      // x(i) /= L(i,i)
+      res[i] /= mat[adr+nnz-1];
+      mjtNum tmp = res[i];
+
+      // x(j) -= L(i,j)*x(i), j=0:i-1
+      for (int j=0; j<nnz-1; j++) {
+        res[colind[adr+j]] -= mat[adr+j]*tmp;
+      }
+    }
+  }
+
+  // vec <- L^-1 vec
+  for (int i=0; i<n; i++) {
+    // get rowadr[i], rownnz[i]
+    const int adr = rowadr[i], nnz = rownnz[i];
+
+    // x(i) -= sum_j L(i,j)*x(j), j=0:i-1
+    if (nnz>1) {
+      res[i] -= mju_dotSparse(mat+adr, res, nnz-1, colind+adr);
+      // modulo AVX, the above line does
+      // for (int j=0; j<nnz-1; j++)
+      //   res[i] -= mat[adr+j]*res[colind[adr+j]];
+    }
+
+
+    // x(i) /= L(i,i)
+    res[i] /= mat[adr+nnz-1];
+  }
+}
+
+
+
+// sparse reverse-order Cholesky rank-one update: L'*L +/- x*x'; return rank
+//  x is sparse, change in sparsity pattern of mat is not allowed
+int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
+                         int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind,
+                         mjData* d) {
+  mjMARKSTACK;
+  int* buf_ind = (int*) mj_stackAlloc(d, n);
+  mjtNum* sparse_buf = mj_stackAlloc(d, n);
+
+  // backpass over rows corresponding to non-zero x(r)
+  int rank = n, i = x_nnz - 1;
+  while (i>=0) {
+    // get rownnz and rowadr for this row
+    int nnz = rownnz[x_ind[i]], adr = rowadr[x_ind[i]];
+
+    // compute quantities
+    mjtNum tmp = mat[adr+nnz-1]*mat[adr+nnz-1] + (flg_plus ? x[i]*x[i] : -x[i]*x[i]);
+    if (tmp<mjMINVAL) {
+      tmp = mjMINVAL;
+      rank--;
+    }
+    mjtNum r = mju_sqrt(tmp);
+    mjtNum c = r / mat[adr+nnz-1];
+    mjtNum s = x[i] / mat[adr+nnz-1];
+
+    // update diagonal
+    mat[adr+nnz-1] = r;
+
+    // update row:  mat(r,1:r-1) = (mat(r,1:r-1) + s*x(1:r-1)) / c
+    int new_nnz = mju_combineSparse(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
+                                    nnz-1, i, colind + adr, x_ind,
+                                    sparse_buf, buf_ind);
+
+    // check for size change
+    if (new_nnz!=nnz-1) {
+      mju_error("Varying sparsity pattern in mju_cholUpdateSparse");
+    }
+
+    // update x:  x(1:r-1) = c*x(1:r-1) - s*mat(r,1:r-1)
+    int new_x_nnz = mju_combineSparse(x, mat+adr, n, c, -s,
+                                      i, nnz-1, x_ind, colind+adr,
+                                      sparse_buf, buf_ind);
+
+    // update i, correct for changing x
+    i = i - 1 + (new_x_nnz - i);
+  }
+
+  mjFREESTACK;
+  return rank;
+}
+
+
+
+//--------------------------- eigen decomposition --------------------------------------------------
+
+// eigenvalue decomposition of symmetric 3x3 matrix
+static const mjtNum eigEPS = 1E-12;
+int mju_eig3(mjtNum* eigval, mjtNum* eigvec, mjtNum quat[4], const mjtNum mat[9]) {
+  mjtNum D[9], tmp[9];
+  mjtNum tau, t, c;
+  int iter, rk, ck, rotk;
+
+  // initialize with unit quaternion
+  quat[0] = 1;
+  quat[1] = quat[2] = quat[3] = 0;
+
+  // Jacobi iteration
+  for (iter=0; iter<500; iter++) {
+    // make quaternion matrix eigvec, compute D = eigvec'*mat*eigvec
+    mju_quat2Mat(eigvec, quat);
+    mju_mulMatTMat(tmp, eigvec, mat, 3, 3, 3);
+    mju_mulMatMat(D, tmp, eigvec, 3, 3, 3);
+
+    // assign eigenvalues
+    eigval[0] = D[0];
+    eigval[1] = D[4];
+    eigval[2] = D[8];
+
+    // find max off-diagonal element, set indices
+    if (fabs(D[1])>fabs(D[2]) && fabs(D[1])>fabs(D[5])) {
+      rk = 0;     // row
+      ck = 1;     // column
+      rotk = 2;   // rotation axis
+    } else if (fabs(D[2])>fabs(D[5])) {
+      rk = 0;
+      ck = 2;
+      rotk = 1;
+    } else {
+      rk = 1;
+      ck = 2;
+      rotk = 0;
+    }
+
+    // terminate if max off-diagonal element too small
+    if (fabs(D[3*rk+ck])<eigEPS) {
+      break;
+    }
+
+    // 2x2 symmetric Schur decomposition
+    tau = (D[4*ck]-D[4*rk])/(2*D[3*rk+ck]);
+    if (tau>=0) {
+      t = 1.0/(tau + mju_sqrt(1 + tau*tau));
+    } else {
+      t = -1.0/(-tau + mju_sqrt(1 + tau*tau));
+    }
+    c = 1.0/mju_sqrt(1 + t*t);
+
+    // terminate if cosine too close to 1
+    if (c>1.0-eigEPS) {
+      break;
+    }
+
+    // express rotation as quaternion
+    tmp[1] = tmp[2] = tmp[3] = 0;
+    tmp[rotk+1] = (tau>=0 ? -mju_sqrt(0.5-0.5*c) : mju_sqrt(0.5-0.5*c));
+    if (rotk==1) {
+      tmp[rotk+1] = -tmp[rotk+1];
+    }
+    tmp[0] = mju_sqrt(1.0 - tmp[rotk+1]*tmp[rotk+1]);
+    mju_normalize4(tmp);
+
+    // accumulate quaternion rotation
+    mju_mulQuat(tmp+4, quat, tmp);
+    mju_copy4(quat, tmp+4);
+    mju_normalize4(quat);
+  }
+
+  // sort eigenvalues in decreasing order (bubblesort: 0, 1, 0)
+  for (int j=0; j<3; j++) {
+    int j1 = j%2;       // lead index
+
+    if (eigval[j1] < eigval[j1+1]) {
+      // swap eigenvalues
+      t = eigval[j1];
+      eigval[j1] = eigval[j1+1];
+      eigval[j1+1] = t;
+
+      // rotate quaternion
+      tmp[0] = 0.707106781186548;     // mju_cos(pi/4) = mju_sin(pi/4)
+      tmp[1] = tmp[2] = tmp[3] = 0;
+      tmp[(j1+2)%3+1] = tmp[0];
+      mju_mulQuat(tmp+4, quat, tmp);
+      mju_copy4(quat, tmp+4);
+      mju_normalize4(quat);
+    }
+  }
+
+  // recompute eigvec
+  mju_quat2Mat(eigvec, quat);
+
+  return iter;
+}
+
+
+
+//---------------------------------- QCQP ----------------------------------------------------------
+
+// solve QCQP in 2 dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+int mju_QCQP2(mjtNum* res, const mjtNum* Ain, const mjtNum* bin,
+              const mjtNum* d, mjtNum r) {
+  mjtNum A11, A22, A12, b1, b2;
+  mjtNum P11, P22, P12, det, detinv, v1, v2, la, val, deriv;
+
+  // scale A,b so that constraint becomes x'*x <= r*r
+  b1 = bin[0]*d[0];
+  b2 = bin[1]*d[1];
+  A11 = Ain[0]*d[0]*d[0];
+  A22 = Ain[3]*d[1]*d[1];
+  A12 = Ain[1]*d[0]*d[1];
+
+  // Newton iteration
+  la = 0;
+  for (int iter=0; iter<20; iter++) {
+    // det(A+la)
+    det = (A11+la)*(A22+la) - A12*A12;
+
+    // check SPD, with 1e-10 threshold
+    if (det<1e-10) {
+      res[0] = 0;
+      res[1] = 0;
+      return 0;
+    }
+
+    // P = inv(A+la)
+    detinv = 1/det;
+    P11 = (A22+la)*detinv;
+    P22 = (A11+la)*detinv;
+    P12 = -A12*detinv;
+
+    // v = -P*b
+    v1 = -P11*b1 - P12*b2;
+    v2 = -P12*b1 - P22*b2;
+
+    // val = v'*v - r*r
+    val = v1*v1 + v2*v2 - r*r;
+
+    // check for convergence, or initial solution inside constraint set
+    if (val<1e-10) {
+      break;
+    }
+
+    // deriv = -2 * v' * P * v
+    deriv = -2.0*(P11*v1*v1 + 2.0*P12*v1*v2 + P22*v2*v2);
+
+    // compute update, exit if too small
+    mjtNum delta = -val/deriv;
+    if (delta<1e-10) {
+      break;
+    }
+
+    // update
+    la += delta;
+  }
+
+  // undo scaling
+  res[0] = v1*d[0];
+  res[1] = v2*d[1];
+
+  return (la!=0);
+}
+
+
+
+// solve QCQP in 3 dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+int mju_QCQP3(mjtNum* res, const mjtNum* Ain, const mjtNum* bin,
+              const mjtNum* d, mjtNum r) {
+  mjtNum A11, A22, A33, A12, A13, A23, b1, b2, b3;
+  mjtNum P11, P22, P33, P12, P13, P23, det, detinv, v1, v2, v3, la, val, deriv;
+
+  // scale A,b so that constraint becomes x'*x <= r*r
+  b1 = bin[0]*d[0];
+  b2 = bin[1]*d[1];
+  b3 = bin[2]*d[2];
+  A11 = Ain[0]*d[0]*d[0];
+  A22 = Ain[4]*d[1]*d[1];
+  A33 = Ain[8]*d[2]*d[2];
+  A12 = Ain[1]*d[0]*d[1];
+  A13 = Ain[2]*d[0]*d[2];
+  A23 = Ain[5]*d[1]*d[2];
+
+  // Newton iteration
+  la = 0;
+  for (int iter=0; iter<20; iter++) {
+    // unscaled P
+    P11 = (A22+la)*(A33+la) - A23*A23;
+    P22 = (A11+la)*(A33+la) - A13*A13;
+    P33 = (A11+la)*(A22+la) - A12*A12;
+    P12 = A13*A23 - A12*(A33+la);
+    P13 = A12*A23 - A13*(A22+la);
+    P23 = A12*A13 - A23*(A11+la);
+
+    // det(A+la)
+    det = (A11+la)*P11 + A12*P12 + A13*P13;
+
+    // check SPD, with 1e-10 threshold
+    if (det<1e-10) {
+      res[0] = 0;
+      res[1] = 0;
+      res[2] = 0;
+      return 0;
+    }
+
+    // detinv
+    detinv = 1/det;
+
+    // final P
+    P11 *= detinv;
+    P22 *= detinv;
+    P33 *= detinv;
+    P12 *= detinv;
+    P13 *= detinv;
+    P23 *= detinv;
+
+    // v = -P*b
+    v1 = -P11*b1 - P12*b2 - P13*b3;
+    v2 = -P12*b1 - P22*b2 - P23*b3;
+    v3 = -P13*b1 - P23*b2 - P33*b3;
+
+    // val = v'*v - r*r
+    val = v1*v1 + v2*v2 + v3*v3 - r*r;
+
+    // check for convergence, or initial solution inside constraint set
+    if (val<1e-10) {
+      break;
+    }
+
+    // deriv = -2 * v' * P * v
+    deriv = -2.0*(P11*v1*v1 + P22*v2*v2 + P33*v3*v3)
+            -4.0*(P12*v1*v2 + P13*v1*v3 + P23*v2*v3);
+
+    // compute update, exit if too small
+    mjtNum delta = -val/deriv;
+    if (delta<1e-10) {
+      break;
+    }
+
+    // update
+    la += delta;
+  }
+
+  // undo scaling
+  res[0] = v1*d[0];
+  res[1] = v2*d[1];
+  res[2] = v3*d[2];
+
+  return (la!=0);
+}
+
+
+
+// solve QCQP in n dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+int mju_QCQP(mjtNum* res, const mjtNum* Ain, const mjtNum* bin,
+             const mjtNum* d, mjtNum r, int n) {
+  mjtNum A[25], Ala[25], b[5];
+  mjtNum la, val, deriv, tmp[5];
+
+  // check size
+  if (n>5) {
+    mju_error("mju_QCQP supports n up to 5");
+  }
+
+  // scale A,b so that constraint becomes x'*x <= r*r
+  for (int i=0; i<n; i++) {
+    b[i] = bin[i] * d[i];
+
+    for (int j=0; j<n; j++) {
+      A[j+i*n] = Ain[j+i*n] * d[i] * d[j];
+    }
+  }
+
+  // Newton iteration
+  la = 0;
+  for (int iter=0; iter<20; iter++) {
+    // make A+la
+    mju_copy(Ala, A, n*n);
+    for (int i=0; i<n; i++) {
+      Ala[i*(n+1)] += la;
+    }
+
+    // factorize, check rank with 1e-10 threshold
+    if (mju_cholFactor(Ala, n, 1e-10) < n) {
+      mju_zero(res, n);
+      return 0;
+    }
+
+    // set res = -Ala \ b
+    mju_cholSolve(res, Ala, b, n);
+    mju_scl(res, res, -1, n);
+
+    // val = b' * Ala^-2 * b - r*r
+    val = mju_dot(res, res, n) - r*r;
+
+    // check for convergence, or initial solution inside constraint set
+    if (val<1e-10) {
+      break;
+    }
+
+    // deriv = -2 * b' * Ala^-3 * b
+    mju_cholSolve(tmp, Ala, res, n);
+    deriv = -2.0 * mju_dot(res, tmp, n);
+
+    // compute update, exit if too small
+    mjtNum delta = -val/deriv;
+    if (delta<1e-10) {
+      break;
+    }
+
+    // update
+    la += delta;
+  }
+
+  // undo scaling
+  for (int i=0; i<n; i++) {
+    res[i] = res[i] * d[i];
+  }
+
+  return (la!=0);
+}

src/engine/engine_util_solve.h

@@ -0,0 +1,73 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_SOLVE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SOLVE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Cholesky decomposition: mat = L*L'; return rank
+MJAPI int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
+
+// Cholesky solve
+MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);
+
+// Cholesky rank-one update: L*L' +/- x*x'; return rank
+MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
+
+// sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank'
+//  mat must have uncompressed layout; rownnz is modified to end at diagonal
+int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
+                         int* rownnz, int* rowadr, int* colind,
+                         mjData* d);
+
+// sparse reverse-order Cholesky solve
+void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
+                         const int* rownnz, const int* rowadr, const int* colind);
+
+// sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank
+//  x is sparse, change in sparsity pattern of mat is not allowed
+int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
+                         int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind,
+                         mjData* d);
+
+// eigenvalue decomposition of symmetric 3x3 matrix
+MJAPI int mju_eig3(mjtNum* eigval, mjtNum* eigvec, mjtNum quat[4], const mjtNum mat[9]);
+
+// solve QCQP in 2 dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+MJAPI int mju_QCQP2(mjtNum* res, const mjtNum* Ain, const mjtNum* bin, const mjtNum* d, mjtNum r);
+
+// solve QCQP in 3 dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+MJAPI int mju_QCQP3(mjtNum* res, const mjtNum* Ain, const mjtNum* bin, const mjtNum* d, mjtNum r);
+
+// solve QCQP in n<=5 dimensions:
+//  min  0.5*x'*A*x + x'*b  s.t.  sum (xi/di)^2 <= r^2
+// return 0 if unconstrained, 1 if constrained
+int mju_QCQP(mjtNum* res, const mjtNum* Ain, const mjtNum* bin, const mjtNum* d, mjtNum r, int n);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_SOLVE_H_

src/engine/engine_util_sparse.c

@@ -0,0 +1,735 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_util_sparse.h"
+
+#include <string.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjtnum.h>
+#include "engine/engine_io.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+
+#ifdef mjUSEPLATFORMSIMD
+  #if defined(__AVX__) && defined(mjUSEDOUBLE)
+    #define mjUSEAVX
+    #include "immintrin.h"
+  #endif
+#endif
+
+//------------------------------ sparse operations -------------------------------------------------
+
+// dot-product, first vector is sparse
+mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
+                     const int nnz1, const int* ind1) {
+  int i = 0;
+  mjtNum res = 0;
+
+#ifdef mjUSEAVX
+  int nnz1_4 = nnz1 - 4;
+
+  // vector part
+  if (nnz1_4>=0) {
+    __m256d sum, prod, val1, val2;
+    __m128d vlow, vhigh, high64;
+
+    // init
+    val2 = _mm256_set_pd(vec2[ind1[3]],
+                         vec2[ind1[2]],
+                         vec2[ind1[1]],
+                         vec2[ind1[0]]);
+    val1 = _mm256_loadu_pd(vec1);
+    sum = _mm256_mul_pd(val1, val2);
+    i = 4;
+
+    // parallel computation
+    while (i<=nnz1_4) {
+      val1 = _mm256_loadu_pd(vec1+i);
+      val2 = _mm256_set_pd(vec2[ind1[i+3]],
+                           vec2[ind1[i+2]],
+                           vec2[ind1[i+1]],
+                           vec2[ind1[i+0]]);
+      prod = _mm256_mul_pd(val1, val2);
+      sum = _mm256_add_pd(sum, prod);
+      i += 4;
+    }
+
+    // reduce
+    vlow = _mm256_castpd256_pd128(sum);
+    vhigh = _mm256_extractf128_pd(sum, 1);
+    vlow = _mm_add_pd(vlow, vhigh);
+    high64 = _mm_unpackhi_pd(vlow, vlow);
+    res = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
+  }
+#endif
+
+  // scalar part
+  for (; i<nnz1; i++) {
+    res += vec1[i] * vec2[ind1[i]];
+  }
+
+  return res;
+}
+
+
+
+// dot-productX3, first vector is sparse; supernode of size 3
+void mju_dotSparseX3(mjtNum* res0, mjtNum* res1, mjtNum* res2,
+                     const mjtNum* vec10, const mjtNum* vec11, const mjtNum* vec12,
+                     const mjtNum* vec2, const int nnz1, const int* ind1) {
+  int i = 0;
+
+  // clear result
+  mjtNum RES0 = 0;
+  mjtNum RES1 = 0;
+  mjtNum RES2 = 0;
+
+#ifdef mjUSEAVX
+  int nnz1_4 = nnz1 - 4;
+
+  // vector part
+  if (nnz1_4>=0) {
+    __m256d sum0, sum1, sum2, prod, val1, val2;
+    __m128d vlow, vhigh, high64;
+
+    // init
+    val2 = _mm256_set_pd(vec2[ind1[3]],
+                         vec2[ind1[2]],
+                         vec2[ind1[1]],
+                         vec2[ind1[0]]);
+    val1 = _mm256_loadu_pd(vec10);
+    sum0 = _mm256_mul_pd(val1, val2);
+    val1 = _mm256_loadu_pd(vec11);
+    sum1 = _mm256_mul_pd(val1, val2);
+    val1 = _mm256_loadu_pd(vec12);
+    sum2 = _mm256_mul_pd(val1, val2);
+    i = 4;
+
+    // parallel computation
+    while (i<=nnz1_4) {
+      // get val2 only once
+      val2 = _mm256_set_pd(vec2[ind1[i+3]],
+                           vec2[ind1[i+2]],
+                           vec2[ind1[i+1]],
+                           vec2[ind1[i+0]]);
+
+      // process each val1
+      val1 = _mm256_loadu_pd(vec10+i);
+      prod = _mm256_mul_pd(val1, val2);
+      sum0 = _mm256_add_pd(sum0, prod);
+
+      val1 = _mm256_loadu_pd(vec11+i);
+      prod = _mm256_mul_pd(val1, val2);
+      sum1 = _mm256_add_pd(sum1, prod);
+
+      val1 = _mm256_loadu_pd(vec12+i);
+      prod = _mm256_mul_pd(val1, val2);
+      sum2 = _mm256_add_pd(sum2, prod);
+
+      i += 4;
+    }
+
+    // reduce
+    vlow   = _mm256_castpd256_pd128(sum0);
+    vhigh  = _mm256_extractf128_pd(sum0, 1);
+    vlow   = _mm_add_pd(vlow, vhigh);
+    high64 = _mm_unpackhi_pd(vlow, vlow);
+    RES0   = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
+
+    vlow   = _mm256_castpd256_pd128(sum1);
+    vhigh  = _mm256_extractf128_pd(sum1, 1);
+    vlow   = _mm_add_pd(vlow, vhigh);
+    high64 = _mm_unpackhi_pd(vlow, vlow);
+    RES1   = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
+
+    vlow   = _mm256_castpd256_pd128(sum2);
+    vhigh  = _mm256_extractf128_pd(sum2, 1);
+    vlow   = _mm_add_pd(vlow, vhigh);
+    high64 = _mm_unpackhi_pd(vlow, vlow);
+    RES2   = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
+  }
+#endif
+
+  // scalar part
+  for (; i<nnz1; i++) {
+    mjtNum v2 = vec2[ind1[i]];
+
+    RES0 += vec10[i] * v2;
+    RES1 += vec11[i] * v2;
+    RES2 += vec12[i] * v2;
+  }
+
+  // copy result
+  *res0 = RES0;
+  *res1 = RES1;
+  *res2 = RES2;
+}
+
+
+
+// dot-product, both vectors are sparse
+mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
+                      const int nnz1, const int* ind1,
+                      const int nnz2, const int* ind2) {
+  int i1 = 0, i2 = 0;
+  mjtNum res = 0;
+
+  // check for empty array
+  if (!nnz1 || !nnz2) {
+    return 0;
+  }
+
+  while (i1<nnz1 && i2<nnz2) {
+    // get current indices
+    int adr1 = ind1[i1], adr2 = ind2[i2];
+
+    // match: accumulate result, advance both
+    if (adr1==adr2) {
+      res += vec1[i1++] * vec2[i2++];
+    }
+
+    // otherwise advance smaller
+    else if (adr1<adr2) {
+      i1++;
+    } else {
+      i2++;
+    }
+  }
+
+  return res;
+}
+
+
+
+// convert matrix from dense to sparse
+void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                      int* rownnz, int* rowadr, int* colind) {
+  int adr = 0;
+
+  // find non-zeros and construct sparse
+  for (int r=0; r<nr; r++) {
+    // init row
+    rownnz[r] = 0;
+    rowadr[r] = adr;
+
+    // find non-zeros
+    for (int c=0; c<nc; c++) {
+      if (mat[r*nc+c]) {
+        // record index and count
+        colind[adr] = c;
+        rownnz[r]++;
+
+        // copy element
+        res[adr++] = mat[r*nc+c];
+      }
+    }
+  }
+}
+
+
+
+// convert matrix from sparse to dense
+void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                      const int* rownnz, const int* rowadr, const int* colind) {
+  // clear
+  mju_zero(res, nr*nc);
+
+  // copy non-zeros
+  for (int r=0; r<nr; r++) {
+    for (int i=0; i<rownnz[r]; i++) {
+      res[r*nc + colind[rowadr[r]+i]] = mat[rowadr[r]+i];
+    }
+  }
+}
+
+
+
+// multiply sparse matrix and dense vector:  res = mat * vec.
+void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                         int nr, const int* rownnz, const int* rowadr,
+                         const int* colind, const int* rowsuper) {
+  // no supernodes, or no AVX
+#ifdef mjUSEAVX
+  if (!rowsuper)
+#endif
+  {
+    // regular sparse dot-product
+    for (int r=0; r<nr; r++) {
+      res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
+    }
+
+    return;
+  }
+
+  // regular or supernode
+  for (int r=0; r<nr; r++) {
+    if (rowsuper[r]) {
+      int rs = rowsuper[r]+1;
+
+      // handle rows in blocks of 3
+      while (rs>=3) {
+        mju_dotSparseX3(res+r, res+r+1, res+r+2,
+                        mat+rowadr[r], mat+rowadr[r+1], mat+rowadr[r+2],
+                        vec, rownnz[r], colind+rowadr[r]);
+
+        r += 3;
+        rs -= 3;
+      }
+
+      // handle remaining rows
+      while (rs>0) {
+        res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
+
+        r++;
+        rs--;
+      }
+
+      // go back one, because of outer for loop
+      r--;
+    }
+
+    else {
+      res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
+    }
+  }
+}
+
+// res = res*scl1 + vec*scl2
+static void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m256d sclpar1, sclpar2, sum, val1, val2;
+
+    // init
+    sclpar1 = _mm256_set1_pd(scl1);
+    sclpar2 = _mm256_set1_pd(scl2);
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm256_loadu_pd(res+i);
+      val2 = _mm256_loadu_pd(vec+i);
+      val1 = _mm256_mul_pd(val1, sclpar1);
+      val2 = _mm256_mul_pd(val2, sclpar2);
+      sum = _mm256_add_pd(val1, val2);
+      _mm256_storeu_pd(res+i, sum);
+      i += 4;
+    }
+  }
+
+  // process remaining
+  int n_i = n - i;
+  if (n_i==3) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+    res[i+1] = res[i+1]*scl1 + vec[i+1]*scl2;
+    res[i+2] = res[i+2]*scl1 + vec[i+2]*scl2;
+  } else if (n_i==2) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+    res[i+1] = res[i+1]*scl1 + vec[i+1]*scl2;
+  } else if (n_i==1) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+  }
+
+#else
+  for (; i<n; i++) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+  }
+#endif
+}
+
+// return 1 if vec1==vec2, 0 otherwise
+static int mju_compare(const int* vec1, const int* vec2, int n) {
+  int i = 0;
+
+#ifdef mjUSEAVX
+  int n_4 = n - 4;
+
+  // vector part
+  if (n_4>=0) {
+    __m128i val1, val2, cmp;
+
+    // parallel computation
+    while (i<=n_4) {
+      val1 = _mm_loadu_si128((const __m128i*)(vec1+i));
+      val2 = _mm_loadu_si128((const __m128i*)(vec2+i));
+      cmp = _mm_cmpeq_epi32(val1, val2);
+      if (_mm_movemask_epi8(cmp)!= 0xFFFF) {
+        return 0;
+      }
+      i += 4;
+    }
+  }
+#endif
+
+  // scalar part
+  for (; i<n; i++) {
+    if (vec1[i]!=vec2[i]) {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
+// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
+int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+                      int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
+                      mjtNum* buf, int* buf_ind) {
+  // check for identical pattern
+  if (dst_nnz==src_nnz) {
+    if (mju_compare(dst_ind, src_ind, dst_nnz)) {
+      // combine mjtNum data directly
+      mju_addToSclScl(dst, src, a, b, dst_nnz);
+      return dst_nnz;
+    }
+  }
+
+  // copy dst into buf
+  if (dst_nnz) {
+    memcpy(buf, dst, dst_nnz*sizeof(mjtNum));
+    memcpy(buf_ind, dst_ind, dst_nnz*sizeof(int));
+  }
+
+  // prepare to merge buf and scr into dst
+  int bi = 0, si = 0, nnz = 0;
+  int buf_nnz = dst_nnz;
+  int badr = bi<buf_nnz ? buf_ind[bi] : n+1;
+  int sadr = si<src_nnz ? src_ind[si] : n+1;
+
+  // merge vectors
+  while (bi<buf_nnz || si<src_nnz) {
+    // both
+    if (badr==sadr) {
+      dst[nnz] = a*buf[bi++] + b*src[si++];
+      dst_ind[nnz++] = badr;
+
+      badr = bi<buf_nnz ? buf_ind[bi] : n+1;
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+    }
+
+    // dst only
+    else if (badr<sadr) {
+      dst[nnz] = a*buf[bi++];
+      dst_ind[nnz++] = badr;
+
+      badr = bi<buf_nnz ? buf_ind[bi] : n+1;
+    }
+
+    // src only
+    else {
+      dst[nnz] = b*src[si++];
+      dst_ind[nnz++] = sadr;
+
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+    }
+  }
+
+  return nnz;
+}
+
+
+
+// incomplete combine sparse: dst = a*dst + b*src at common indices
+void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+                          int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind) {
+  // check for identical pattern
+  if (dst_nnz==src_nnz) {
+    if (mju_compare(dst_ind, src_ind, dst_nnz)) {
+      // combine mjtNum data directly
+      mju_addToSclScl(dst, src, a, b, dst_nnz);
+      return;
+    }
+  }
+
+  // scale dst by a
+  if (a!=1) {
+    mju_scl(dst, dst, a, dst_nnz);
+  }
+
+  // prepare to merge
+  int di = 0, si = 0;
+  int dadr = di<dst_nnz ? dst_ind[di] : n+1;
+  int sadr = si<src_nnz ? src_ind[si] : n+1;
+
+  // add src*b at common indices
+  while (di<dst_nnz) {
+    // both
+    if (dadr==sadr) {
+      dst[di++] += b*src[si++];
+
+      dadr = di<dst_nnz ? dst_ind[di] : n+1;
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+    }
+
+    // dst only
+    else if (dadr<sadr) {
+      di++;
+      dadr = di<dst_nnz ? dst_ind[di] : n+1;
+    }
+
+    // src only
+    else {
+      si++;
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+    }
+  }
+}
+
+
+
+// compress layout of sparse matrix
+void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind) {
+  rowadr[0] = 0;
+  int adr = rownnz[0];
+  for (int r=1; r<nr; r++) {
+    // save old rowadr, record new
+    int rowadr1 = rowadr[r];
+    rowadr[r] = adr;
+
+    // shift mat and mat_colind
+    for (int adr1=rowadr1; adr1<rowadr1+rownnz[r]; adr1++) {
+      mat[adr] = mat[adr1];
+      colind[adr] = colind[adr1];
+      adr++;
+    }
+  }
+}
+
+// transpose sparse matrix
+void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                         int* res_rownnz, int* res_rowadr, int* res_colind,
+                         const int* rownnz, const int* rowadr, const int* colind) {
+  // clear counters for transposed
+  memset(res_rownnz, 0, nc*sizeof(int));
+
+  // set uncompressed layout
+  for (int rt=0; rt<nc; rt++) {
+    res_rowadr[rt] = rt*nr;
+  }
+
+  // scan original, compute uncompressed
+  for (int r=0; r<nr; r++) {
+    for (int ci=0; ci<rownnz[r]; ci++) {
+      // get rt=c
+      int rt = colind[rowadr[r]+ci];
+
+      // record index ct=r, assuming uncompressed res_rowadr[rt]=rt*nr
+      res_colind[rt*nr + res_rownnz[rt]] = r;
+
+      // copy data
+      res[rt*nr + res_rownnz[rt]] = mat[rowadr[r]+ci];
+
+      // increase counter for rt
+      res_rownnz[rt]++;
+    }
+  }
+
+  // compress
+  mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind);
+}
+
+
+
+// construct row supernodes
+void mju_superSparse(int nr, int* rowsuper,
+                     const int* rownnz, const int* rowadr, const int* colind) {
+  // no rows: nothing to do
+  if (!nr) {
+    return;
+  }
+
+  // find match to child
+  for (int r=0; r<nr-1; r++) {
+    // different number of nonzeros: cannot be a match
+    if (rownnz[r]!=rownnz[r+1]) {
+      rowsuper[r] = 0;
+    }
+
+    // same number of nonzeros: compare colind vectors
+    else {
+      rowsuper[r] = mju_compare(colind+rowadr[r], colind+rowadr[r+1], rownnz[r]);
+    }
+  }
+
+  // clear last (by definition)
+  rowsuper[nr-1] = 0;
+
+  // accumulate in reverse
+  for (int r=nr-2; r>=0; r--) {
+    if (rowsuper[r]) {
+      rowsuper[r] += rowsuper[r+1];
+    }
+  }
+}
+
+// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
+void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
+                        const mjtNum* diag, int nr, int nc,
+                        int* res_rownnz, int* res_rowadr, int* res_colind,
+                        const int* rownnz, const int* rowadr,
+                        const int* colind, const int* rowsuper,
+                        const int* rownnzT, const int* rowadrT,
+                        const int* colindT, const int* rowsuperT,
+                        mjData* d) {
+  // allocate space for accumulation buffer and matT
+  mjMARKSTACK;
+  int* chain = (int*) mj_stackAlloc(d, 2*nc);
+  mjtNum* buffer = mj_stackAlloc(d, nc);
+
+  // set uncompressed layout
+  for (int r=0; r<nc; r++) {
+    res_rowadr[r] = r*nc;
+  }
+
+  // compute lower-triangular uncompressed layout (nc per row)
+  for (int r=0; r<nc; r++) {
+    // copy chain from parent
+    if (rowsuperT && r>0 && rowsuperT[r-1]>0) {
+      // copy parent chain
+      res_rownnz[r] = res_rownnz[r-1];
+      memcpy(res_colind+res_rowadr[r], res_colind+res_rowadr[r-1],
+             res_rownnz[r]*sizeof(int));
+
+      // add diagonal if rowT is not empty
+      if (rownnzT[r]) {
+        res_colind[res_rowadr[r]+res_rownnz[r]] = r;
+        res_rownnz[r]++;
+      }
+    }
+
+    // construct chain
+    else {
+      // clear chain accumulation buffers
+      int nchain = 0;
+      int inew = 0, iold = nc;
+      int lastadded = -1;
+
+      // for each nonzero c in matT_row(r), add nonzeros of mat_row(c) to chain(r)
+      for (int i=0; i<rownnzT[r]; i++) {
+        // save c
+        int c = colindT[rowadrT[r]+i];
+
+        // skip if a chain from same supernode was already added
+        if (rowsuper && lastadded>=0 && (c-lastadded)<=rowsuper[lastadded]) {
+          continue;
+        } else {
+          lastadded = c;
+        }
+
+        // swap chains
+        int adr = inew;
+        inew = iold;
+        iold = adr;
+
+        // merge chains
+        int nnewchain = 0;
+        adr = 0;
+        int end = rowadr[c]+rownnz[c];
+        for (int adr1=rowadr[c]; adr1<end; adr1++) {
+          // save column index from mat
+          int col_mat = colind[adr1];
+
+          // skip column indices in chain smaller than col_mat
+          while (adr<nchain && chain[iold + adr]<col_mat && chain[iold + adr]<=r) {
+            chain[inew + nnewchain++] = chain[iold + adr++];
+          }
+
+          // only lower-triangular
+          if (col_mat>r) {
+            break;
+          }
+
+          // existing element: advance chain
+          if (adr<nchain && chain[iold + adr]==col_mat) {
+            adr++;
+          }
+
+          // add column index from matT
+          chain[inew + nnewchain++] = col_mat;
+        }
+
+        // append the rest of the master chain
+        while (adr<nchain && chain[iold + adr]<=r) {
+          chain[inew + nnewchain++] = chain[iold + adr++];
+        }
+
+        // assign newchain
+        nchain = nnewchain;
+      }
+
+      // copy chain
+      res_rownnz[r] = nchain;
+      if (nchain) {
+        memcpy(res_colind+res_rowadr[r], chain+inew, nchain*sizeof(int));
+      }
+    }
+  }
+
+  // compute matrix data given uncompressed layout
+  for (int r=0; r<nc; r++) {
+    // clear buffer[colind] for this chain
+    int adr = res_rowadr[r];
+    for (int i=0; i<res_rownnz[r]; i++) {
+      buffer[res_colind[adr+i]] = 0;
+    }
+
+    // res_row(r) = sum_c ( matT(r,c) * diag(c) * mat_row(c) )
+    for (int i=0; i<rownnzT[r]; i++) {
+      // save c and matT(r,c)*diag(c)
+      int c = colindT[rowadrT[r]+i];
+      mjtNum matTrc = matT[rowadrT[r]+i];
+      if (diag) {
+        matTrc *= diag[c];
+      }
+
+      // process row
+      int end = rowadr[c]+rownnz[c];
+      for (int adr=rowadr[c]; adr<end; adr++) {
+        // get column index from mat, only lower-triangular
+        int adr1;
+        if ((adr1=colind[adr])>r) {
+          break;
+        }
+
+        // add to buffer
+        buffer[adr1] += matTrc*mat[adr];
+      }
+    }
+
+    // copy buffer
+    adr = res_rowadr[r];
+    for (int i=0; i<res_rownnz[r]; i++) {
+      res[adr+i] = buffer[res_colind[adr+i]];
+    }
+  }
+
+  // make symmetric; uncompressed layout
+  for (int r=1; r<nc; r++) {
+    int end = nc*r+res_rownnz[r]-1;
+    for (int adr=nc*r; adr<end; adr++) {
+      // add to row given by column index
+      int adr1 = nc*res_colind[adr] + res_rownnz[res_colind[adr]]++;
+      res[adr1] = res[adr];
+      res_colind[adr1] = r;
+    }
+  }
+
+  mjFREESTACK;
+}

src/engine/engine_util_sparse.h

@@ -0,0 +1,86 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjtnum.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//------------------------------ sparse operations -------------------------------------------------
+
+// dot-product, first vector is sparse
+mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
+                     const int nnz1, const int* ind1);
+
+// dot-product, both vectors are sparse
+mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
+                      const int nnz1, const int* ind1,
+                      const int nnz2, const int* ind2);
+
+// convert matrix from dense to sparse
+void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                      int* rownnz, int* rowadr, int* colind);
+
+// convert matrix from sparse to dense
+void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                      const int* rownnz, const int* rowadr, const int* colind);
+
+// multiply sparse matrix and dense vector:  res = mat * vec
+void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                         int nr, const int* rownnz, const int* rowadr,
+                         const int* colind, const int* rowsuper);
+
+// compress layout of sparse matrix
+void mju_compressSparse(mjtNum* mat, int nr, int nc,
+                        int* rownnz, int* rowadr, int* colind);
+
+// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
+int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+                      int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
+                      mjtNum* buf, int* buf_ind);
+
+// incomplete combine sparse: dst = a*dst + b*src at common indices
+void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+                          int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind);
+
+// transpose sparse matrix
+void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                         int* res_rownnz, int* res_rowadr, int* res_colind,
+                         const int* rownnz, const int* rowadr, const int* colind);
+
+// construct row supernodes
+void mju_superSparse(int nr, int* rowsuper,
+                     const int* rownnz, const int* rowadr, const int* colind);
+
+// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
+void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
+                        const mjtNum* diag, int nr, int nc,
+                        int* res_rownnz, int* res_rowadr, int* res_colind,
+                        const int* rownnz, const int* rowadr,
+                        const int* colind, const int* rowsuper,
+                        const int* rownnzT, const int* rowadrT,
+                        const int* colindT, const int* rowsuperT,
+                        mjData* d);
+
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_

src/engine/engine_util_spatial.c

@@ -0,0 +1,485 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_util_spatial.h"
+
+#include <math.h>
+
+#include <mujoco/mjmodel.h>
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+
+
+//------------------------------ quaternion operations ---------------------------------------------
+
+// rotate vector by quaternion
+void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]) {
+  // null quat: copy vec
+  if (quat[0]==1 && quat[1]==0 && quat[2]==0 && quat[3]==0) {
+    mju_copy3(res, vec);
+  }
+
+  // regular processing
+  else {
+    mjtNum mat[9];
+    mju_quat2Mat(mat, quat);
+    mju_rotVecMat(res, vec, mat);
+  }
+}
+
+
+
+// negate quaternion
+void mju_negQuat(mjtNum res[4], const mjtNum quat[4]) {
+  res[0] = quat[0];
+  res[1] = -quat[1];
+  res[2] = -quat[2];
+  res[3] = -quat[3];
+}
+
+
+
+// multiply quaternions
+void mju_mulQuat(mjtNum res[4], const mjtNum qa[4], const mjtNum qb[4]) {
+  res[0] = qa[0]*qb[0] - qa[1]*qb[1] - qa[2]*qb[2] - qa[3]*qb[3];
+  res[1] = qa[0]*qb[1] + qa[1]*qb[0] + qa[2]*qb[3] - qa[3]*qb[2];
+  res[2] = qa[0]*qb[2] - qa[1]*qb[3] + qa[2]*qb[0] + qa[3]*qb[1];
+  res[3] = qa[0]*qb[3] + qa[1]*qb[2] - qa[2]*qb[1] + qa[3]*qb[0];
+}
+
+
+
+// multiply quaternion and axis
+void mju_mulQuatAxis(mjtNum res[4], const mjtNum quat[4], const mjtNum axis[3]) {
+  res[0] = - (quat[1]*axis[0] + quat[2]*axis[1] + quat[3]*axis[2]);
+  res[1] = quat[0]*axis[0] + quat[2]*axis[2] - quat[3]*axis[1];
+  res[2] = quat[0]*axis[1] + quat[3]*axis[0] - quat[1]*axis[2];
+  res[3] = quat[0]*axis[2] + quat[1]*axis[1] - quat[2]*axis[0];
+}
+
+
+
+// convert axisAngle to quaternion
+void mju_axisAngle2Quat(mjtNum res[4], const mjtNum axis[3], mjtNum angle) {
+  // zero angle: null quat
+  if (angle==0) {
+    res[0] = 1;
+    res[1] = 0;
+    res[2] = 0;
+    res[3] = 0;
+  }
+
+  // regular processing
+  else {
+    mjtNum s = mju_sin(angle*0.5);
+    res[0] = mju_cos(angle*0.5);
+    res[1] = axis[0]*s;
+    res[2] = axis[1]*s;
+    res[3] = axis[2]*s;
+  }
+}
+
+
+
+// convert quaternion (corresponding to orientation difference) to 3D velocity
+void mju_quat2Vel(mjtNum res[3], const mjtNum quat[4], mjtNum dt) {
+  mjtNum axis[3] = {quat[1], quat[2], quat[3]};
+  mjtNum sin_a_2 = mju_normalize3(axis);
+  mjtNum speed = 2 * mju_atan2(sin_a_2, quat[0]);
+
+  // when axis-angle is larger than pi, rotation is in the opposite direction
+  if (speed>mjPI) {
+    speed -= 2*mjPI;
+  }
+  speed /= dt;
+
+  mju_scl3(res, axis, speed);
+}
+
+
+
+// Subtract quaternions, express as 3D velocity: qb*quat(res) = qa.
+void mju_subQuat(mjtNum res[3], const mjtNum qa[4], const mjtNum qb[4]) {
+  // qdif = neg(qb)*qa
+  mjtNum qneg[4], qdif[4];
+  mju_negQuat(qneg, qb);
+  mju_mulQuat(qdif, qneg, qa);
+
+  // convert to 3D velocity
+  mju_quat2Vel(res, qdif, 1);
+}
+
+
+
+// convert quaternion to 3D rotation matrix
+void mju_quat2Mat(mjtNum res[9], const mjtNum quat[4]) {
+  // null quat: identity
+  if (quat[0]==1 && quat[1]==0 && quat[2]==0 && quat[3]==0) {
+    res[0] = 1;
+    res[1] = 0;
+    res[2] = 0;
+    res[3] = 0;
+    res[4] = 1;
+    res[5] = 0;
+    res[6] = 0;
+    res[7] = 0;
+    res[8] = 1;
+  }
+
+  // regular processing
+  else {
+    const mjtNum q00 = quat[0]*quat[0];
+    const mjtNum q01 = quat[0]*quat[1];
+    const mjtNum q02 = quat[0]*quat[2];
+    const mjtNum q03 = quat[0]*quat[3];
+    const mjtNum q11 = quat[1]*quat[1];
+    const mjtNum q12 = quat[1]*quat[2];
+    const mjtNum q13 = quat[1]*quat[3];
+    const mjtNum q22 = quat[2]*quat[2];
+    const mjtNum q23 = quat[2]*quat[3];
+    const mjtNum q33 = quat[3]*quat[3];
+
+    res[0] = q00 + q11 - q22 - q33;
+    res[4] = q00 - q11 + q22 - q33;
+    res[8] = q00 - q11 - q22 + q33;
+
+    res[1] = 2*(q12 - q03);
+    res[2] = 2*(q13 + q02);
+    res[3] = 2*(q12 + q03);
+    res[5] = 2*(q23 - q01);
+    res[6] = 2*(q13 - q02);
+    res[7] = 2*(q23 + q01);
+  }
+}
+
+
+
+// convert 3D rotation matrix to quaterion
+void mju_mat2Quat(mjtNum quat[4], const mjtNum mat[9]) {
+  // q0 largest
+  if (mat[0]+mat[4]+mat[8]>0) {
+    quat[0] = 0.5 * mju_sqrt(1 + mat[0] + mat[4] + mat[8]);
+    quat[1] = 0.25 * (mat[7] - mat[5]) / quat[0];
+    quat[2] = 0.25 * (mat[2] - mat[6]) / quat[0];
+    quat[3] = 0.25 * (mat[3] - mat[1]) / quat[0];
+  }
+
+  // q1 largest
+  else if (mat[0]>mat[4] && mat[0]>mat[8]) {
+    quat[1] = 0.5 * mju_sqrt(1 + mat[0] - mat[4] - mat[8]);
+    quat[0] = 0.25 * (mat[7] - mat[5]) / quat[1];
+    quat[2] = 0.25 * (mat[1] + mat[3]) / quat[1];
+    quat[3] = 0.25 * (mat[2] + mat[6]) / quat[1];
+  }
+
+  // q2 largest
+  else if (mat[4]>mat[8]) {
+    quat[2] = 0.5 * mju_sqrt(1 - mat[0] + mat[4] - mat[8]);
+    quat[0] = 0.25 * (mat[2] - mat[6]) / quat[2];
+    quat[1] = 0.25 * (mat[1] + mat[3]) / quat[2];
+    quat[3] = 0.25 * (mat[5] + mat[7]) / quat[2];
+  }
+
+  // q3 largest
+  else {
+    quat[3] = 0.5 * mju_sqrt(1 - mat[0] - mat[4] + mat[8]);
+    quat[0] = 0.25 * (mat[3] - mat[1]) / quat[3];
+    quat[1] = 0.25 * (mat[2] + mat[6]) / quat[3];
+    quat[2] = 0.25 * (mat[5] + mat[7]) / quat[3];
+  }
+
+  mju_normalize4(quat);
+}
+
+
+
+// time-derivative of quaternion, given 3D rotational velocity
+void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]) {
+  res[0] = 0.5*(-vel[0]*quat[1] - vel[1]*quat[2] - vel[2]*quat[3]);
+  res[1] = 0.5*( vel[0]*quat[0] + vel[1]*quat[3] - vel[2]*quat[2]);
+  res[2] = 0.5*(-vel[0]*quat[3] + vel[1]*quat[0] + vel[2]*quat[1]);
+  res[3] = 0.5*( vel[0]*quat[2] - vel[1]*quat[1] + vel[2]*quat[0]);
+}
+
+
+
+// integrate quaterion given 3D angular velocity
+void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale) {
+  mjtNum angle, tmp[4], qrot[4];
+
+  // form local rotation quaternion, apply
+  mju_copy3(tmp, vel);
+  angle = scale * mju_normalize3(tmp);
+  mju_axisAngle2Quat(qrot, tmp, angle);
+  mju_mulQuat(tmp, quat, qrot);
+  mju_normalize4(tmp);
+  mju_copy4(quat, tmp);
+}
+
+
+
+// compute quaternion performing rotation from z-axis to given vector
+void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]) {
+  mjtNum axis[3], a, vn[3] = {vec[0], vec[1], vec[2]}, z[3] = {0, 0, 1};
+
+  // set default result to no-rotation quaternion
+  quat[0] = 1;
+  mju_zero3(quat+1);
+
+  // normalize vector; if too small, no rotation
+  if (mju_normalize3(vn)<mjMINVAL) {
+    return;
+  }
+
+  // compute angle and axis
+  mju_cross(axis, z, vn);
+  a = mju_normalize3(axis);
+
+  // almost parallel
+  if (fabs(a)<mjMINVAL) {
+    // opposite: 180 deg rotation around x axis
+    if (mju_dot3(vn, z) < 0) {
+      quat[0] = 0;
+      quat[1] = 1;
+    }
+
+    return;
+  }
+
+  // make quaterion from angle and axis
+  a = mju_atan2(a, mju_dot3(vn, z));
+  mju_axisAngle2Quat(quat, axis, a);
+}
+
+
+
+//------------------------------ pose operations (quat, pos) ---------------------------------------
+
+// multiply two poses
+void mju_mulPose(mjtNum posres[3], mjtNum quatres[4],
+                 const mjtNum pos1[3], const mjtNum quat1[4],
+                 const mjtNum pos2[3], const mjtNum quat2[4]) {
+  // quatres = quat1*quat2
+  mju_mulQuat(quatres, quat1, quat2);
+  mju_normalize4(quatres);
+
+  // posres = quat1*pos2 + pos1
+  mju_rotVecQuat(posres, pos2, quat1);
+  mju_addTo3(posres, pos1);
+}
+
+
+
+// negate pose
+void mju_negPose(mjtNum posres[3], mjtNum quatres[4], const mjtNum pos[3], const mjtNum quat[4]) {
+  // qres = neg(quat)
+  mju_negQuat(quatres, quat);
+
+  // pres = -neg(quat)*pos
+  mju_rotVecQuat(posres, pos, quatres);
+  mju_scl3(posres, posres, -1);
+}
+
+
+
+// transform vector by pose
+void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[4], const mjtNum vec[3]) {
+  // res = quat*vec + pos
+  mju_rotVecQuat(res, vec, quat);
+  mju_addTo3(res, pos);
+}
+
+
+
+//------------------------------ spatial algebra ---------------------------------------------------
+
+// vector cross-product, 3D
+void mju_cross(mjtNum res[3], const mjtNum a[3], const mjtNum b[3]) {
+  res[0] = a[1]*b[2] - a[2]*b[1];
+  res[1] = a[2]*b[0] - a[0]*b[2];
+  res[2] = a[0]*b[1] - a[1]*b[0];
+}
+
+
+
+// cross-product for motion vector
+void mju_crossMotion(mjtNum res[6], const mjtNum vel[6], const mjtNum v[6]) {
+  res[0] = -vel[2]*v[1] + vel[1]*v[2];
+  res[1] =  vel[2]*v[0] - vel[0]*v[2];
+  res[2] = -vel[1]*v[0] + vel[0]*v[1];
+  res[3] = -vel[2]*v[4] + vel[1]*v[5];
+  res[4] =  vel[2]*v[3] - vel[0]*v[5];
+  res[5] = -vel[1]*v[3] + vel[0]*v[4];
+
+  res[3] += -vel[5]*v[1] + vel[4]*v[2];
+  res[4] +=  vel[5]*v[0] - vel[3]*v[2];
+  res[5] += -vel[4]*v[0] + vel[3]*v[1];
+}
+
+
+
+// cross-product for force vectors
+void mju_crossForce(mjtNum res[6], const mjtNum vel[6], const mjtNum f[6]) {
+  res[0] = -vel[2]*f[1] + vel[1]*f[2];
+  res[1] =  vel[2]*f[0] - vel[0]*f[2];
+  res[2] = -vel[1]*f[0] + vel[0]*f[1];
+  res[3] = -vel[2]*f[4] + vel[1]*f[5];
+  res[4] =  vel[2]*f[3] - vel[0]*f[5];
+  res[5] = -vel[1]*f[3] + vel[0]*f[4];
+
+  res[0] += -vel[5]*f[4] + vel[4]*f[5];
+  res[1] +=  vel[5]*f[3] - vel[3]*f[5];
+  res[2] += -vel[4]*f[3] + vel[3]*f[4];
+}
+
+
+
+// express inertia in com-based frame
+void mju_inertCom(mjtNum res[10], const mjtNum inert[3], const mjtNum mat[9],
+                  const mjtNum dif[3], mjtNum mass) {
+  // tmp = diag(inert) * mat'  (mat is local-to-global rotation)
+  mjtNum tmp[9] = {mat[0]*inert[0], mat[3]*inert[0], mat[6]*inert[0],
+                   mat[1]*inert[1], mat[4]*inert[1], mat[7]*inert[1],
+                   mat[2]*inert[2], mat[5]*inert[2], mat[8]*inert[2]
+                  };
+
+  // res_rot = mat * diag(inert) * mat'
+  res[0] = mat[0]*tmp[0] + mat[1]*tmp[3] + mat[2]*tmp[6];
+  res[1] = mat[3]*tmp[1] + mat[4]*tmp[4] + mat[5]*tmp[7];
+  res[2] = mat[6]*tmp[2] + mat[7]*tmp[5] + mat[8]*tmp[8];
+  res[3] = mat[0]*tmp[1] + mat[1]*tmp[4] + mat[2]*tmp[7];
+  res[4] = mat[0]*tmp[2] + mat[1]*tmp[5] + mat[2]*tmp[8];
+  res[5] = mat[3]*tmp[2] + mat[4]*tmp[5] + mat[5]*tmp[8];
+
+  // res_rot -= mass * dif_cross * dif_cross
+  res[0] += mass*(dif[1]*dif[1] + dif[2]*dif[2]);
+  res[1] += mass*(dif[0]*dif[0] + dif[2]*dif[2]);
+  res[2] += mass*(dif[0]*dif[0] + dif[1]*dif[1]);
+  res[3] -= mass*dif[0]*dif[1];
+  res[4] -= mass*dif[0]*dif[2];
+  res[5] -= mass*dif[1]*dif[2];
+
+  // res_tran = mass * dif
+  res[6] = mass*dif[0];
+  res[7] = mass*dif[1];
+  res[8] = mass*dif[2];
+
+  // res_mass = mass
+  res[9] = mass;
+}
+
+
+
+// multiply 6D vector (rotation, translation) by 6D inertia matrix
+void mju_mulInertVec(mjtNum res[6], const mjtNum i[10], const mjtNum v[6]) {
+  res[0] = i[0]*v[0] + i[3]*v[1] + i[4]*v[2] - i[8]*v[4] + i[7]*v[5];
+  res[1] = i[3]*v[0] + i[1]*v[1] + i[5]*v[2] + i[8]*v[3] - i[6]*v[5];
+  res[2] = i[4]*v[0] + i[5]*v[1] + i[2]*v[2] - i[7]*v[3] + i[6]*v[4];
+  res[3] = i[8]*v[1] - i[7]*v[2] + i[9]*v[3];
+  res[4] = i[6]*v[2] - i[8]*v[0] + i[9]*v[4];
+  res[5] = i[7]*v[0] - i[6]*v[1] + i[9]*v[5];
+}
+
+
+
+// express motion axis in com-based frame
+void mju_dofCom(mjtNum res[6], const mjtNum axis[3], const mjtNum offset[3]) {
+  // hinge
+  if (offset) {
+    mju_copy3(res, axis);
+    mju_cross(res+3, axis, offset);
+  }
+
+  // slide
+  else {
+    mju_zero3(res);
+    mju_copy3(res+3, axis);
+  }
+}
+
+
+
+// multiply dof matrix (6-by-n, transposed) by vector (n-by-1)
+void mju_mulDofVec(mjtNum* res, const mjtNum* dof, const mjtNum* vec, int n) {
+  if (n==1) {
+    mju_scl(res, dof, vec[0], 6);
+  } else if (n<=0) {
+    mju_zero(res, 6);
+  } else {
+    mju_mulMatTVec(res, dof, vec, n, 6);
+  }
+}
+
+
+
+// transform 6D motion or force vector between frames
+//  rot is 3-by-3 matrix; flg_force determines vector type (motion or force)
+void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_force,
+                          const mjtNum newpos[3], const mjtNum oldpos[3],
+                          const mjtNum rotnew2old[9]) {
+  mjtNum cros[3], dif[3], tran[6];
+
+  // apply translation
+  mju_copy(tran, vec, 6);
+  mju_sub3(dif, newpos, oldpos);
+  if (flg_force) {
+    mju_cross(cros, dif, vec+3);
+    mju_sub3(tran, vec, cros);
+  } else {
+    mju_cross(cros, dif, vec);
+    mju_sub3(tran+3, vec+3, cros);
+  }
+
+  // apply rotation if provided
+  if (rotnew2old) {
+    mju_rotVecMatT(res, tran, rotnew2old);
+    mju_rotVecMatT(res+3, tran+3, rotnew2old);
+  }
+
+  // otherwise copy
+  else {
+    mju_copy(res, tran, 6);
+  }
+}
+
+
+
+// make 3D frame given X axis (normal) and possibly Y axis (tangent 1)
+void mju_makeFrame(mjtNum frame[9]) {
+  mjtNum tmp[3];
+
+  // normalize xaxis
+  if (mju_normalize3(frame) < 0.5) {
+    mju_error("xaxis of contact frame undefined");
+  }
+
+  // if yaxis undefined, set yaxis to (0,1,0) if possible, otherwize (0,0,1)
+  if (mju_norm3(frame+3) < 0.5) {
+    mju_zero3(frame+3);
+
+    if (frame[1]<0.5 && frame[1]>-0.5) {
+      frame[4] = 1;
+    } else {
+      frame[5] = 1;
+    }
+  }
+
+  // make yaxis orthogonal to xaxis
+  mju_scl3(tmp, frame, mju_dot3(frame, frame+3));
+  mju_subFrom3(frame+3, tmp);
+  mju_normalize3(frame+3);
+
+  // zaxis = cross(xaxis, yaxis)
+  mju_cross(frame+6, frame, frame+3);
+}

src/engine/engine_util_spatial.h

@@ -0,0 +1,115 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPATIAL_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPATIAL_H_
+
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+//------------------------------ quaternion operations ---------------------------------------------
+
+// rotate vector by quaternion
+MJAPI void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]);
+
+// compute conjugate quaternion, corresponding to opposite rotation
+MJAPI void mju_negQuat(mjtNum res[4], const mjtNum quat[4]);
+
+// multiply quaternions
+MJAPI void mju_mulQuat(mjtNum res[4], const mjtNum quat1[4], const mjtNum quat2[4]);
+
+// multiply quaternion and axis
+MJAPI void mju_mulQuatAxis(mjtNum res[4], const mjtNum quat[4], const mjtNum axis[3]);
+
+// convert axisAngle to quaternion
+MJAPI void mju_axisAngle2Quat(mjtNum res[4], const mjtNum axis[3], mjtNum angle);
+
+// convert quaternion (corresponding to orientation difference) to 3D velocity
+MJAPI void mju_quat2Vel(mjtNum res[3], const mjtNum quat[4], mjtNum dt);
+
+// subtract quaternions, convert to 3D velocity: qb*quat(res) = qa
+MJAPI void mju_subQuat(mjtNum res[3], const mjtNum qa[4], const mjtNum qb[4]);
+
+// convert quaternion to 3D rotation matrix
+MJAPI void mju_quat2Mat(mjtNum res[9], const mjtNum quat[4]);
+
+// convert 3D rotation matrix to quaterion
+MJAPI void mju_mat2Quat(mjtNum quat[4], const mjtNum mat[9]);
+
+// time-derivative of quaternion, given 3D rotational velocity
+MJAPI void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]);
+
+// integrate quaterion given 3D angular velocity
+MJAPI void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale);
+
+// compute quaternion performing rotation from z-axis to given vector
+MJAPI void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]);
+
+
+//------------------------------ pose operations (pos, quat) ---------------------------------------
+
+// multiply two poses
+MJAPI void mju_mulPose(mjtNum posres[3], mjtNum quatres[4],
+                       const mjtNum pos1[3], const mjtNum quat1[4],
+                       const mjtNum pos2[3], const mjtNum quat2[4]);
+
+// compute conjugate pose, corresponding to the opposite spatial transformation
+MJAPI void mju_negPose(mjtNum posres[3], mjtNum quatres[4],
+                       const mjtNum pos[3], const mjtNum quat[4]);
+
+// transform vector by pose
+MJAPI void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[4],
+                          const mjtNum vec[3]);
+
+
+//------------------------------ spatial algebra ---------------------------------------------------
+
+// vector cross-product, 3D
+MJAPI void mju_cross(mjtNum res[3], const mjtNum a[3], const mjtNum b[3]);
+
+// cross-product for motion vector
+void mju_crossMotion(mjtNum res[6], const mjtNum vel[6], const mjtNum v[6]);
+
+// cross-product for force vectors
+void mju_crossForce(mjtNum res[6], const mjtNum vel[6], const mjtNum f[6]);
+
+// express inertia in com-based frame
+void mju_inertCom(mjtNum res[10], const mjtNum inert[3], const mjtNum mat[9],
+                  const mjtNum dif[3], mjtNum mass);
+
+// express motion axis in com-based frame
+void mju_dofCom(mjtNum res[6], const mjtNum axis[3], const mjtNum offset[3]);
+
+// multiply 6D vector (rotation, translation) by 6D inertia matrix
+void mju_mulInertVec(mjtNum res[6], const mjtNum inert[10], const mjtNum vec[6]);
+
+// multiply dof matrix by vector
+void mju_mulDofVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);
+
+// coordinate transform of 6D motion or force vector in rotation:translation format
+//  rotnew2old is 3-by-3, NULL means no rotation; flg_force specifies force or motion type
+MJAPI void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_force,
+                                const mjtNum newpos[3], const mjtNum oldpos[3],
+                                const mjtNum rotnew2old[9]);
+
+// make 3D frame given X axis (and possibly Y axis)
+void mju_makeFrame(mjtNum frame[9]);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPATIAL_H_

src/engine/engine_vfs.c

@@ -0,0 +1,211 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_vfs.h"
+
+#include <string.h>
+#include <stdlib.h>
+
+#include "engine/engine_array_safety.h"
+#include "engine/engine_file.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+
+// strip path prefix from filename
+static void vfs_strippath(char* newname, const char* oldname) {
+  int i, sz = strlen(oldname);
+
+  // find last delimiter
+  for (i=sz-1; i>=0; i--) {
+    if (oldname[i]=='\\' || oldname[i]=='/') {
+      break;
+    }
+  }
+
+  // check resulting length
+  if (sz-(i+1)>=mjMAXVFSNAME) {
+    mju_error("Filename too long in VFS");
+  }
+  if (sz-(i+1)<=0) {
+    mju_error("Empty filename in VFS");
+  }
+
+  // copy
+  mju_strncpy(newname, oldname+i+1, mjMAXVFSNAME);
+
+  // make lowercase
+  for (i=strlen(newname)-1; i>=0; i--) {
+    if (newname[i]>='A' && newname[i]<='Z') {
+      newname[i] = (char)(((int)newname[i]) +'a' - 'A');
+    }
+  }
+}
+
+
+
+// initialize to empty (no deallocation)
+void mj_defaultVFS(mjVFS* vfs) {
+  memset(vfs, 0, sizeof(mjVFS));
+}
+
+
+
+// add file to VFS, return 0: success, 1: full, 2: repeated name, -1: failed to load
+int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename) {
+  // check vfs size
+  if (vfs->nfile>=mjMAXVFS-1) {
+    return 1;
+  }
+
+  // make full name
+  char fullname[1000];
+  if (directory) {
+    mjSTRNCPY(fullname, directory);
+    mjSTRNCAT(fullname, filename);
+  } else {
+    mjSTRNCPY(fullname, filename);
+  }
+
+  // strip path
+  char newname[mjMAXVFSNAME];
+  vfs_strippath(newname, filename);
+
+  // check for repeated name
+  for (int i=0; i<vfs->nfile; i++) {
+    if (strncmp(newname, vfs->filename[i], mjMAXVFSNAME)==0) {
+      return 2;
+    }
+  }
+
+  // assign name
+  mjSTRNCPY(vfs->filename[vfs->nfile], newname);
+
+  // allocate and read
+  int filesize = 0;
+  vfs->filedata[vfs->nfile] = mju_fileToMemory(filename, &filesize);
+  if (!vfs->filedata[vfs->nfile]) {
+    return -1;
+  }
+  // assign size and count
+  vfs->filesize[vfs->nfile] = filesize;
+  vfs->nfile++;
+
+  return 0;
+}
+
+
+
+// make empty file in VFS, return 0: success, 1: full, 2: repeated name
+int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize) {
+  // check vfs size
+  if (vfs->nfile>=mjMAXVFS-1) {
+    return 1;
+  }
+
+  // check filesize
+  if (filesize<=0) {
+    mju_error("mj_makeEmptyFileVFS expects positive filesize");
+  }
+
+  // strip path
+  char newname[mjMAXVFSNAME];
+  vfs_strippath(newname, filename);
+
+  // check for repeated name
+  for (int i=0; i<vfs->nfile; i++) {
+    if (strncmp(newname, vfs->filename[i], mjMAXVFSNAME)==0) {
+      return 2;
+    }
+  }
+
+  // assign name
+  mjSTRNCPY(vfs->filename[vfs->nfile], newname);
+
+  // allocate and clear
+  vfs->filedata[vfs->nfile] = mju_malloc(filesize);
+  if (!vfs->filedata[vfs->nfile]) {
+    mju_error("mj_makeEmptyFileVFS: could not allocate memory");
+  }
+  memset(vfs->filedata[vfs->nfile], 0, filesize);
+
+  // assign size and count
+  vfs->filesize[vfs->nfile] = filesize;
+  vfs->nfile++;
+
+  return 0;
+}
+
+
+
+// return file index in VFS, or -1 if not found in VFS
+int mj_findFileVFS(const mjVFS* vfs, const char* filename) {
+  // strip path
+  char newname[mjMAXVFSNAME];
+  vfs_strippath(newname, filename);
+
+  // find specific file
+  for (int i=0; i<vfs->nfile; i++) {
+    if (strncmp(newname, vfs->filename[i], mjMAXVFSNAME)==0) {
+      return i;
+    }
+  }
+
+  return -1;
+}
+
+
+
+// delete file from VFS, return 0: success, -1: not found in VFS
+int mj_deleteFileVFS(mjVFS* vfs, const char* filename) {
+  // strip path
+  char newname[mjMAXVFSNAME];
+  vfs_strippath(newname, filename);
+
+  // find specified file
+  for (int i=0; i<vfs->nfile; i++) {
+    if (strncmp(newname, vfs->filename[i], mjMAXVFSNAME)==0) {
+      // free buffer
+      mju_free(vfs->filedata[i]);
+
+      // scroll remaining files forward
+      while (i<vfs->nfile-1) {
+        mjSTRNCPY(vfs->filename[i], vfs->filename[i+1]);
+        vfs->filesize[i] = vfs->filesize[i+1];
+        vfs->filedata[i] = vfs->filedata[i+1];
+      }
+
+      // set last to 0, for style
+      vfs->filename[vfs->nfile-1][0] = 0;
+      vfs->filesize[vfs->nfile-1] = 0;
+      vfs->filedata[vfs->nfile-1] = NULL;
+
+      // decrease counter
+      vfs->nfile--;
+      return 0;
+    }
+  }
+
+  return -1;
+}
+
+
+
+// delete all files from VFS
+void mj_deleteVFS(mjVFS* vfs) {
+  for (int i=0; i<vfs->nfile; i++) {
+    mju_free(vfs->filedata[i]);
+  }
+
+  memset(vfs, 0, sizeof(mjVFS));
+}

src/engine/engine_vfs.h

@@ -0,0 +1,47 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_VFS_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_VFS_H_
+
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// initialize to empty (no deallocation)
+MJAPI void mj_defaultVFS(mjVFS* vfs);
+
+// add file to VFS, return 0: success, 1: full, 2: repeated name, -1: not found on disk
+MJAPI int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
+
+// make empty file in VFS, return 0: success, 1: full, 2: repeated name
+MJAPI int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
+
+// return file index in VFS, or -1 if not found in VFS
+MJAPI int mj_findFileVFS(const mjVFS* vfs, const char* filename);
+
+// delete file from VFS, return 0: success, -1: not found in VFS
+MJAPI int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
+
+// delete all files from VFS
+MJAPI void mj_deleteVFS(mjVFS* vfs);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_VFS_H_

src/engine/engine_vis_init.c

@@ -0,0 +1,349 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_vis_init.h"
+
+#include <math.h>
+#include <string.h>
+
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+#include "engine/engine_array_safety.h"
+#include "engine/engine_macro.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+
+#ifdef _MSC_VER
+#pragma warning (disable: 4305) // tell MSVC to not complain that float x = 0.1 should be 0.1f
+#endif
+
+
+//--------------------------------- Strings --------------------------------------------------------
+
+// label names
+const char* mjLABELSTRING[mjNLABEL] = {
+  "None",
+  "Body",
+  "Joint",
+  "Geom",
+  "Site",
+  "Camera",
+  "Light",
+  "Tendon",
+  "Actuator",
+  "Constraint",
+  "Skin",
+  "Selection",
+  "SelPoint",
+  "ContactForce"
+};
+
+
+// frame names
+const char* mjFRAMESTRING[mjNFRAME] = {
+  "None",
+  "Body",
+  "Geom",
+  "Site",
+  "Camera",
+  "Light",
+  "Contact",
+  "World"
+};
+
+
+// visual opptions: {name, initial value, shortcut}
+const char* mjVISSTRING[mjNVISFLAG][3] = {
+  {"Convex &Hull",    "0", "H"},
+  {"Te&xture",        "1", "X"},
+  {"&Joint",          "0", "J"},
+  {"Act&uator",       "0", "U"},
+  {"Camera",          "0", "Q"},
+  {"Light",           "0", "Z"},
+  {"Tendon",          "1", "V"},
+  {"Range Finder",    "1", "Y"},
+  {"Co&nstraint",     "0", "N"},
+  {"&Inertia",        "0", "I"},
+  {"Scale Inertia",   "0", "'"},
+  {"Pertur&b Force",  "0", "B"},
+  {"Perturb &Object", "1", "O"},
+  {"&Contact Point",  "0", "C"},
+  {"Contact &Force",  "0", "F"},
+  {"Contact S&plit",  "0", "P"},
+  {"&Transparent",    "0", "T"},
+  {"&Auto Connect",   "0", "A"},
+  {"Center of &Mass", "0", "M"},
+  {"S&elect Point",   "0", "E"},
+  {"Static Bo&dy",    "1", "D"},
+  {"Skin",            "1", ";"}
+};
+
+
+// render options: {name, initial value, shortcut}
+const char* mjRNDSTRING[mjNRNDFLAG][3] = {
+  {"Shadow",      "1", "S"},
+  {"Wireframe",   "0", "W"},
+  {"Reflection",  "1", "R"},
+  {"Additive",    "0", "L"},
+  {"Skybox",      "1", "K"},
+  {"Fog",         "0", "G"},
+  {"Haze",        "1", "/"},
+  {"Segment",     "0", ","},
+  {"Id Color",    "0", "."}
+};
+
+
+
+//--------------------------------- Implementation -------------------------------------------------
+
+
+// allocate and init abstract scene
+void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom) {
+  // free previous
+  mjv_freeScene(scn);
+
+  // allocate geom buffers
+  if (maxgeom>0) {
+    // allocate
+    scn->maxgeom = maxgeom;
+    scn->geoms = (mjvGeom*) mju_malloc(maxgeom*sizeof(mjvGeom));
+    scn->geomorder = (int*) mju_malloc(maxgeom*sizeof(int));
+
+    // check allocation
+    if (!scn->geoms || !scn->geomorder) {
+      mju_error("Could not allocate geom buffers");
+    }
+  }
+
+  // set default OpenGL options
+  for (int i=0; i<mjNRNDFLAG; i++) {
+    scn->flags[i] = (mjRNDSTRING[i][1][0]=='1');
+  }
+
+  // set default model transformation
+  scn->scale = 1;
+  scn->rotate[0] = 1;
+
+  // set number of skins
+  if (m) {
+    scn->nskin = m->nskin;
+  } else {
+    scn->nskin = 0;
+  }
+
+  // allocate skin data
+  if (scn->nskin) {
+    // compute number of vertices in all skins
+    int nskin = m->nskin;
+    int totvert = 0;
+    for (int i=0; i<nskin; i++) {
+      totvert += m->skin_vertnum[i];
+    }
+
+    // allocate
+    scn->skinfacenum = (int*) mju_malloc(nskin*sizeof(int));
+    scn->skinvertadr = (int*) mju_malloc(nskin*sizeof(int));
+    scn->skinvertnum = (int*) mju_malloc(nskin*sizeof(int));
+    scn->skinvert    = (float*) mju_malloc(3*totvert*sizeof(float));
+    scn->skinnormal  = (float*) mju_malloc(3*totvert*sizeof(float));
+
+    // check allocation
+    if (!scn->skinfacenum ||
+        !scn->skinvertadr ||
+        !scn->skinvertnum ||
+        !scn->skinvert    ||
+        !scn->skinnormal) {
+      mju_error("Could not allocate skin buffers");
+    }
+
+    // copy constant data
+    for (int i=0; i<nskin; i++) {
+      scn->skinfacenum[i] = m->skin_facenum[i];
+      scn->skinvertadr[i] = m->skin_vertadr[i];
+      scn->skinvertnum[i] = m->skin_vertnum[i];
+    }
+  }
+
+  // mjvGeom, mjvLight, mjvGLCamera objects are invalid
+}
+
+
+
+// free abstract scene
+void mjv_freeScene(mjvScene* scn) {
+  // free buffers allocated by mjv_makeScene
+  mju_free(scn->geoms);
+  mju_free(scn->geomorder);
+  mju_free(scn->skinfacenum);
+  mju_free(scn->skinvertadr);
+  mju_free(scn->skinvertnum);
+  mju_free(scn->skinvert);
+  mju_free(scn->skinnormal);
+
+  // clear data structure
+  mjv_defaultScene(scn);
+}
+
+
+
+// set default scene
+void mjv_defaultScene(mjvScene* scn) {
+  memset(scn, 0, sizeof(mjvScene));
+}
+
+
+
+// set default visualization options
+void mjv_defaultOption(mjvOption* vopt) {
+  vopt->label = mjLABEL_NONE;
+  vopt->frame = mjFRAME_NONE;
+
+  for (int i=0; i<mjNGROUP; i++) {
+    int state = (i<3 ? 1 : 0);
+    vopt->geomgroup[i] = state;
+    vopt->sitegroup[i] = state;
+    vopt->jointgroup[i] = state;
+    vopt->tendongroup[i] = state;
+    vopt->actuatorgroup[i] = state;
+  }
+
+  for (int i=0; i<mjNVISFLAG; i++) {
+    vopt->flags[i] = (mjVISSTRING[i][1][0]=='1');
+  }
+}
+
+
+
+// set default camera
+void mjv_defaultCamera(mjvCamera* cam) {
+  memset(cam, 0, sizeof(mjvCamera));
+
+  cam->type        = mjCAMERA_FREE;
+  cam->fixedcamid  = -1;
+  cam->trackbodyid = -1;
+  cam->distance    = 2;
+  cam->azimuth     = 90;
+  cam->elevation   = -45;
+}
+
+
+
+// set default perturbation
+void mjv_defaultPerturb(mjvPerturb* pert) {
+  memset(pert, 0, sizeof(mjvPerturb));
+
+  pert->skinselect = -1;
+  pert->refquat[0] = 1;
+  pert->scale = 1;
+}
+
+
+
+// predefined line colors
+static const float _linergb[8][3] = {
+  {1.0, 0.3, 0.3},
+  {0.1, 1.0, 0.1},
+  {0.3, 0.3, 1.0},
+  {0.1, 1.0, 1.0},
+  {1.0, 0.2, 1.0},
+  {1.0, 1.0, 0.1},
+  {1.0, 0.6, 0.2},
+  {0.6, 0.7, 0.4}
+};
+
+
+
+// set default figure
+void mjv_defaultFigure(mjvFigure* fig) {
+  // set everything to zero
+  memset(fig, 0, sizeof(mjvFigure));
+
+  // disable highlight
+  fig->highlightid = -1;
+
+  // set enable flags
+  fig->flg_legend = 1;
+  fig->flg_ticklabel[0] = 1;
+  fig->flg_ticklabel[1] = 1;
+  fig->flg_extend = 1;
+
+  // set style
+  fig->linewidth = 3;
+  fig->gridwidth = 1;
+  fig->gridsize[0] = 2;
+  fig->gridsize[1] = 2;
+  fig->gridrgb[0] = 0.4f;
+  fig->gridrgb[1] = 0.4f;
+  fig->gridrgb[2] = 0.4f;
+  fig->figurergba[3] = 1;
+  fig->panergba[3] = 1;
+  fig->legendrgba[3] = 0.3f;
+  fig->textrgb[0] = 1;
+  fig->textrgb[1] = 1;
+  fig->textrgb[2] = 1;
+  fig->range[0][0] = 0;
+  fig->range[0][1] = 1;
+  fig->range[1][0] = 0;
+  fig->range[1][1] = 1;
+  mjSTRNCPY(fig->xformat, "%.0f");
+  mjSTRNCPY(fig->yformat, "%.2g");
+  mjSTRNCPY(fig->minwidth, "XXX");
+
+  // set line colors
+  for (int n=0; n<mjMAXLINE; n++) {
+    // predefined colors
+    if (n<8) {
+      fig->linergb[n][0] = _linergb[n][0];
+      fig->linergb[n][1] = _linergb[n][1];
+      fig->linergb[n][2] = _linergb[n][2];
+    }
+
+    // automatically generated colors: Halton sequence
+    else {
+      fig->linergb[n][0] = 0.1f + 0.8f*mju_Halton(n, 2);
+      fig->linergb[n][1] = 0.1f + 0.8f*mju_Halton(n, 3);
+      fig->linergb[n][2] = 0.1f + 0.8f*mju_Halton(n, 5);
+    }
+  }
+}
+
+
+
+// compute rbound for mjvGeom
+float mjv_rbound(const mjvGeom* geom) {
+  // model geom: return
+  if (geom->objtype==mjOBJ_GEOM) {
+    return geom->modelrbound;
+  }
+
+  // compute rbound according to type
+  const float* s = geom->size;
+  switch (geom->type) {
+  case mjGEOM_SPHERE:
+    return s[0];
+
+  case mjGEOM_CAPSULE:
+    return (s[0]+s[2]);
+
+  case mjGEOM_CYLINDER:
+    return sqrtf(s[0]*s[0] + s[2]*s[2]);
+
+  case mjGEOM_BOX:
+    return sqrtf(s[0]*s[0] + s[1]*s[1] + s[2]*s[2]);
+    break;
+
+  default:  // not accurate for arrows, but they are not transparent
+    return mjMAX(s[0], mjMAX(s[1], s[2]));
+  }
+}

src/engine/engine_vis_init.h

@@ -0,0 +1,62 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_VIS_INIT_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_VIS_INIT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// strings
+MJAPI extern const char* mjLABELSTRING[mjNLABEL];
+MJAPI extern const char* mjFRAMESTRING[mjNFRAME];
+MJAPI extern const char* mjVISSTRING[mjNVISFLAG][3];
+MJAPI extern const char* mjRNDSTRING[mjNRNDFLAG][3];
+
+
+// set default scene
+MJAPI void mjv_defaultScene(mjvScene* scn);
+
+// allocate and init abstract scene
+MJAPI void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom);
+
+// free abstract scene
+MJAPI void mjv_freeScene(mjvScene* scn);
+
+// set default visualization options
+MJAPI void mjv_defaultOption(mjvOption* vopt);
+
+// set default camera
+MJAPI void mjv_defaultCamera(mjvCamera* cam);
+
+// set default perturbation
+MJAPI void mjv_defaultPerturb(mjvPerturb* pert);
+
+// set default figure
+MJAPI void mjv_defaultFigure(mjvFigure* fig);
+
+// compute rbound for mjvGeom
+float mjv_rbound(const mjvGeom* geom);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_VIS_INIT_H_

src/engine/engine_vis_interact.c

@@ -0,0 +1,767 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_vis_interact.h"
+
+#include <math.h>
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+#include "engine/engine_ray.h"
+#include "engine/engine_support.h"
+#include "engine/engine_util_blas.h"
+#include "engine/engine_util_errmem.h"
+#include "engine/engine_util_misc.h"
+#include "engine/engine_util_spatial.h"
+
+// transform pose from room to model space
+void mjv_room2model(mjtNum* modelpos, mjtNum* modelquat, const mjtNum* roompos,
+                    const mjtNum* roomquat, const mjvScene* scn) {
+  mjtNum translate[3], rotate[4], invpos[3], invquat[4];
+
+  // check scale
+  if (scn->scale<mjMINVAL) {
+    mju_error("mjvScene scale too small in mjv_room2model");
+  }
+
+  // enabled: transform
+  if (scn->enabletransform) {
+    // convert translate, rotate to mjtNum
+    mju_f2n(translate, scn->translate, 3);
+    mju_f2n(rotate, scn->rotate, 4);
+
+    // invert model pose (without scale)
+    mju_negPose(invpos, invquat, translate, rotate);
+
+    // map from room to model space
+    mju_mulPose(modelpos, modelquat, invpos, invquat, roompos, roomquat);
+
+    // divide position by scale
+    mju_scl3(modelpos, modelpos, 1.0/scn->scale);
+  }
+
+  // disabled: copy
+  else {
+    mju_copy3(modelpos, roompos);
+    mju_copy4(modelquat, roomquat);
+  }
+}
+
+
+
+// transform pose from model to room space
+void mjv_model2room(mjtNum* roompos, mjtNum* roomquat, const mjtNum* modelpos,
+                    const mjtNum* modelquat, const mjvScene* scn) {
+  mjtNum translate[3], rotate[4];
+
+  // check scale
+  if (scn->scale<mjMINVAL) {
+    mju_error("mjvScene scale too small in mjv_model2room");
+  }
+
+  // enabled: transform
+  if (scn->enabletransform) {
+    // convert translate, rotate to mjtNum
+    mju_f2n(translate, scn->translate, 3);
+    mju_f2n(rotate, scn->rotate, 4);
+
+    // map from model to room space
+    mju_mulPose(roompos, roomquat, translate, rotate, modelpos, modelquat);
+
+    // scale position
+    mju_scl3(roompos, roompos, scn->scale);
+  }
+
+  // disabled: copy
+  else {
+    mju_copy3(roompos, modelpos);
+    mju_copy4(roomquat, modelquat);
+  }
+}
+
+
+
+// get camera info in model space: average left and right OpenGL cameras
+void mjv_cameraInModel(mjtNum* headpos, mjtNum* forward, mjtNum* up, const mjvScene* scn) {
+  mjtNum pos[3], fwd[3], u[3], quat[4];
+  mjtNum modelpos[3], modelquat[4], modelmat[9];
+
+  // check znear
+  if (scn->camera[0].frustum_near<mjMINVAL || scn->camera[1].frustum_near<mjMINVAL) {
+    mju_error("mjvScene frustum_near too small in mjv_cameraInModel");
+  }
+
+  // clear results
+  if (headpos) {
+    mju_zero3(headpos);
+  }
+  if (forward) {
+    mju_zero3(forward);
+  }
+  if (up) {
+    mju_zero3(up);
+  }
+
+  // average over cameras
+  for (int n=0; n<2; n++) {
+    // convert pos, fwd, u
+    mju_f2n(pos, scn->camera[n].pos, 3);
+    mju_f2n(fwd, scn->camera[n].forward, 3);
+    mju_f2n(u, scn->camera[n].up, 3);
+
+    // normalize just in case
+    mju_normalize3(fwd);
+    mju_normalize3(u);
+
+    // make orientation matrix: x = left, y = up, z = forward
+    mjtNum left[3];
+    mju_cross(left, u, fwd);
+    mju_normalize3(left);
+    mjtNum mat[9] = {
+      left[0], u[0], fwd[0],
+      left[1], u[1], fwd[1],
+      left[2], u[2], fwd[2]
+    };
+    mju_mat2Quat(quat, mat);
+
+    // convert to model space, make orientation matrix
+    mjv_room2model(modelpos, modelquat, pos, quat, scn);
+    mju_quat2Mat(modelmat, modelquat);
+
+    // finalize results
+    if (headpos) {
+      mju_addToScl3(headpos, modelpos, 0.5);
+    }
+    if (forward) {
+      forward[0] += 0.5*modelmat[2];
+      forward[1] += 0.5*modelmat[5];
+      forward[2] += 0.5*modelmat[8];
+    }
+    if (up) {
+      up[0] += 0.5*modelmat[1];
+      up[1] += 0.5*modelmat[4];
+      up[2] += 0.5*modelmat[7];
+    }
+  }
+
+  // normalize forward and up
+  if (forward) {
+    mju_normalize3(forward);
+  }
+  if (up) {
+    mju_normalize3(up);
+  }
+}
+
+
+
+// get camera info in room space: average left and right OpenGL cameras
+void mjv_cameraInRoom(mjtNum* headpos, mjtNum* forward, mjtNum* up, const mjvScene* scn) {
+  mjtNum pos[3], fwd[3], u[3];
+
+  // check znear
+  if (scn->camera[0].frustum_near<mjMINVAL || scn->camera[1].frustum_near<mjMINVAL) {
+    mju_error("mjvScene frustum_near too small in mjv_cameraInRoom");
+  }
+
+  // clear results
+  if (headpos) {
+    mju_zero3(headpos);
+  }
+  if (forward) {
+    mju_zero3(forward);
+  }
+  if (up) {
+    mju_zero3(up);
+  }
+
+  // average over cameras
+  for (int n=0; n<2; n++) {
+    // convert pos, fwd, u
+    mju_f2n(pos, scn->camera[n].pos, 3);
+    mju_f2n(fwd, scn->camera[n].forward, 3);
+    mju_f2n(u, scn->camera[n].up, 3);
+
+    // finalize results
+    if (headpos) {
+      mju_addToScl3(headpos, pos, 0.5);
+    }
+    if (forward) {
+      mju_addToScl3(forward, fwd, 0.5);
+    }
+    if (up) {
+      mju_addToScl3(up, u, 0.5);
+    }
+  }
+
+  // normalize
+  if (forward) {
+    mju_normalize3(forward);
+  }
+  if (up) {
+    mju_normalize3(up);
+  }
+}
+
+
+
+// get frustum height at unit distance from camera; average left and right OpenGL cameras
+mjtNum mjv_frustumHeight(const mjvScene* scn) {
+  mjtNum height;
+
+  // check znear
+  if (scn->camera[0].frustum_near<mjMINVAL || scn->camera[1].frustum_near<mjMINVAL) {
+    mju_error("mjvScene frustum_near too small in mjv_frustumHeight");
+  }
+
+  // add normalized height for left and right cameras
+  height = (scn->camera[0].frustum_top-scn->camera[0].frustum_bottom)/scn->camera[0].frustum_near +
+           (scn->camera[1].frustum_top-scn->camera[1].frustum_bottom)/scn->camera[1].frustum_near;
+
+  // average
+  return 0.5*height;
+}
+
+
+
+// rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_y)
+MJAPI void mjv_alignToCamera(mjtNum* res, const mjtNum* vec, const mjtNum* forward) {
+  mjtNum xaxis[2], yaxis[2];
+
+  // fotward-aligned y-axis
+  mju_copy(yaxis, forward, 2);
+  mju_normalize(yaxis, 2);
+
+  // corresponding x-axis
+  xaxis[0] = yaxis[1];
+  xaxis[1] = -yaxis[0];
+
+  // apply horizontal rotation
+  res[0] = vec[0]*xaxis[0] + vec[1]*yaxis[0];
+  res[1] = vec[0]*xaxis[1] + vec[1]*yaxis[1];
+  res[2] = vec[2];
+}
+
+
+
+// convert 2D mouse motion to z-aligned 3D world coordinates
+static void convert2D(mjtNum* res, int action, mjtNum dx, mjtNum dy, const mjtNum* forward) {
+  mjtNum vec[3];
+
+  // construct 3D vector
+  switch (action) {
+  case mjMOUSE_ROTATE_V:
+    vec[0] = dy;
+    vec[1] = 0;
+    vec[2] = dx;
+    break;
+
+  case mjMOUSE_ROTATE_H:
+    vec[0] = dy;
+    vec[1] = dx;
+    vec[2] = 0;
+    break;
+
+  case mjMOUSE_MOVE_V:
+    vec[0] = dx;
+    vec[1] = 0;
+    vec[2] = -dy;
+    break;
+
+  case mjMOUSE_MOVE_H:
+    vec[0] = dx;
+    vec[1] = -dy;
+    vec[2] = 0;
+    break;
+
+  case mjMOUSE_ZOOM:
+    break;
+
+  default:
+    mju_error_i("Unexpected mouse action %d in convert2D", action);
+  }
+
+  // call 3D converter
+  mjv_alignToCamera(res, vec, forward);
+}
+
+
+
+// move camera with mouse; action is mjtMouse
+void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                    const mjvScene* scn, mjvCamera* cam) {
+  mjtNum headpos[3], forward[3];
+  mjtNum vec[3], dif[3], scl;
+
+  // fixed camera: nothing to do
+  if (cam->type==mjCAMERA_FIXED) {
+    return;
+  }
+
+  // process action
+  switch (action) {
+  case mjMOUSE_ROTATE_V:
+  case mjMOUSE_ROTATE_H:
+    cam->azimuth -= reldx * 180.0;
+    cam->elevation -= reldy * 180.0;
+    break;
+
+  case mjMOUSE_MOVE_V:
+  case mjMOUSE_MOVE_H:
+    // do not move lookat point of tracking camera
+    if (cam->type==mjCAMERA_TRACKING) {
+      return;
+    }
+
+    // get camera info and align
+    mjv_cameraInModel(headpos, forward, NULL, scn);
+    convert2D(vec, action, reldx, reldy, forward);
+
+    // compute scaling: rendered lookat displacement = mouse displacement
+    mju_sub3(dif, cam->lookat, headpos);
+    scl = mjv_frustumHeight(scn) * mju_dot3(dif, forward);
+
+    // move lookat point in opposite direction
+    mju_addToScl3(cam->lookat, vec, -scl);
+    break;
+
+  case mjMOUSE_ZOOM:
+    cam->distance -= mju_log(1 + cam->distance/m->stat.extent/3) * reldy * 9 * m->stat.extent;
+    break;
+
+  default:
+    mju_error_i("Unexpected action %d in mjv_moveCamera", action);
+  }
+
+  // clamp camera parameters
+  if (cam->azimuth > 180) {
+    cam->azimuth -= 360;
+  }
+  if (cam->azimuth < -180) {
+    cam->azimuth += 360;
+  }
+  if (cam->elevation > 89) {
+    cam->elevation = 89;
+  }
+  if (cam->elevation < -89) {
+    cam->elevation = -89;
+  }
+  if (cam->distance < 0.01*m->stat.extent) {
+    cam->distance = 0.01*m->stat.extent;
+  }
+  if (cam->distance > 100*m->stat.extent) {
+    cam->distance = 100*m->stat.extent;
+  }
+}
+
+
+
+// move perturb object with mouse; action is mjtMouse
+void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
+                     mjtNum reldy, const mjvScene* scn, mjvPerturb* pert) {
+  int sel = pert->select;
+  mjtNum forward[3], vec[3], dif[3], scl, q1[4], q2[4], xiquat[4];
+
+  // get camera info and align
+  mjv_cameraInModel(NULL, forward, NULL, scn);
+  convert2D(vec, action, reldx, reldy, forward);
+
+  // process action
+  switch (action) {
+  case mjMOUSE_MOVE_V:
+  case mjMOUSE_MOVE_H:
+    mju_addToScl3(pert->refpos, vec, pert->scale);
+    break;
+
+  case mjMOUSE_ROTATE_V:
+  case mjMOUSE_ROTATE_H:
+    // normalize vector, get length
+    scl = mju_normalize3(vec);
+
+    // make quaternion and apply
+    mju_axisAngle2Quat(q1, vec, scl*mjPI*2);
+    mju_mulQuat(q2, q1, pert->refquat);
+    mju_copy4(pert->refquat, q2);
+    mju_normalize4(pert->refquat);
+
+    // compute xiquat
+    mju_mulQuat(xiquat, d->xquat+4*sel, m->body_iquat+4*sel);
+
+    // limit rotation relative to selected body
+    if (sel>0 && sel<m->nbody) {
+      // q2 = neg(selbody) * refquat
+      mju_negQuat(q1, xiquat);
+      mju_mulQuat(q2, q1, pert->refquat);
+
+      // convert q2 to axis-angle
+      mju_quat2Vel(dif, q2, 1);
+      scl = mju_normalize3(dif);
+
+      // check limit: +/- 90 deg allowed
+      if (scl<-mjPI*0.5 || scl>mjPI*0.5) {
+        // clamp angle
+        scl = mju_max(-mjPI*0.5, mju_min(mjPI*0.5, scl));
+
+        // reconstruct q2
+        mju_axisAngle2Quat(q2, dif, scl);
+
+        // set refquat = selbody * q2_new
+        mju_mulQuat(pert->refquat, xiquat, q2);
+      }
+    }
+    break;
+
+  case mjMOUSE_ZOOM:
+    break;
+
+  default:
+    mju_error_i("Unexpected mouse action %d in mjv_movePerturb", action);
+  }
+}
+
+
+
+// move model with mouse; action is mjtMouse
+void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                   const mjtNum roomup[3], mjvScene* scn) {
+  mjtNum roomforward[3], roomright[3], camforward[3];
+  mjtNum vec[3], scl, quat[4], rotate[4], result[4];
+
+  // transformation disabled: nothing to do
+  if (!scn->enabletransform) {
+    return;
+  }
+
+  // get camera forward in room space
+  mjv_cameraInRoom(NULL, camforward, NULL, scn);
+
+  // make orthogonal to roomright
+  mju_addScl3(roomforward, camforward, roomup, -mju_dot3(camforward, roomup));
+  mju_normalize3(roomforward);
+
+  // compute roomright
+  mju_cross(roomright, roomforward, roomup);
+  mju_normalize3(roomright);
+
+  // process action
+  switch (action) {
+  case mjMOUSE_ROTATE_V:
+  case mjMOUSE_ROTATE_H:
+    // construct rotation vector
+    for (int i=0; i<3; i++) {
+      if (action==mjMOUSE_ROTATE_V) {
+        vec[i] = roomup[i]*reldx + roomright[i]*reldy;
+      } else {
+        vec[i] = roomforward[i]*reldx + roomright[i]*reldy;
+      }
+    }
+
+    // make quaternion from angle-axis
+    scl = mju_normalize3(vec);
+    mju_axisAngle2Quat(quat, vec, scl*mjPI);
+
+    // get current model rotation
+    mju_f2n(rotate, scn->rotate, 4);
+
+    // compose rotation, normalize and and set
+    mju_mulQuat(result, quat, rotate);
+    mju_normalize4(result);
+    mju_n2f(scn->rotate, result, 4);
+    break;
+
+  case mjMOUSE_MOVE_V:
+    for (int i=0; i<3; i++) {
+      scn->translate[i] += (float)(roomright[i]*reldx - roomup[i]*reldy) * m->stat.extent;
+    }
+    break;
+
+  case mjMOUSE_MOVE_H:
+    for (int i=0; i<3; i++) {
+      scn->translate[i] += (float)(roomright[i]*reldx - roomforward[i]*reldy) * m->stat.extent;
+    }
+    break;
+
+  case mjMOUSE_ZOOM:
+    scn->scale += (float)(mju_log(1 + scn->scale/3) * reldy * 3);
+    if (scn->scale<0.01f) {
+      scn->scale = 0.01f;
+    } else if (scn->scale>100.0f) {
+      scn->scale = 100.0f;
+    }
+    break;
+
+  default:
+    mju_error_i("Unexpected action %d in mjv_moveModel", action);
+  }
+}
+
+
+
+// copy perturb pos,quat from selected body; set scale for perturbation
+void mjv_initPerturb(const mjModel* m, const mjData* d, const mjvScene* scn, mjvPerturb* pert) {
+  int sel = pert->select;
+  mjtNum headpos[3], forward[3], dif[3];
+
+  // invalid selected body: return
+  if (sel<=0 || sel>=m->nbody) {
+    return;
+  }
+
+  // copy
+  mju_copy3(pert->refpos, d->xipos + 3*sel);
+  mju_mulQuat(pert->refquat, d->xquat + 4*sel, m->body_iquat + 4*sel);
+
+  // get camera info
+  mjv_cameraInModel(headpos, forward, NULL, scn);
+
+  // compute scaling: rendered pert->refpos displacement = mouse displacement
+  mju_sub3(dif, pert->refpos, headpos);
+  pert->scale = mjv_frustumHeight(scn) * mju_dot3(dif, forward);
+}
+
+
+
+// set perturb pos,quat in d->mocap when selected body is mocap, and in d->qpos otherwise
+//  d->qpos written only if flg_paused and subtree root for selected body has free joint
+void mjv_applyPerturbPose(const mjModel* m, mjData* d, const mjvPerturb* pert,
+                          int flg_paused) {
+  int rootid = 0, sel = pert->select;
+  mjtNum pos1[3], quat1[4], pos2[3], quat2[4], refpos[3], refquat[4];
+  mjtNum *Rpos, *Rquat, *Cpos, *Cquat;
+
+  // exit if nothing to do
+  if (sel<=0 || sel>=m->nbody || !(pert->active | pert->active2)) {
+    return;
+  }
+
+  // get rootid above selected body
+  rootid = m->body_rootid[sel];
+
+  // transform refpos,refquat from I-frame to X-frame of body[sel]
+  mju_negPose(pos1, quat1, m->body_ipos+3*sel, m->body_iquat+4*sel);
+  mju_mulPose(refpos, refquat, pert->refpos, pert->refquat, pos1, quat1);
+
+  // mocap body
+  if (m->body_mocapid[sel]>=0) {
+    // copy ref pose into mocap pose
+    mju_copy3(d->mocap_pos + 3*m->body_mocapid[sel], refpos);
+    mju_copy4(d->mocap_quat + 4*m->body_mocapid[sel], refquat);
+  }
+
+  // floating body, paused
+  else if (flg_paused && m->body_jntnum[sel]==1 &&
+           m->jnt_type[m->body_jntadr[sel]]==mjJNT_FREE) {
+    // copy ref pose into qpos
+    mju_copy3(d->qpos + m->jnt_qposadr[m->body_jntadr[sel]], refpos);
+    mju_copy4(d->qpos + m->jnt_qposadr[m->body_jntadr[sel]] + 3, refquat);
+  }
+
+  // child of floating body, paused
+  else if (flg_paused && m->body_jntnum[rootid]==1 &&
+           m->jnt_type[m->body_jntadr[rootid]]==mjJNT_FREE) {
+    // get pointers to root
+    Rpos = d->qpos + m->jnt_qposadr[m->body_jntadr[rootid]];
+    Rquat = Rpos + 3;
+
+    // get pointers to child
+    Cpos = d->xpos + 3*sel;
+    Cquat = d->xquat + 4*sel;
+
+    // set root <- ref*neg(child)*root
+    mju_negPose(pos1, quat1, Cpos, Cquat);                      // neg(child)
+    mju_mulPose(pos2, quat2, pos1, quat1, Rpos, Rquat);         // neg(child)*root
+    mju_mulPose(Rpos, Rquat, refpos, refquat, pos2, quat2);     // ref*neg(child)*root
+  }
+}
+
+
+
+// set perturb force,torque in d->xfrc_applied, if selected body is dynamic
+void mjv_applyPerturbForce(const mjModel* m, mjData* d, const mjvPerturb* pert) {
+  mjtNum xiquat[4], difquat[4], bvel[6], mass, stiffness, *result;
+  int sel = pert->select;
+
+  // exit if nothing to do
+  if (sel<0 ||sel>=m->nbody || !(pert->active | pert->active2)) {
+    return;
+  }
+
+  // get pointer to body xfrc_applied
+  result = d->xfrc_applied + 6*sel;
+
+  // global selbody velocity
+  mj_objectVelocity(m, d, mjOBJ_BODY, sel, bvel, 0);
+
+  // spring perturbation, with critical damping
+  //  - force
+  stiffness = m->vis.map.stiffness;
+  mass = 1.0/mju_max(mjMINVAL, m->body_invweight0[2*sel]);
+  mju_sub3(result, pert->refpos, d->xipos+3*sel);
+  mju_scl3(result, result, stiffness*mass);
+  mju_addToScl3(result, bvel+3, -sqrtf(stiffness)*mass);
+
+  //  - torque
+  stiffness = m->vis.map.stiffnessrot;
+  mass = 1.0/mju_max(mjMINVAL, m->body_invweight0[2*sel+1]);
+  mju_mulQuat(xiquat, d->xquat+4*sel, m->body_iquat+4*sel);
+  mju_negQuat(xiquat, xiquat);
+  mju_mulQuat(difquat, pert->refquat, xiquat);
+  mju_quat2Vel(result+3, difquat, 1.0/(stiffness*mass));
+  mju_addToScl3(result+3, bvel, -sqrtf(stiffness)*mass);
+
+  // mask
+  if (!((pert->active | pert->active2) & mjPERT_TRANSLATE)) {
+    mju_zero3(result);
+  }
+  if (!((pert->active | pert->active2) & mjPERT_ROTATE)) {
+    mju_zero3(result+3);
+  }
+}
+
+
+
+// return the average of two OpenGL cameras
+mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera* cam2) {
+  mjtNum pos[3], forward[3], up[3], projection, tmp1[3], tmp2[3];
+  mjvGLCamera cam;
+
+  // compute pos
+  mju_f2n(tmp1, cam1->pos, 3);
+  mju_f2n(tmp2, cam2->pos, 3);
+  mju_add3(pos, tmp1, tmp2);
+  mju_scl3(pos, pos, 0.5);
+
+  // compute forward
+  mju_f2n(tmp1, cam1->forward, 3);
+  mju_f2n(tmp2, cam2->forward, 3);
+  mju_add3(forward, tmp1, tmp2);
+  mju_normalize3(forward);
+
+  // compute up, make it orthogonal to forward
+  mju_f2n(tmp1, cam1->up, 3);
+  mju_f2n(tmp2, cam2->up, 3);
+  mju_add3(up, tmp1, tmp2);
+  projection = mju_dot3(up, forward);
+  mju_addToScl3(up, forward, -projection);
+  mju_normalize3(up);
+
+  // assign 3d quantities
+  mju_n2f(cam.pos, pos, 3);
+  mju_n2f(cam.forward, forward, 3);
+  mju_n2f(cam.up, up, 3);
+
+  // average frustum
+  cam.frustum_bottom = 0.5f * (cam1->frustum_bottom + cam2->frustum_bottom);
+  cam.frustum_top    = 0.5f * (cam1->frustum_top + cam2->frustum_top);
+  cam.frustum_center = 0.5f * (cam1->frustum_center + cam2->frustum_center);
+  cam.frustum_near   = 0.5f * (cam1->frustum_near + cam2->frustum_near);
+  cam.frustum_far    = 0.5f * (cam1->frustum_far + cam2->frustum_far);
+
+  return cam;
+}
+
+
+
+// Select geom or skin with mouse, return bodyid; -1: none selected.
+int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
+               mjtNum aspectratio, mjtNum relx, mjtNum rely,
+               const mjvScene* scn, mjtNum selpnt[3], int geomid[1], int skinid[1]) {
+  // get average camera
+  mjvGLCamera cam = mjv_averageCamera(scn->camera, scn->camera+1);
+
+  // get camera pose in model space
+  mjtNum pos[3], forward[3], up[3], left[3];
+  mjv_cameraInModel(pos, forward, up, scn);
+  mju_cross(left, up, forward);
+  mju_normalize3(left);
+
+  // compute frustum halfwidth so as to match viewport aspect ratio
+  mjtNum halfwidth = 0.5*aspectratio*(cam.frustum_top - cam.frustum_bottom);
+
+  // construct ray
+  mjtNum ray[3];
+  mju_scl3(ray, forward, cam.frustum_near);
+  mju_addToScl3(ray, up, cam.frustum_bottom + rely*(cam.frustum_top-cam.frustum_bottom));
+  mju_addToScl3(ray, left, -(cam.frustum_center + (2*relx-1)*halfwidth));
+  mju_normalize3(ray);
+
+  // find intersection with geoms
+  *geomid = -1;
+  mjtNum geomdist = mj_ray(m, d, pos, ray, vopt->geomgroup,
+                           vopt->flags[mjVIS_STATIC], -1, geomid);
+
+  // find intersection with skins
+  int bodyid = -1;
+  mjtNum skindist = -1;
+  *skinid = -1;
+  if (vopt->flags[mjVIS_SKIN]) {
+    for (int i=0; i<m->nskin; i++) {
+      // process one skin
+      int vertid;
+      mjtNum newdist = mju_raySkin(m->skin_facenum[i], m->skin_vertnum[i],
+                                   m->skin_face + 3*m->skin_faceadr[i],
+                                   scn->skinvert + 3*m->skin_vertadr[i],
+                                   pos, ray, &vertid);
+
+      // update if closer intersection found
+      if (newdist>=0 && (newdist<skindist || skindist<0)) {
+        // assign result
+        skindist = newdist;
+
+        // find body with largest weight for this vertex
+        float bestweight = -1;
+        for (int j=m->skin_boneadr[i];
+             j<m->skin_boneadr[i]+m->skin_bonenum[i];
+             j++) {
+          for (int k=m->skin_bonevertadr[j];
+               k<m->skin_bonevertadr[j]+m->skin_bonevertnum[j];
+               k++) {
+            // get vertex id and weight
+            int vid = m->skin_bonevertid[k];
+            float vweight = m->skin_bonevertweight[k];
+
+            // update if matching id and bigger weight
+            if (vid==vertid && vweight>bestweight) {
+              bestweight = vweight;
+              bodyid = m->skin_bonebodyid[j];
+              *skinid = i;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  // no intersection
+  if (geomdist<0 && skindist<0) {
+    return -1;
+  }
+
+  // geom only, or geom closer than skin
+  else if (geomdist>=0 && (skindist<0 || skindist>geomdist)) {
+    mju_addScl3(selpnt, pos, ray, geomdist);
+    *skinid = -1;
+    return m->geom_bodyid[*geomid];
+  }
+
+  // skin
+  else {
+    mju_addScl3(selpnt, pos, ray, skindist);
+    *geomid = -1;
+    return bodyid;
+  }
+}

src/engine/engine_vis_interact.h

@@ -0,0 +1,85 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_VIS_INTERACT_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_VIS_INTERACT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// transform pose from room to model space
+MJAPI void mjv_room2model(mjtNum* modelpos, mjtNum* modelquat, const mjtNum* roompos,
+                          const mjtNum* roomquat, const mjvScene* scn);
+
+// transform pose from model to room space
+MJAPI void mjv_model2room(mjtNum* roompos, mjtNum* roomquat, const mjtNum* modelpos,
+                          const mjtNum* modelquat, const mjvScene* scn);
+
+// get camera info in model space: average left and right OpenGL cameras
+MJAPI void mjv_cameraInModel(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+                             const mjvScene* scn);
+
+// get camera info in room space: average left and right OpenGL cameras
+MJAPI void mjv_cameraInRoom(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+                            const mjvScene* scn);
+
+// get frustum height at unit distance from camera; average left and right OpenGL cameras
+MJAPI mjtNum mjv_frustumHeight(const mjvScene* scn);
+
+// rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_y)
+MJAPI void mjv_alignToCamera(mjtNum* res, const mjtNum* vec, const mjtNum* forward);
+
+// move camera with mouse; action is mjtMouse
+MJAPI void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                          const mjvScene* scn, mjvCamera* cam);
+
+// move perturb object with mouse; action is mjtMouse
+MJAPI void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
+                           mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
+
+// move model with mouse; action is mjtMouse
+MJAPI void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                         const mjtNum roomup[3], mjvScene* scn);
+
+// copy perturb pos,quat from selected body; set scale perturbation
+MJAPI void mjv_initPerturb(const mjModel* m, const mjData* d,
+                           const mjvScene* scn, mjvPerturb* pert);
+
+// set perturb pos,quat in d->mocap when selected body is mocap, and in d->qpos otherwise
+//  d->qpos written only if flg_paused and subtree root for selected body has free joint
+MJAPI void mjv_applyPerturbPose(const mjModel* m, mjData* d, const mjvPerturb* pert,
+                                int flg_paused);
+
+// set perturb force,torque in d->xfrc_applied, if selected body is dynamic
+MJAPI void mjv_applyPerturbForce(const mjModel* m, mjData* d, const mjvPerturb* pert);
+
+// return the average of two OpenGL cameras
+MJAPI mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera* cam2);
+
+// Select geom or skin with mouse, return bodyid; -1: none selected.
+MJAPI int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
+                     mjtNum aspectratio, mjtNum relx, mjtNum rely,
+                     const mjvScene* scn, mjtNum selpnt[3], int geomid[1], int skinid[1]);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_VIS_INTERACT_H_

src/engine/engine_vis_visualize.h

@@ -0,0 +1,58 @@
+// Copyright 2021 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_VIS_VISUALIZE_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_VIS_VISUALIZE_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// set (type, size, pos, mat) connector-type geom between given points
+//  assume that mjv_initGeom was already called to set all other properties
+MJAPI void mjv_makeConnector(mjvGeom* geom, int type, mjtNum width,
+                             mjtNum a0, mjtNum a1, mjtNum a2,
+                             mjtNum b0, mjtNum b1, mjtNum b2);
+
+// initialize given fields when not NULL, set the rest to their default values
+MJAPI void mjv_initGeom(mjvGeom* geom, int type, const mjtNum* size,
+                        const mjtNum* pos, const mjtNum* mat, const float* rgba);
+
+// update entire scene
+MJAPI void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
+                           const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
+
+// add geoms from selected categories to existing scene
+MJAPI void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
+                        const mjvPerturb* pert, int catmask, mjvScene* scn);
+
+// make list of lights only
+MJAPI void mjv_makeLights(const mjModel* m, mjData* d, mjvScene* scn);
+
+// update camera only
+MJAPI void mjv_updateCamera(const mjModel* m, mjData* d, mjvCamera* cam, mjvScene* scn);
+
+// update skins only
+MJAPI void mjv_updateSkin(const mjModel* m, mjData* d, mjvScene* scn);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_VIS_VISUALIZE_H_