Version 2.1.2: Python bindings, OBJ assets support, bugfixes.

PiperOrigin-RevId: 434731612
Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0

doc/APIreference.rst

@@ -23,7 +23,7 @@ MuJoCo defines a large number of primitive types described here. Except for :ref
 rest of the API does not use these enum types directly. Instead it uses ints, and only the documentation/comments state
 that certain ints correspond to certain enum types. This is because we want the API to be compiler-independent, and the
 C standard does not dictate how many bytes must be used to represent an enum type. Nevertheless we recommend using these
-types when calling the API functions (and letting the compiler do the enum-to-int type cast.)
+types when calling the API functions (and letting the compiler do the enum-to-int type cast).
 
 .. _mjtNum:
 
@@ -38,7 +38,8 @@ mjtNum
        typedef float mjtNum;
    #endif
 
-| Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+| Defined in `mjtnum.h <https://github.com/deepmind/mujoco/blob/main/include/mjtnum.h>`_
+
 | This is the floating-point type used throughout the simulator. If the symbol ``mjUSEDOUBLE`` is defined in
   ``mjmodel.h``, this type is defined as ``double``, otherwise it is defined as ``float``. Currently only the
   double-precision version of MuJoCo is distributed, although the entire code base works with single-precision as well.
@@ -47,7 +48,7 @@ mjtNum
   write code that works with either single or double precision. To this end we provide math utility functions that are
   always defined with the correct floating-point type.
 
-| Note that changing ``mjUSEDOUBLE`` in ``mjmodel.h`` will not change how the library was compiled, and instead will
+| Note that changing ``mjUSEDOUBLE`` in ``mjtnum.h`` will not change how the library was compiled, and instead will
   result in numerous link errors. In general, the header files distributed with precompiled MuJoCo should never be
   changed by the user.
 
@@ -61,6 +62,7 @@ mjtByte
    typedef unsigned char mjtByte;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Byte type used to represent boolean variables.
 
 .. _mjtDisableBit:
@@ -89,6 +91,7 @@ mjtDisableBit
    } mjtDisableBit;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Constants which are powers of 2. They are used as bitmasks for the field ``disableflags`` of :ref:`mjOption`.
   At runtime this field is ``m->opt.disableflags``. The number of these constants is given by ``mjNDISABLE`` which is
   also the length of the global string array :ref:`mjDISABLESTRING` with text descriptions of these
@@ -112,6 +115,7 @@ mjtEnableBit
    } mjtEnableBit;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Constants which are powers of 2. They are used as bitmasks for the field ``enableflags`` of :ref:`mjOption`.
   At runtime this field is ``m->opt.enableflags``. The number of these constants is given by ``mjNENABLE`` which is also
   the length of the global string array :ref:`mjENABLESTRING` with text descriptions of these flags.
@@ -132,6 +136,7 @@ mjtJoint
    } mjtJoint;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Primitive joint types. These values are used in ``m->jnt_type``. The numbers in the comments indicate how many
   positional coordinates each joint type has. Note that ball joints and rotational components of free joints are
   represented as unit quaternions - which have 4 positional coordinates but 3 degrees of freedom each.
@@ -169,6 +174,7 @@ mjtGeom
    } mjtGeom;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Geometric types supported by MuJoCo. The first group are "official" geom types that can be used in the model. The
   second group are geom types that cannot be used in the model but are used by the visualizer to add decorative
   elements. These values are used in ``m->geom_type`` and ``m->site_type``.
@@ -190,6 +196,7 @@ mjtCamLight
    } mjtCamLight;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Dynamic modes for cameras and lights, specifying how the camera/light position and orientation are computed. These
   values are used in ``m->cam_mode`` and ``m->light_mode``.
 
@@ -208,6 +215,7 @@ mjtTexture
    } mjtTexture;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Texture types, specifying how the texture will be mapped. These values are used in ``m->tex_type``.
 
 .. _mjtIntegrator:
@@ -224,6 +232,7 @@ mjtIntegrator
    } mjtIntegrator;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Numerical integrator types. These values are used in ``m->opt.integrator``.
 
 .. _mjtCollision:
@@ -241,7 +250,8 @@ mjtCollision
    } mjtCollision;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
-| Collision modes specifying how candidate geom pair are generated for near-phase collision checking. These values are
+
+| Collision modes specifying how candidate geom pairs are generated for near-phase collision checking. These values are
   used in ``m->opt.collision``.
 
 .. _mjtCone:
@@ -258,6 +268,7 @@ mjtCone
    } mjtCone;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Available friction cone types. These values are used in ``m->opt.cone``.
 
 .. _mjtJacobian:
@@ -275,6 +286,7 @@ mjtJacobian
    } mjtJacobian;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Available Jacobian types. These values are used in ``m->opt.jacobian``.
 
 .. _mjtSolver:
@@ -292,6 +304,7 @@ mjtSolver
    } mjtSolver;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Available constraint solver algorithms. These values are used in ``m->opt.solver``.
 
 .. _mjtEq:
@@ -311,6 +324,7 @@ mjtEq
    } mjtEq;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Equality constraint types. These values are used in ``m->eq_type``.
 
 .. _mjtWrap:
@@ -331,6 +345,7 @@ mjtWrap
    } mjtWrap;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Tendon wrapping object types. These values are used in ``m->wrap_type``.
 
 .. _mjtTrn:
@@ -352,6 +367,7 @@ mjtTrn
    } mjtTrn;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Actuator transmission types. These values are used in ``m->actuator_trntype``.
 
 .. _mjtDyn:
@@ -371,6 +387,7 @@ mjtDyn
    } mjtDyn;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Actuator dynamics types. These values are used in ``m->actuator_dyntype``.
 
 .. _mjtGain:
@@ -388,6 +405,7 @@ mjtGain
    } mjtGain;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Actuator gain types. These values are used in ``m->actuator_gaintype``.
 
 .. _mjtBias:
@@ -406,6 +424,7 @@ mjtBias
    } mjtBias;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Actuator bias types. These values are used in ``m->actuator_biastype``.
 
 .. _mjtObj:
@@ -444,6 +463,7 @@ mjtObj
    } mjtObj;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | MuJoCo object types. These values are used in the support functions :ref:`mj_name2id` and
   :ref:`mj_id2name` to convert between object names and integer ids.
 
@@ -467,6 +487,7 @@ mjtConstraint
    } mjtConstraint;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Constraint types. These values are not used in mjModel, but are used in the mjData field ``d->efc_type`` when the list
   of active constraints is constructed at each simulation time step.
 
@@ -487,6 +508,7 @@ mjtConstraintState
    } mjtConstraintState;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | These values are used by the solver internally to keep track of the constraint states.
 
 .. _mjtSensor:
@@ -550,6 +572,7 @@ mjtSensor
    } mjtSensor;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | Sensor types. These values are used in ``m->sensor_type``.
 
 .. _mjtStage:
@@ -568,6 +591,7 @@ mjtStage
    } mjtStage;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | These are the compute stages for the skipstage parameters of :ref:`mj_forwardSkip` and
   :ref:`mj_inverseSkip`.
 
@@ -587,6 +611,7 @@ mjtDataType
    } mjtDataType;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | These are the possible sensor data types, used in ``mjData.sensor_datatype``.
 
 .. _mjtWarning:
@@ -611,6 +636,7 @@ mjtWarning
    } mjtWarning;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | Warning types. The number of warning types is given by ``mjNWARNING`` which is also the length of the array
   ``mjData.warning``.
 
@@ -646,6 +672,7 @@ mjtTimer
    } mjtTimer;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | Timer types. The number of timer types is given by ``mjNTIMER`` which is also the length of the array
   ``mjData.timer``, as well as the length of the string array :ref:`mjTIMERSTRING` with timer names.
 
@@ -665,6 +692,7 @@ mjtCatBit
    } mjtCatBit;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the available categories of geoms in the abstract visualizer. The bitmask can be used in the function
   :ref:`mjr_render` to specify which categories should be rendered.
 
@@ -687,8 +715,9 @@ mjtMouse
    } mjtMouse;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the mouse actions that the abstract visualizer recognizes. It is up to the user to intercept mouse events
-  and translate them into these actions, as illustrated in simulate.cc.
+  and translate them into these actions, as illustrated in ``simulate.cc``.
 
 .. _mjtPertBit:
 
@@ -704,8 +733,9 @@ mjtPertBit
    } mjtPertBit;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These bitmasks enable the translational and rotational components of the mouse perturbation. For the regular mouse,
-  only one can be enabled at a time. For the 3D mouse (SpaceNavigator) both can be enabled simultaneously. Tehy are used
+  only one can be enabled at a time. For the 3D mouse (SpaceNavigator) both can be enabled simultaneously. They are used
   in ``mjvPerturb.active``.
 
 .. _mjtCamera:
@@ -724,6 +754,7 @@ mjtCamera
    } mjtCamera;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the possible camera types, used in ``mjvCamera.type``.
 
 .. _mjtLabel:
@@ -754,6 +785,7 @@ mjtLabel
    } mjtLabel;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the abstract visualization elements that can have text labels. Used in ``mjvOption.label``.
 
 .. _mjtFrame:
@@ -777,6 +809,7 @@ mjtFrame
    } mjtFrame;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the MuJoCo objects whose spatial frames can be rendered. Used in ``mjvOption.frame``.
 
 .. _mjtVisFlag:
@@ -815,6 +848,7 @@ mjtVisFlag
    } mjtVisFlag;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are indices in the array ``mjvOption.flags``, whose elements enable/disable the visualization of the
   corresponding model or decoration element.
 
@@ -841,6 +875,7 @@ mjtRndFlag
    } mjtRndFlag;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are indices in the array ``mjvScene.flags``, whose elements enable/disable OpenGL rendering effects.
 
 .. _mjtStereo:
@@ -858,6 +893,7 @@ mjtStereo
    } mjtStereo;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | These are the possible stereo rendering types. They are used in ``mjvScene.stereo``.
 
 .. _mjtGridPos:
@@ -876,6 +912,7 @@ mjtGridPos
    } mjtGridPos;
 
 | Defined in `mjrender.h <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h>`_
+
 | These are the possible grid positions for text overlays. They are used as an argument to the function
   :ref:`mjr_overlay`.
 
@@ -893,6 +930,7 @@ mjtFramebuffer
    } mjtFramebuffer;
 
 | Defined in `mjrender.h <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h>`_
+
 | These are the possible framebuffers. They are used as an argument to the function :ref:`mjr_setBuffer`.
 
 .. _mjtFontScale:
@@ -913,6 +951,7 @@ mjtFontScale
    } mjtFontScale;
 
 | Defined in `mjrender.h <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h>`_
+
 | These are the possible font sizes. The fonts are predefined bitmaps stored in the dynamic library at three different
   sizes.
 
@@ -931,6 +970,7 @@ mjtFont
    } mjtFont;
 
 | Defined in `mjrender.h <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h>`_
+
 | These are the possible font types.
 
 .. _mjtButton:
@@ -949,6 +989,7 @@ mjtButton
    } mjtButton;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | Mouse button IDs used in the UI framework.
 
 .. _mjtEvent:
@@ -970,6 +1011,7 @@ mjtEvent
    } mjtEvent;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | Event types used in the UI framework.
 
 .. _mjtItem:
@@ -1002,6 +1044,7 @@ mjtItem
    } mjtItem;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | Item types used in the UI framework.
 
 .. _tyFunction:
@@ -1118,6 +1161,7 @@ mjVFS
    typedef struct _mjVFS mjVFS;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | This is the data structure with the virtual file system. It can only be constructed programmatically, and does not
   have an analog in MJCF.
 
@@ -1167,6 +1211,7 @@ mjOption
    typedef struct _mjOption mjOption;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | This is the data structure with simulation options. It corresponds to the MJCF element
   :ref:`option <option>`. One instance of it is embedded in mjModel.
 
@@ -1272,6 +1317,7 @@ mjVisual
    typedef struct _mjVisual mjVisual;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | This is the data structure with abstract visualization options. It corresponds to the MJCF element
   :ref:`visual <visual>`. One instance of it is embedded in mjModel.
 
@@ -1293,6 +1339,7 @@ mjStatistic
    typedef struct _mjStatistic mjStatistic;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | This is the data structure with model statistics precomputed by the compiler or set by the user. It corresponds to the
   MJCF element :ref:`statistic <statistic>`. One instance of it is embedded in mjModel.
 
@@ -1641,6 +1688,8 @@ mjModel
        int*      sensor_needstage;     // required compute stage (mjtStage)        (nsensor x 1)
        int*      sensor_objtype;       // type of sensorized object (mjtObj)       (nsensor x 1)
        int*      sensor_objid;         // id of sensorized object                  (nsensor x 1)
+       int*      sensor_reftype;       // type of reference frame (mjtObj)         (nsensor x 1)
+       int*      sensor_refid;         // id of reference frame; -1: global frame  (nsensor x 1)
        int*      sensor_dim;           // number of scalar outputs                 (nsensor x 1)
        int*      sensor_adr;           // address in sensor array                  (nsensor x 1)
        mjtNum*   sensor_cutoff;        // cutoff for real and positive; 0: ignore  (nsensor x 1)
@@ -1699,6 +1748,7 @@ mjModel
    typedef struct _mjModel mjModel;
 
 | Defined in `mjmodel.h <https://github.com/deepmind/mujoco/blob/main/include/mjmodel.h>`_
+
 | This is the main data structure holding the MuJoCo model. It is treated as constant by the simulator.
 
 .. _mjContact:
@@ -1739,6 +1789,7 @@ mjContact
    typedef struct _mjContact mjContact;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | This is the data structure holding information about one contact. ``mjData.contact`` is a preallocated array of
   mjContact data structures, populated at runtime with the contacts found by the collision detector. Additional contact
   information is then filled-in by the simulator.
@@ -1758,6 +1809,7 @@ mjWarningStat
    typedef struct _mjWarningStat mjWarningStat;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | This is the data structure holding information about one warning type. ``mjData.warning`` is a preallocated array of
   mjWarningStat data structures, one for each warning type.
 
@@ -1776,6 +1828,7 @@ mjTimerStat
    typedef struct _mjTimerStat mjTimerStat;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | This is the data structure holding information about one timer. ``mjData.timer`` is a preallocated array of
   mjTimerStat data structures, one for each timer type.
 
@@ -1799,6 +1852,7 @@ mjSolverStat
    typedef struct _mjSolverStat mjSolverStat;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | This is the data structure holding information about one solver iteration. ``mjData.solver`` is a preallocated array
   of mjSolverStat data structures, one for each iteration of the solver, up to a maximum of mjNSOLVER. The actual number
   of solver iterations is given by ``mjData.solver_iter``.
@@ -1861,19 +1915,19 @@ mjData
        mjtNum*   qfrc_applied;         // applied generalized force                (nv x 1)
        mjtNum*   xfrc_applied;         // applied Cartesian force/torque           (nbody x 6)
 
-       // dynamics
-       mjtNum*   qacc;                 // acceleration                             (nv x 1)
-       mjtNum*   act_dot;              // time-derivative of actuator activation   (na x 1)
-
        // mocap data
        mjtNum*  mocap_pos;             // positions of mocap bodies                (nmocap x 3)
        mjtNum*  mocap_quat;            // orientations of mocap bodies             (nmocap x 4)
 
+       // dynamics
+       mjtNum*   qacc;                 // acceleration                             (nv x 1)
+       mjtNum*   act_dot;              // time-derivative of actuator activation   (na x 1)
+
        // user data
-       mjtNum*  userdata;              // user data, not touched by engine         (nuserdata x 1)
+       mjtNum*   userdata;              // user data, not touched by engine        (nuserdata x 1)
 
        // sensors
-       mjtNum*  sensordata;            // sensor data array                        (nsensordata x 1)
+       mjtNum*   sensordata;            // sensor data array                       (nsensordata x 1)
 
        //-------------------------------- POSITION dependent
 
@@ -2005,6 +2059,7 @@ mjData
    typedef struct _mjData mjData;
 
 | Defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mjdata.h>`_
+
 | This is the main data structure holding the simulation state. It is the workspace where all functions read their
   modifiable inputs and write their outputs.
 
@@ -2028,6 +2083,7 @@ mjvPerturb
    typedef struct _mjvPerturb mjvPerturb;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This is the data structure holding information about mouse perturbations.
 
 .. _mjvCamera:
@@ -2053,6 +2109,7 @@ mjvCamera
    typedef struct _mjvCamera mjvCamera;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This is the data structure describing one abstract camera.
 
 .. _mjvGLCamera:
@@ -2079,6 +2136,7 @@ mjvGLCamera
    typedef struct _mjvGLCamera mjvGLCamera;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This is the data structure describing one OpenGL camera.
 
 .. _mjvGeom:
@@ -2121,6 +2179,7 @@ mjvGeom
    typedef struct _mjvGeom mjvGeom;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This is the data structure describing one abstract visualization geom - which could correspond to a model geom or to a
   decoration element constructed by the visualizer.
 
@@ -2148,6 +2207,7 @@ mjvLight
    typedef struct _mjvLight mjvLight;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This is the data structure describing one OpenGL light.
 
 .. _mjvOption:
@@ -2171,6 +2231,7 @@ mjvOption
    typedef struct _mjvOption mjvOption;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This structure contains options that enable and disable the visualization of various elements.
 
 .. _mjvScene:
@@ -2198,7 +2259,7 @@ mjvScene
 
        // OpenGL lights
        int      nlight;                // number of lights currently in buffer
-       mjvLight lights[8];             // buffer for lights
+       mjvLight lights[mjMAXLIGHT];    // buffer for lights
 
        // OpenGL cameras
        mjvGLCamera camera[2];          // left and right camera
@@ -2216,6 +2277,7 @@ mjvScene
    typedef struct _mjvScene mjvScene;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This structure contains everything needed to render the 3D scene in OpenGL.
 
 .. _mjvFigure:
@@ -2269,6 +2331,7 @@ mjvFigure
    typedef struct _mjvFigure mjvFigure;
 
 | Defined in `mjvisualize.h <https://github.com/deepmind/mujoco/blob/main/include/mjvisualize.h>`_
+
 | This structure contains everything needed to render a 2D plot in OpenGL. The buffers for line points etc. are
   preallocated, and the user has to populate them before calling the function :ref:`mjr_figure` with this
   data structure as an argument.
@@ -2290,6 +2353,7 @@ mjrRect
    typedef struct _mjrRect mjrRect;
 
 | Defined in `mjrender.h (57) <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h#L57>`_
+
 | This structure specifies a rectangle.
 
 .. _mjrContext:
@@ -2384,6 +2448,7 @@ mjrContext
    typedef struct _mjrContext mjrContext;
 
 | Defined in `mjrender.h (67) <https://github.com/deepmind/mujoco/blob/main/include/mjrender.h#L67>`_
+
 | This structure contains the custom OpenGL rendering context, with the ids of all OpenGL resources uploaded to the GPU.
 
 .. _mjuiState:
@@ -2434,6 +2499,7 @@ mjuiState
    typedef struct _mjuiState mjuiState;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure contains the keyboard and mouse state used by the UI framework.
 
 .. _mjuiThemeSpacing:
@@ -2460,6 +2526,7 @@ mjuiThemeSpacing
    typedef struct _mjuiThemeSpacing mjuiThemeSpacing;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines the spacing of UI items in the theme.
 
 .. _mjuiThemeColor:
@@ -2496,6 +2563,7 @@ mjuiThemeColor
    typedef struct _mjuiThemeColor mjuiThemeColor;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines the colors of UI items in the theme.
 
 .. _mjuiItem:
@@ -2553,6 +2621,7 @@ mjuiItem
    typedef struct _mjuiItem mjuiItem;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines one UI item.
 
 .. _mjuiSection:
@@ -2579,6 +2648,7 @@ mjuiSection
    typedef struct _mjuiSection mjuiSection;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines one section of the UI.
 
 .. _mjUI:
@@ -2625,6 +2695,7 @@ mjUI
    typedef struct _mjUI mjUI;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines the entire UI.
 
 .. _mjuiDef:
@@ -2645,6 +2716,7 @@ mjuiDef
    typedef struct _mjuiDef mjuiDef;
 
 | Defined in `mjui.h <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`_
+
 | This structure defines one entry in the definition table used for simplified UI construction.
 
 .. _tyXMacro:
@@ -2717,7 +2789,7 @@ Global variables
 Error callbacks
 ^^^^^^^^^^^^^^^
 
-All user callbacks (i.e. global function pointers whose name starts with 'mjcb') are initially set to NULL, which
+All user callbacks (i.e., global function pointers whose name starts with 'mjcb') are initially set to NULL, which
 disables them and allows the default processing to take place. To install a callback, simply set the corresponding
 global pointer to a user function of the correct type. Keep in mind that these are global and not model-specific. So if
 you are simulating multiple models in parallel, they use the same set of callbacks.
@@ -2732,7 +2804,7 @@ mju_user_error
    extern void (*mju_user_error)(const char*);
 
 This is called from within the main error function :ref:`mju_error`. When installed, this function overrides the default
-error processing. Once it prints error mesages (or whatever else the user wants to do), it must **exit** the program.
+error processing. Once it prints error messages (or whatever else the user wants to do), it must **exit** the program.
 MuJoCo is written with the assumption that mju_error will not return. If it does, the behavior of the software is
 undefined.
 
@@ -2878,7 +2950,7 @@ mjcb_time
    extern mjfTime mjcb_time;
 
 Installing this callback enables the built-in profiler, and keeps timing statistics in ``mjData.timer``. The return type
-is mjtNum, while the time units are up to the user. simulate.cc assumes the unit is 1 millisecond. In order to be
+is mjtNum, while the time units are up to the user. ``simulate.cc`` assumes the unit is 1 millisecond. In order to be
 useful, the callback should use high-resolution timers with at least microsecond precision. This is because the
 computations being timed are very fast.
 
@@ -3022,7 +3094,7 @@ mjVISSTRING
 | [1]: the string "0" or "1" indicating if the flag is on or off by default, as set by
   :ref:`mjv_defaultOption`;
 
-| [2]: one-character string with a suggested keyboard shortcut, used in simulate.cc.
+| [2]: one-character string with a suggested keyboard shortcut, used in ``simulate.cc``.
 
 .. _mjRNDSTRING:
 
@@ -3156,8 +3228,8 @@ API functions
 The main header `mujoco.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco.h>`_ exposes a very large number
 of functions. However the functions that most users are likely to need are a small fraction. For example,
 ``simulate.cc`` which is as elaborate as a MuJoCo application is likely to get, calls around 40 of these functions,
-while basic.cc calls around 20. The rest are explosed just in case someone has a use for them. This includes us as users
+while ``basic.cc`` calls around 20. The rest are explosed just in case someone has a use for them. This includes us as
-of MuJoCo -- we do our own work with the public library instead of relying on internal builds.
+users of MuJoCo -- we do our own work with the public library instead of relying on internal builds.
 
 .. _Activation:
 
@@ -3672,9 +3744,9 @@ mj_printFormattedModel
 
 .. code-block:: C
 
-   void mj_printFormattedModel(const mjModel* m, const char* filename, const char* float_format_str);
+   void mj_printFormattedModel(const mjModel* m, const char* filename, const char* float_format);
 
-Print ``mjModel`` to text file, specifying format. ``float_format_str`` must be a valid printf-style format string for a
+Print ``mjModel`` to text file, specifying format. ``float_format`` must be a valid printf-style format string for a
 single float value.
 
 .. _mj_printModel:
@@ -3688,6 +3760,18 @@ mj_printModel
 
 Print model to text file.
 
+.. _mj_printFormattedData:
+
+mj_printFormattedData
+~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: C
+
+   void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, const char* float_format);
+
+Print ``mjData`` to text file, specifying format. ``float_format`` must be a valid printf-style format string for a
+single float value.
+
 .. _mj_printData:
 
 mj_printData
@@ -4140,7 +4224,7 @@ mj_constraintUpdate
 .. code-block:: C
 
    void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
-                            mjtNum* cost, int flg_coneHessian);
+                            mjtNum cost[1], int flg_coneHessian);
 
 Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set \*cost = s(jar)
 where jar = Jac*qacc-aref.
@@ -4380,8 +4464,8 @@ mj_applyFT
 .. code-block:: C
 
    void mj_applyFT(const mjModel* m, mjData* d,
-                   const mjtNum* force, const mjtNum* torque,
+                   const mjtNum[3] force, const mjtNum[3] torque,
-                   const mjtNum* point, int body, mjtNum* qfrc_target);
+                   const mjtNum[3] point, int body, mjtNum* qfrc_target);
 
 This function can be used to apply a Cartesian force and torque to a point on a body, and add the result to the vector
 mjData.qfrc_applied of all applied forces. Note that the function requires a pointer to this vector, because sometimes
@@ -4395,7 +4479,7 @@ mj_objectVelocity
 .. code-block:: C
 
    void mj_objectVelocity(const mjModel* m, const mjData* d,
-                          int objtype, int objid, mjtNum* res, int flg_local);
+                          int objtype, int objid, mjtNum[6] res, int flg_local);
 
 Compute object 6D velocity in object-centered frame, world/local orientation.
 
@@ -4407,7 +4491,7 @@ mj_objectAcceleration
 .. code-block:: C
 
    void mj_objectAcceleration(const mjModel* m, const mjData* d,
-                              int objtype, int objid, mjtNum* res, int flg_local);
+                              int objtype, int objid, mjtNum[6] res, int flg_local);
 
 Compute object 6D acceleration in object-centered frame, world/local orientation.
 
@@ -4418,7 +4502,7 @@ mj_contactForce
 
 .. code-block:: C
 
-   void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum* result);
+   void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum[6] result);
 
 Extract 6D force:torque given contact id, in the contact frame.
 
@@ -4468,8 +4552,8 @@ mj_local2Global
 
 .. code-block:: C
 
-   void mj_local2Global(mjData* d, mjtNum* xpos, mjtNum* xmat,
+   void mj_local2Global(mjData* d, mjtNum[3] xpos, mjtNum[9] xmat,
-                        const mjtNum* pos, const mjtNum* quat,
+                        const mjtNum[3] pos, const mjtNum[4] quat,
                         int body, mjtByte sameframe);
 
 Map from body local to global Cartesian coordinates.
@@ -4539,9 +4623,9 @@ mj_ray
 
 .. code-block:: C
 
-   mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum* vec,
+   mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum[3] pnt, const mjtNum[3] vec,
                  const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
-                 int* geomid);
+                 int geomid[1]);
 
 Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude. Return geomid and distance (x) to
 nearest surface, or -1 if no intersection.
@@ -4559,7 +4643,7 @@ mj_rayHfield
 .. code-block:: C
 
    mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
-                       const mjtNum* pnt, const mjtNum* vec);
+                       const mjtNum[3] pnt, const mjtNum[3] vec);
 
 Interect ray with hfield, return nearest distance or -1 if no intersection.
 
@@ -4571,7 +4655,7 @@ mj_rayMesh
 .. code-block:: C
 
    mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
-                     const mjtNum* pnt, const mjtNum* vec);
+                     const mjtNum[3] pnt, const mjtNum[3] vec);
 
 Interect ray with mesh, return nearest distance or -1 if no intersection.
 
@@ -4582,8 +4666,8 @@ mju_rayGeom
 
 .. code-block:: C
 
-   mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+   mjtNum mju_rayGeom(const mjtNum[3] pos, const mjtNum[9] mat, const mjtNum[3] size,
-                      const mjtNum* pnt, const mjtNum* vec, int geomtype);
+                      const mjtNum[3] pnt, const mjtNum[3] vec, int geomtype);
 
 Interect ray with pure geom, return nearest distance or -1 if no intersection.
 
@@ -4595,7 +4679,7 @@ mju_raySkin
 .. code-block:: C
 
    mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
-                      const mjtNum* pnt, const mjtNum* vec, int* vertid);
+                      const mjtNum[3] pnt, const mjtNum[3] vec, int vertid[1]);
 
 Interect ray with skin, return nearest vertex id.
 
@@ -4605,7 +4689,7 @@ Interaction
 ^^^^^^^^^^^
 
 These function implement abstract mouse interactions, allowing control over cameras and perturbations. Their use is well
-illustrated in simulate.cc.
+illustrated in ``simulate.cc``.
 
 .. _mjv_defaultCamera:
 
@@ -4636,8 +4720,8 @@ mjv_room2model
 
 .. code-block:: C
 
-   void mjv_room2model(mjtNum* modelpos, mjtNum* modelquat, const mjtNum* roompos,
+   void mjv_room2model(mjtNum[3] modelpos, mjtNum[4] modelquat, const mjtNum[3] roompos,
-                       const mjtNum* roomquat, const mjvScene* scn);
+                       const mjtNum[4] roomquat, const mjvScene* scn);
 
 Transform pose from room to model space.
 
@@ -4648,8 +4732,8 @@ mjv_model2room
 
 .. code-block:: C
 
-   void mjv_model2room(mjtNum* roompos, mjtNum* roomquat, const mjtNum* modelpos,
+   void mjv_model2room(mjtNum[3] roompos, mjtNum[4] roomquat, const mjtNum[3] modelpos,
-                       const mjtNum* modelquat, const mjvScene* scn);
+                       const mjtNum[4] modelquat, const mjvScene* scn);
 
 Transform pose from model to room space.
 
@@ -4660,7 +4744,7 @@ mjv_cameraInModel
 
 .. code-block:: C
 
-   void mjv_cameraInModel(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+   void mjv_cameraInModel(mjtNum[3] headpos, mjtNum[3] forward, mjtNum[3] up,
                           const mjvScene* scn);
 
 Get camera info in model space; average left and right OpenGL cameras.
@@ -4672,7 +4756,7 @@ mjv_cameraInRoom
 
 .. code-block:: C
 
-   void mjv_cameraInRoom(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+   void mjv_cameraInRoom(mjtNum[3] headpos, mjtNum[3] forward, mjtNum[3] up,
                          const mjvScene* scn);
 
 Get camera info in room space; average left and right OpenGL cameras.
@@ -4695,7 +4779,7 @@ mjv_alignToCamera
 
 .. code-block:: C
 
-   void mjv_alignToCamera(mjtNum* res, const mjtNum* vec, const mjtNum* forward);
+   void mjv_alignToCamera(mjtNum[3] res, const mjtNum[3] vec, const mjtNum[3] forward);
 
 Rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_y).
 
@@ -4731,7 +4815,7 @@ mjv_moveModel
 .. code-block:: C
 
    void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
-                      const mjtNum* roomup, mjvScene* scn);
+                      const mjtNum roomup[3], mjvScene* scn);
 
 Move model with mouse; action is mjtMouse.
 
@@ -4791,13 +4875,13 @@ mjv_select
 
    int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
                   mjtNum aspectratio, mjtNum relx, mjtNum rely,
-                  const mjvScene* scn, mjtNum* selpnt, int* geomid, int* skinid);
+                  const mjvScene* scn, mjtNum[3] selpnt, int geomid[1], int skinid[1]);
 
 This function is used for mouse selection. Previously selection was done via OpenGL, but as of MuJoCo 1.50 it relies on
 ray intersections which are much more efficient. aspectratio is the viewport width/height. relx and rely are the
 relative coordinates of the 2D point of interest in the viewport (usually mouse cursor). The function returns the id of
 the geom under the specified 2D point, or -1 if there is no geom (note that they skybox if present is not a model geom).
-The 3D coordinates of the clicked point are returned in selpnt. See simulate.cc for an illustration.
+The 3D coordinates of the clicked point are returned in selpnt. See ``simulate.cc`` for an illustration.
 
 .. _Visualization-api:
 
@@ -4806,7 +4890,7 @@ Visualization
 
 The functions in this section implement abstract visualization. The results are used by the OpenGL rendered, and can
 also be used by users wishing to implement their own rendered, or hook up MuJoCo to advanced rendering tools such as
-Unity or Unreal Engine. See simulate.cc for illustration of how to use these functions.
+Unity or Unreal Engine. See ``simulate.cc`` for illustration of how to use these functions.
 
 .. _mjv_defaultOption:
 
@@ -4837,8 +4921,8 @@ mjv_initGeom
 
 .. code-block:: C
 
-   void mjv_initGeom(mjvGeom* geom, int type, const mjtNum* size,
+   void mjv_initGeom(mjvGeom* geom, int type, const mjtNum[3] size,
-                     const mjtNum* pos, const mjtNum* mat, const float* rgba);
+                     const mjtNum[3] pos, const mjtNum[9] mat, const float[4] rgba);
 
 Initialize given geom fields when not NULL, set the rest to their default values.
 
@@ -4951,7 +5035,7 @@ Update skins.
 OpenGL rendering
 ^^^^^^^^^^^^^^^^
 
-These functions expose the OpenGL renderer. See simulate.cc for illustration of how to use these functions.
+These functions expose the OpenGL renderer. See ``simulate.cc`` for illustration of how to use these functions.
 
 .. _mjr_defaultContext:
 
@@ -6185,7 +6269,7 @@ mju_cholFactor
 
    int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
 
-Cholesky decomposition: mat = L*L'; return rank.
+Cholesky decomposition: mat = L*L'; return rank, decomposition performed in-place into mat.
 
 .. _mju_cholSolve:
 
@@ -6214,7 +6298,7 @@ mju_eig3
 
 .. code-block:: C
 
-   int mju_eig3(mjtNum* eigval, mjtNum* eigvec, mjtNum* quat, const mjtNum* mat);
+   int mju_eig3(mjtNum[3] eigval, mjtNum[9] eigvec, mjtNum[4] quat, const mjtNum[9] mat);
 
 Eigenvalue decomposition of symmetric 3x3 matrix.
 

doc/XMLreference.rst

@@ -1214,6 +1214,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`framequat`      | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1223,6 +1225,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`framexaxis`     | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1232,6 +1236,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`frameyaxis`     | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1241,6 +1247,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`framezaxis`     | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1250,6 +1258,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`framelinvel`    | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1259,6 +1269,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`frameangvel`    | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1268,6 +1280,8 @@ XML schema
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`cutoff`            | :at:`noise`             | :at:`user`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`reftype`           | :at:`refname`           |                         | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 | |_2|:el:`framelinacc`    | \* | .. table::                                                                         |
 |                          |    |    :class: mjcf-attributes                                                         |
@@ -1641,10 +1655,11 @@ adjust it properly through the XML.
    around this convention (both the camera and perturbation commands are based on it) so we do not recommend deviating
    from it.
 :at:`wind`: :at-val:`real(3), "0 0 0"`
-   Velocity vector of the medium (i.e. wind). This vector is subtracted from the 3D translational velocity of each body,
+   Velocity vector of the medium (i.e., wind). This vector is subtracted from the 3D translational velocity of each
-   and the result is used to compute viscous, lift and drag forces acting on the body; recall :ref:`Passive forces
+   body, and the result is used to compute viscous, lift and drag forces acting on the body; recall :ref:`Passive forces
    <gePassive>` in the Computation chapter. The magnitude of these forces scales with the values of the next two
    attributes.
+
 :at:`magnetic`: :at-val:`real(3), "0 -0.5 0"`
    Global magnetic flux. This vector is used by magnetometer sensors, which are defined as sites and return the magnetic
    flux at the site position expressed in the site frame.
@@ -1744,7 +1759,7 @@ from its default.
    This flag disables the clamping of control inputs to all actuators, even if the actuator-specific attributes are set
    to enable clamping.
 :at:`warmstart`: :at-val:`[disable, enable], "enable"`
-   This flag disables warm-starting of the constraint solver. By default the solver uses the solution (i.e. the
+   This flag disables warm-starting of the constraint solver. By default the solver uses the solution (i.e., the
    constraint force) from the previous time step to initialize the iterative optimization. This feature should be
    disabled when evaluating the dynamics at a collection of states that do not form a trajectory - in which case warm
    starts make no sense and are likely to slow down the solver.
@@ -1785,7 +1800,7 @@ fields of mjOption which can be modified at runtime, sizes are structural parame
 compilation.
 
 :at:`njmax`: :at-val:`int, "-1"`
-   This and the next two attributes specify the maximum sizes of the dynamic arrays in mjData, i.e. arrays whose
+   This and the next two attributes specify the maximum sizes of the dynamic arrays in mjData, i.e., arrays whose
    effective length varies at runtime. This attribute specifies the maximum number of scalar constraints (or
    equivalently, rows of the constraint Jacobian) that can be handled at runtime. If the number of active constraints is
    about to exceed this maximum (usually because too many contacts become active) the extra constraints are discarded
@@ -1801,7 +1816,7 @@ compilation.
    warning is generated. The actual number of contacts is stored in mjData.ncon. If this value is negative, the compiler
    will use a heuristic to guess an appropriate number.
 :at:`nstack`: :at-val:`int, "-1"`
-   This attribute specifies the size of the pre-allocated stack in mjData, in units of sizeof(mjtNum) which is currently
+   This attribute specifies the size of the preallocated stack in mjData, in units of sizeof(mjtNum) which is currently
    defined as double; thus the size in bytes is 8 times larger. The custom stack is used by all MuJoCo functions that
    need dynamically allocated memory. We do not use heap memory allocation at runtime, so as to speed up processing as
    well as avoid heap fragmentation. Note that the internal allocator keeps track of how much stack space has ever been
@@ -1814,23 +1829,23 @@ compilation.
    The number of key frames allocated in mjModel is the larger of this value and the number of :ref:`key <key>` elements
    below. Note that the interactive simulator has the ability to take snapshots of the system state and save them as key
    frames.
-:at:`nuser_body`: :at-val:`int, "0"`
+:at:`nuser_body`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each body. See also :ref:`User parameters <CUser>`.
    The parameter values are set via the user attribute of the :ref:`body <body>` element. These values are not accessed
    by MuJoCo. They can be used to define element properties needed in user callbacks and other custom code.
-:at:`nuser_jnt`: :at-val:`int, "0"`
+:at:`nuser_jnt`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`joint <joint>`.
-:at:`nuser_geom`: :at-val:`int, "0"`
+:at:`nuser_geom`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`geom <geom>`.
-:at:`nuser_site"`: :at-val:`int, "0"`
+:at:`nuser_site`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`site <site>`.
-:at:`nuser_cam"`: :at-val:`int, "0"`
+:at:`nuser_cam`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`camera <camera>`.
-:at:`nuser_tendon`: :at-val:`int, "0"`
+:at:`nuser_tendon`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`tendon <tendon>`.
-:at:`nuser_actuator`: :at-val:`int, "0"`
+:at:`nuser_actuator`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`actuator <actuator>`.
-:at:`nuser_sensor`: :at-val:`int, "0"`
+:at:`nuser_sensor`: :at-val:`int, "-1"`
    The number of custom user parameters added to the definition of each :ref:`sensor <sensor>`.
 
 .. _visual:
@@ -1843,7 +1858,7 @@ compilation.
   yields a list of geometric entities for subsequent rendering. The settings here are global, in contrast with the
   element-specific visual settings. The global and element-specific settings refer to non-overlapping properties. Some
   of the global settings affect properties such as triangulation of geometric primitives that cannot be set per element.
-  Other global settings affect the properties of decorative objects, i.e. objects such as contact points and force
+  Other global settings affect the properties of decorative objects, i.e., objects such as contact points and force
   arrows which do not correspond to model elements. The visual settings are grouped semantically into several
   subsections.
 | This element is a good candidate for the :ref:`file include <CInclude>` mechanism. One can create an XML file with
@@ -1858,7 +1873,7 @@ While all settings in mjVisual are global, the settings here could not be fit in
 is effectively a miscellaneous subsection.
 
 :at:`fovy`: :at-val:`real, "45"`
-   This attribute specifies the vertical field of view of the free camera, i.e. the camera that is always available in
+   This attribute specifies the vertical field of view of the free camera, i.e., the camera that is always available in
    the visualizer even if no cameras are explicitly defined in the model. It is always expressed in degrees, regardless
    of the setting of the angle attribute of :ref:`compiler <compiler>`, and is also represented in the low level model
    in degrees. This is because we pass it to OpenGL which uses degrees. The same convention applies to the fovy
@@ -2013,7 +2028,7 @@ documented below.
 :at:`light`: :at-val:`real, "0.3"`
    The size of the decorative object used to represent model lights in the rendering.
 :at:`selectpoint`: :at-val:`real, "0.2"`
-   The radius of the sphere used to render the selection point (i.e. the point where the user left-double-clicked to
+   The radius of the sphere used to render the selection point (i.e., the point where the user left-double-clicked to
    select a body). Note that the local and global coordinates of this point can be printed in the 3D view by activating
    the corresponding rendering flags. In this way, the coordinates of points of interest can be found.
 :at:`jointlength`: :at-val:`real, "1.0"`
@@ -2093,7 +2108,7 @@ disables the rendering of the corresponding object.
    Color of slider-crank mechanisms.
 :at:`crankbroken`: :at-val:`real(4), "0.9 0 0 1"`
    Color used to render the crank of slide-crank mechanisms, in model configurations where the specified rod length
-   cannot be maintained, i.e. it is "broken".
+   cannot be maintained, i.e., it is "broken".
 
 .. _statistic:
 
@@ -2458,7 +2473,7 @@ chapter.
    different from "none", the texture is treated as procedural and any file names are ignored. The keywords have the
    following meaning:
    The **gradient** type generates a color gradient from rgb1 to rgb2. The interpolation in color space is done through
-   a sigmoid function. For cube and skybox textures the gradient is along the +Y axis, i.e. from top to bottom for
+   a sigmoid function. For cube and skybox textures the gradient is along the +Y axis, i.e., from top to bottom for
    skybox rendering.
 
    The **checker** type generates a 2-by-2 checker pattern with alternating colors given by rgb1 to rgb2. This is
@@ -2486,11 +2501,11 @@ chapter.
    texture size and probability need to be adjusted jointly. Together with a gradient skybox texture, this can create
    the appearance of a night sky with stars.
 :at:`width`: :at-val:`int, "0"`
-   The width of the procedural texture, i.e. the number of columns in the image. For cube and skybox procedural textures
+   The width of the procedural texture, i.e., the number of columns in the image. For cube and skybox procedural
-   the width and height must be equal. Larger values usually result in higher quality images, although in some cases
+   textures the width and height must be equal. Larger values usually result in higher quality images, although in some
-   (e.g. checker patterns) small values are sufficient.
+   cases (e.g. checker patterns) small values are sufficient.
 :at:`height`: :at-val:`int, "0"`
-   The height of the procedural texture, i.e. the number of rows in the image.
+   The height of the procedural texture, i.e., the number of rows in the image.
 :at:`hflip`: :at-val:`[false, true], "false"`
    If true, images loaded from file are flipped in the horizontal direction. Does not affect procedural textures.
 :at:`vflip`: :at-val:`[false, true], "false"`
@@ -2548,7 +2563,7 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe
    with the file contents.
 :at:`nrow`: :at-val:`int, "0"`
    This attribute and the next are used to allocate a height field in mjModel and leave the elevation data undefined
-   (i.e. set to 0). This attribute specifies the number of rows in the elevation data matrix. The default value of 0
+   (i.e., set to 0). This attribute specifies the number of rows in the elevation data matrix. The default value of 0
    means that the data will be loaded from a file, which will be used to infer the size of the matrix.
 :at:`ncol`: :at-val:`int, "0"`
    This attribute specifies the number of columns in the elevation data matrix.
@@ -2962,7 +2977,7 @@ axes of inertia of the body. Thus the inertia matrix is diagonal in this frame.
 :at:`fullinertia`: :at-val:`real(6), optional`
    Full inertia matrix M. Since M is 3-by-3 and symmetric, it is specified using only 6 numbers in the following order:
    M(1,1), M(2,2), M(3,3), M(1,2), M(1,3), M(2,3). The compiler computes the eigenvalue decomposition of M and sets the
-   frame orientation and diagonal inertia accordingly. If non-positive eigenvalues are encountered (i.e. if M is not
+   frame orientation and diagonal inertia accordingly. If non-positive eigenvalues are encountered (i.e., if M is not
    positive definite) a compile error is generated.
 
 .. _joint:
@@ -3058,7 +3073,7 @@ unit quaternions.
    Armature inertia (or rotor inertia, or reflected inertia) of all degrees of freedom created by this joint. These are
    constants added to the diagonal of the inertia matrix in generalized coordinates. They make the simulation more
    stable, and often increase physical realism. This is because when a motor is attached to the system with a
-   transmission that amplifies the motor force by c, the inertia of the rotor (i.e. the moving part of the motor) is
+   transmission that amplifies the motor force by c, the inertia of the rotor (i.e., the moving part of the motor) is
    amplified by c*c. The same holds for gears in the early stages of planetary gear boxes. These extra inertias often
    dominate the inertias of the robot parts that are represented explicitly in the model, and the armature attribute is
    the way to model them.
@@ -3280,7 +3295,7 @@ mjModel. If the XML model is saved, it will appear as a regular joint of type "f
    in :ref:`CSolver`. The quantity this function is applied to is the distance between
    the two geoms minus the margin plus the gap.
 :at:`gap`: :at-val:`real, "0"`
-   This attribute is used to enable the generation of inactive contacts, i.e. contacts that are ignored by the
+   This attribute is used to enable the generation of inactive contacts, i.e., contacts that are ignored by the
    constraint solver but are included in mjData.contact for the purpose of custom computations. When this value is
    positive, geom distances between margin and margin-gap correspond to such inactive contacts.
 :at:`fromto`: :at-val:`real(6), optional`
@@ -3380,7 +3395,7 @@ and the +Y axis points up. Thus the frame position and orientation are the key a
 :at:`mode`: :at-val:`[fixed, track, trackcom, targetbody, targetbodycom], "fixed"`
    This attribute specifies how the camera position and orientation in world coordinates are computed in forward
    kinematics (which in turn determine what the camera sees). "fixed" means that the position and orientation specified
-   below are fixed relative to the parent (i.e. the body where the camera is defined). "track" means that the camera
+   below are fixed relative to the parent (i.e., the body where the camera is defined). "track" means that the camera
    position is at a constant offset from the parent in world coordinates, while the camera orientation is constant in
    world coordinates. These constants are determined by applying forward kinematics in qpos0 and treating the camera as
    fixed. Tracking can be used for example to position a camera above a body, point it down so it sees the body, and
@@ -3447,7 +3462,7 @@ the direction specified by the dir attribute. It does not have a full spatial fr
    computed by the compiler but can also be overridden by specifying the extent attribute of :ref:`statistic
    <statistic>`. Internally the shadow-mapping mechanism renders the scene from the light viewpoint (as if it were a
    camera) into a depth texture, and then renders again from the camera viewpoint, using the depth texture to create
-   shadows. The internal rendering pass uses the same near and far clipping planes as regular rendering, i.e. these
+   shadows. The internal rendering pass uses the same near and far clipping planes as regular rendering, i.e., these
    clipping planes bound the cone or box shadow volume in the light direction. As a result, some shadows (especially
    those very close to the light) may be clipped.
 :at:`active`: :at-val:`[false, true], "true"`
@@ -3765,7 +3780,7 @@ element.
 :at:`margin`: :at-val:`real, "0"`
    Distance threshold below which contacts are detected and included in the global array mjData.contact.
 :at:`gap`: :at-val:`real, "0"`
-   This attribute is used to enable the generation of inactive contacts, i.e. contacts that are ignored by the
+   This attribute is used to enable the generation of inactive contacts, i.e., contacts that are ignored by the
    constraint solver but are included in mjData.contact for the purpose of custom computations. When this value is
    positive, geom distances between margin and margin-gap correspond to such inactive contacts.
 
@@ -3850,7 +3865,7 @@ of the other body, without any joint elements in the child body.
    and changing the corresponding component of mjModel.eq_active at runtime can be used to fix the body temporarily.
 :at:`relpose`: :at-val:`real(7), "0 1 0 0 0 0 0"`
    This attribute specifies the relative pose (3D position followed by 4D quaternion orientation) of body2 relative to
-   body1. If the quaternion part (i.e. last 4 components of the vector) are all zeros, as in the default setting, this
+   body1. If the quaternion part (i.e., last 4 components of the vector) are all zeros, as in the default setting, this
    attribute is ignored and the relative pose is the one corresponding to the model reference pose in qpos0. The unusual
    default is because all equality constraint types share the same default for their numeric parameters.
 
@@ -4145,7 +4160,7 @@ specify them independently.
    If specified, the actuator acts on the given tendon. The actuator length equals the tendon length times the gear
    ratio. Both spatial and fixed tendons can be used.
 :at:`cranksite`: :at-val:`string, optional`
-   If specified, the actuator acts on a slider-crank mechanism which is implicitly determined by the actuator (i.e. it
+   If specified, the actuator acts on a slider-crank mechanism which is implicitly determined by the actuator (i.e., it
    is not a separate model element). The specified site corresponds to the pin joining the crank and the connecting rod.
    The actuator length equals the position of the slider-crank mechanism times the gear ratio.
 :at:`slidersite`: :at-val:`string, required for slider-crank transmission`
@@ -4217,7 +4232,7 @@ specify them independently.
   such shortcut is encountered, the parser creates a :el:`general` actuator and sets its dynprm, gainprm and biasprm
   attributes to the internal defaults shown above, regardless of any default settings. It then adjusts dyntype, gaintype
   and biastype depending on the shortcut, parses any custom attributes (beyond the common ones), and translates them
-  into regular attributes (i.e. attributes of the :el:`general` actuator type) as explained here.
+  into regular attributes (i.e., attributes of the :el:`general` actuator type) as explained here.
 | This element creates a direct-drive actuator. The underlying :el:`general` attributes are set as follows:
 
 ========= ======= ========= =======
@@ -4464,7 +4479,7 @@ This element creates a 3-axis force sensor. The sensor outputs three numbers, wh
 child and a parent body, expressed in the site frame defining the sensor. The convention is that the site is attached to
 the child body, and the force points from the child towards the parent. To change the sign of the sensor reading, use
 the scale attribute. The computation here takes into account all forces acting on the system, including contacts as well
-as external perturbations. Using this sensor often requires creating a dummy body welded to its parent (i.e. having no
+as external perturbations. Using this sensor often requires creating a dummy body welded to its parent (i.e., having no
 joint elements).
 
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
@@ -4724,7 +4739,8 @@ This element creates a tendon limit sensor for constraint force.
 :el-prefix:`sensor/` **framepos** (*)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-This element creates a sensor that returns the 3D position of the spatial frame of the object, in global coordinates.
+This element creates a sensor that returns the 3D position of the spatial frame of the object, in global coordinates or
+optionally with respect to a given frame-of-reference.
 
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
@@ -4734,6 +4750,12 @@ This element creates a sensor that returns the 3D position of the spatial frame
    the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
    The name of the object to which the sensor is attached.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   The type of object to which the frame-of-reference is attached. The semantics are identical to the :at:`objtype`
+   attribute. If :at:`reftype` and :at:`refname` are given, the sensor values will be measured with respect to this
+   frame. If they are not given, sensor values will be measured with respect to the global frame.
+:at:`refname`: :at-val:`string`
+   The name of the object to which the frame-of-reference is attached.
 
 .. _sensor-framequat:
 
@@ -4746,11 +4768,13 @@ object, in global coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-framexaxis:
 
@@ -4763,11 +4787,13 @@ object, in global coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-frameyaxis:
 
@@ -4780,11 +4806,13 @@ object, in global coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-framezaxis:
 
@@ -4797,11 +4825,13 @@ object, in global coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-framelinvel:
 
@@ -4814,11 +4844,13 @@ coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-frameangvel:
 
@@ -4831,11 +4863,13 @@ coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`reftype`: :at-val:`[body, xbody, geom, site, camera]`
+   See :ref:`framepos<sensor-framepos>` sensor.
+:at:`refname`: :at-val:`string`
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-framelinacc:
 
@@ -4848,11 +4882,9 @@ coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-frameangacc:
 
@@ -4865,11 +4897,9 @@ coordinates.
 :at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
    See :ref:`CSensor`.
 :at:`objtype`: :at-val:`[body, xbody, geom, site, camera], required`
-   The type of object to which the sensor is attached. This must be an object type that has a spatial frame. "body"
+   See :ref:`framepos<sensor-framepos>` sensor.
-   refers to the inertial frame of the body, while "xbody" refers to the regular frame of the body (usually centered at
-   the joint with the parent body).
 :at:`objname`: :at-val:`string, required`
-   The name of the object to which the sensor is attached.
+   See :ref:`framepos<sensor-framepos>` sensor.
 
 .. _sensor-subtreecom:
 

doc/changelog.rst

@@ -2,6 +2,83 @@
 Changelog
 =========
 
+Version 2.1.2 (Date TBD, 2022)
+------------------------------
+
+New modules
+^^^^^^^^^^^
+
+1. Added new :doc:`Python bindings<python>`, which can be installed via ``pip install mujoco``,
+   and imported as ``import mujoco``.
+#. Added new :doc:`Unity plug-in<unity>`.
+#. Added a new ``introspect`` module, which provides reflection-like capability for MuJoCo's public API, currently
+   describing functions and enums. While implemented in Python, this module is expected to be generally useful for
+   automatic code generation targeting multiple languages. (This is not shipped as part of the ``mujoco`` Python
+   bindings package.)
+
+API changes
+^^^^^^^^^^^
+
+4. Moved definition of ``mjtNum`` floating point type into a new header
+   `mjtnum.h <https://github.com/deepmind/mujoco/blob/main/include/mjtnum.h>`_.
+#. Renamed header `mujoco_export.h` to :ref:`mjexport.h<inHeader>`.
+#. Added ``mj_printFormattedData``, which accepts a format string for floating point numbers, for example to increase
+   precision.
+
+General
+^^^^^^^
+
+7. MuJoCo can load `OBJ <https://en.wikipedia.org/wiki/Wavefront_.obj_file>`_ mesh files.
+
+   a. Meshes containing polygons with more than 4 vertices are not supported.
+   #. In OBJ files containing multiple object groups, any groups after the first one will be ignored.
+
+#. Added optional frame-of-reference specification to :ref:`framepos<sensor-framepos>`,
+   :ref:`framequat<sensor-framequat>`, :ref:`framexaxis<sensor-framexaxis>`, :ref:`frameyaxis<sensor-frameyaxis>`,
+   :ref:`framezaxis<sensor-framezaxis>`, :ref:`framelinvel<sensor-framelinvel>`, and
+   :ref:`frameangvel<sensor-frameangvel>` sensors. The frame-of-reference is specified by new :at:`reftype` and
+   :at:`refname` attributes.
+
+#. Sizes of :ref:`user parameters <CUser>` are now automatically inferred.
+
+   a. Declarations of user parameters in the top-level :ref:`size <size>` clause (e.g. :at:`nuser_body`,
+      :at:`nuser_jnt`, etc.) now accept a value of -1, which is the default. This will automatically set the value to
+      the length of the maximum associated :at:`user` attribute defined in the model.
+   #. Setting a value smaller than -1 will lead to a compiler error (previously a segfault).
+   #. Setting a value to a length smaller than some :at:`user` attribute defined in the model will lead to an error
+      (previously additional values were ignored).
+
+#. Increased the maximum number of lights in an :ref:`mjvScene` from 8 to 100.
+
+#. Saved XML files only contain explicit :ref:`inertial <inertial>` elements if the original XML included them. Inertias
+   that were automatically inferred by the compiler's :ref:`inertiafromgeom <compiler>` mechanism remain unspecified.
+
+#. User-selected geoms are always rendered as opaque. This is useful in interactive visualizers.
+
+#. Static geoms now respect their :ref:`geom group<geom>` for visualisation. Until this change rendering of static geoms
+   could only be toggled using the :ref:`mjVIS_STATIC<mjtVisFlag>` visualisation flag . After this change, both the geom
+   group and the visualisation flag need to be enabled for the geom to be rendered.
+
+#. Pointer parameters in function declarations in :ref:`mujoco.h<inHeader>` that are supposed to represent fixed-length
+   arrays are now spelled as arrays with extents, e.g. ``mjtNum quat[4]`` rather than ``mjtNum* quat``. From the
+   perspective of C and C++, this is a non-change since array types in function signatures decay to pointer types.
+   However, it allows autogenerated code to be aware of expected input shapes.
+
+Bug Fixes
+^^^^^^^^^
+
+15. ``mj_loadXML`` and ``mj_saveLastXML`` are now locale-independent. The Unity plugin should now work correctly for
+    users whose system locales use commas as decimal separators.
+#.  XML assets in VFS no longer need to end in a null character. Instead, the file size is determined by the size
+    parameter of the corresponding VFS entry.
+#.  Fix a vertex buffer object memory leak in ``mjrContext`` when skins are used.
+#.  Camera quaternions are now normalized during XML compilation.
+
+Binary build
+^^^^^^^^^^^^
+
+18. Windows binaries are now built with Clang.
+
 Version 2.1.1 (Dec. 16, 2021)
 -----------------------------
 
@@ -45,7 +122,7 @@ Bug Fixes
    :ref:`weld <equality-weld>` constraints.
 
    .. note::
-      Forces generated by :ref:`spatial tendons <spatial>` which are outside the kinematic tree (i.e. between bodies
+      Forces generated by :ref:`spatial tendons <spatial>` which are outside the kinematic tree (i.e., between bodies
       which have no ancestral relationship) are still not taken into account by force and torque sensors. This remains a
       future work item.
 
@@ -118,7 +195,7 @@ New features
 General
 ^^^^^^^
 
-3. The pre-allocated sizes in the virtual file system (VFS) increased to 2000 and 1000, to allow for larger projects.
+3. The preallocated sizes in the virtual file system (VFS) increased to 2000 and 1000, to allow for larger projects.
 #. The C structs in the ``mjuiItem`` union are now named, for compatibility.
 #. Fixed: ``mjcb_contactfilter`` type is ``mjfConFilt`` (was ``mjfGeneric``).
 #. Fixed: The array of sensors in ``mjCModel`` was not cleared.

doc/computation.rst

@@ -189,7 +189,7 @@ change at runtime. In that case there is still a fixed enumeration order (corres
 elements appear in ``mjModel``) but any inactive constraints are omitted.
 
 The number of position coordinates :math:`n_Q` is larger than the number of degrees of freedom :math:`n_V` whenever
-quaternions are used to represent 3D orientations. This occurs when the model contains ball joints or free joints (i.e.
+quaternions are used to represent 3D orientations. This occurs when the model contains ball joints or free joints (i.e.,
 in most models). In that case :math:`\dot{q}` does not equal :math:`v`, at least not in the usual sense. Instead one has
 to consider the group of rigid body orientations :math:`SO(3)` - which has the geometry of a unit sphere in 4D space.
 Velocities live in the 3D tangent space to this sphere. This is taken into account by all internal computations. For
@@ -326,7 +326,7 @@ Passive forces
 Passive forces are defined as forces that depend only on position and velocity, and not on control in forward dynamics
 or acceleration in inverse dynamics. As a result, such forces are inputs to both the forward and inverse dynamics
 computations, and are identical in both cases. They are stored in ``mjData.qfrc_passive``. The passive forces computed
-by MuJoCo are also passive in the sense of physics, i.e. they do not increase energy, however the user can install the
+by MuJoCo are also passive in the sense of physics, i.e., they do not increase energy, however the user can install the
 callback :ref:`mjcb_passive` and add forces to ``mjData.qfrc_passive`` that may increase energy. This will not interfere
 with MuJoCo's operation as long as such user forces depend only on position and velocity.
 
@@ -460,7 +460,7 @@ constraint contributes :math:`\dim(r)` elements to the total constraint count :m
 properties of quaternions, differentiation with respect to :math:`q` produces vectors of size :math:`n_V` rather than
 :math:`n_Q`.
 
-Among other applications, equality constraints can be used to create "loop joints", i.e. joints that cannot be modeled
+Among other applications, equality constraints can be used to create "loop joints", i.e., joints that cannot be modeled
 via the kinematic tree. Gaming engines represent all joints in this way. The same can be done in MuJoCo but is not
 recommended - because it leads to both slower and less accurate simulation, effectively turning MuJoCo into a gaming
 engine. The only reason to represent joints with equality constraints would be to model soft joints - which can be done
@@ -562,7 +562,7 @@ with a 1 at the joint address. For tendons this is known as the moment arm vecto
    spatial tendons this could be used to model friction between the tendon and the surfaces it wraps around. Such
    friction will be load-independent though. To construct a more detailed model of this phenomenon, create several small
    floating spheres and connect them with tendons in series. Then the contacts between the spheres and the surrounding
-   surfaces will have load-dependent (i.e. Coulomb) friction, but this is less efficient to simulate.
+   surfaces will have load-dependent (i.e., Coulomb) friction, but this is less efficient to simulate.
 
 .. _coLimit:
 
@@ -588,7 +588,7 @@ because solving for equality constraint forces is generally faster.
 ``joint`` : 1 or 2
    Limits can be defined for scalar joints (hinge and slide) as well as for ball joints. Scalar joints are treated as
    described above. Ball joint limits are applied to the exponential-map or angle-axis representation of the joint
-   quaternion, i.e. the vector :math:`(\theta x, \theta y, \theta z)` where :math:`\theta` is the rotation angle and
+   quaternion, i.e., the vector :math:`(\theta x, \theta y, \theta z)` where :math:`\theta` is the rotation angle and
    :math:`(x, y, z)` is the unit vector corresponding to the rotation axis. The limit is applied to the absolute value
    of the rotation angle :math:`\theta`. At runtime the limit is determined by the larger of the two range parameters.
    For semantic clarity however, one should use the second range parameter to specify the limit and set the first range
@@ -1022,7 +1022,7 @@ compute it.
 
 Note that the quadratic term in the inverse problem is weighted by :math:`R` instead of :math:`A+R`. This tells us two
 things. First, in the limit :math:`R \to 0` corresponding to hard constraints the inverse is no longer defined, as one
-would expect. Second and more useful, the inverse problem is diagonal, i.e. it decouples into independent optimization
+would expect. Second and more useful, the inverse problem is diagonal, i.e., it decouples into independent optimization
 problems over the individual constraint forces. The only remaining coupling is due to the constraint set :math:`\Omega`,
 but that set is also decoupled over the conceptual constraints discussed earlier. It turns out that all these
 independent optimization problems can be solved analytically. The only non-trivial case is the elliptic friction cone
@@ -1110,7 +1110,7 @@ Substituting :math:`f^+` in the constraint dynamics :eq:`eq:identity` and rearra
 
 Thus the constrained acceleration interpolates between the unconstrained and the reference acceleration. In particular,
 in the limit :math:`R \to 0` we have a hard constraint and :math:`a^1 = a^*`, while in the limit :math:`R \to \infty` we
-have have an infinitely soft constraint (i.e. no constraint) and :math:`a^1 = a^0`. It is then natural to introduce a
+have have an infinitely soft constraint (i.e., no constraint) and :math:`a^1 = a^0`. It is then natural to introduce a
 model parameter which directly controls the interpolation. We call this parameter *impedance* and denote it :math:`d`.
 It is a vector with dimensionality :math:`n_C` satisfying :math:`0<d<1` element-wise. Once it is specified, we compute
 the diagonal elements of the regularizer as

doc/conf.py

@@ -1,4 +1,4 @@
-## Copyright 2021 DeepMind Technologies Limited
+# 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.
@@ -8,13 +8,11 @@
 #
 # 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.
+# 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.
 """Configuration file for the Sphinx documentation builder."""
 
-import doctest
-import inspect
 import os
 import sys
 

doc/index.rst

@@ -13,5 +13,6 @@
    XMLreference
    programming
    APIreference
+   python
    unity
    changelog

doc/modeling.rst

@@ -45,7 +45,7 @@ XML file has a unique top-level element. This element must be :el:`mujoco` for M
 Compiling models
 ~~~~~~~~~~~~~~~~
 
-Once a high-level mjCModel is created - by loading an MJCF file or an URDF file, or programmatically when such
+Once a high-level mjCModel is created - by loading an MJCF file or a URDF file, or programmatically when such
 functionality becomes available - it is compiled into mjModel. Even though loading and compilation are presently
 combined in one step, compilation is independent of loading, meaning that the compiler works in the same way
 regardless of how mjCModel was created. Both the parser and the compiler perform extensive error checking, and abort
@@ -159,7 +159,7 @@ cylinder  1 0 0 1
 
 
 The box uses the top-level defaults class "main" to set the values of its undefined attributes, because no other class
-was specified. The body specifies childclass "sub", causing all children of this body (and all their children etc) to
+was specified. The body specifies childclass "sub", causing all children of this body (and all their children etc.) to
 use class "sub" unless specified otherwise. So the ellipsoid uses class "sub". The sphere has explicitly defined rgba
 which overrides the default settings. The cylinder specifies defaults class "main", and so it uses "main" instead of
 "sub", even though the latter was specified in the childclass attribute of the body containing the geom.
@@ -187,7 +187,7 @@ geoms attached to the body. The undefined state cannot be entered in the XML fil
 defined in a given class, it cannot be undefined in that class or in any of its child classes. So if the goal is to
 leave a certain attribute undefined in a given model element, it must be undefined in the active defaults class.
 
-A final twist here are actuators. They are different because some of the actuator-related elements are actually
+A final twist here is actuators. They are different because some of the actuator-related elements are actually
 shortcuts, and shortcuts interact with the defaults setting mechanism in a non-obvious way. This is explained in the
 :ref:`Actuator shortcuts <CActuator>` section below.
 
@@ -234,7 +234,7 @@ relative to the body.
 
 In principle the user always has a choice between local and global coordinates, but in practice this choice is viable
 only when using geometric primitives rather than meshes. For meshes, the 3D vertex positions are expressed in either
-local and global coordinates depending on how the mesh was designed - effectively forcing the user to adopt the same
+local or global coordinates depending on how the mesh was designed - effectively forcing the user to adopt the same
 convention for the entire model. The alternative would be to pre-process the mesh data outside MuJoCo so as to change
 coordinates, but that effort is rarely justified.
 
@@ -308,7 +308,7 @@ First we explain the setting of the impedance d. Recall that d must lie between
 to the range [:ref:`mjMINIMP mjMAXIMP <glNumeric>`] which is currently set to [0.0001 0.9999]. It
 causes the solver to interpolate between the unforced acceleration a0 and reference acceleration aref. Small values of
 d correspond to soft/weak constraints while large values of d correspond to strong/hard constraints. The user can set
-d to a constant, or take advantage of its interpolating property and make it position-dependent, i.e. a function of r.
+d to a constant, or take advantage of its interpolating property and make it position-dependent, i.e., a function of r.
 Position-dependent impedance can be used to model soft contact layers around objects, or define equality constraints
 that become stronger with larger violation (so as to approximate backlash for example). The shape of the function d(r)
 is determined by the element-specific parameter vector :at:`solimp`.
@@ -408,11 +408,11 @@ margin, gap
    margin and gap are distance properties and a one-sided specification makes little sense.
 solref, solimp
    If one of the two geoms has higher priority, its solref and solimp parameters are used. If both geoms have the same
-   priority, the weighted average is used. The weights are proportional to the solmix attributes, i.e. weight1 = solmix1
+   priority, the weighted average is used. The weights are proportional to the solmix attributes, i.e., weight1 =
-   / (solmix1 + solmix2) and similarly for weight2. There is one important exception to this weighted averaging rule. If
+   solmix1 / (solmix1 + solmix2) and similarly for weight2. There is one important exception to this weighted averaging
-   solref for either geom is non-positive, i.e. it relies on the new direct format introduced in MuJoCo 2.0, then the
+   rule. If solref for either geom is non-positive, i.e., it relies on the new direct format introduced in MuJoCo 2.0,
-   element-wise minimum is used regardless of solmix. This is because averaging solref parameters in different formats
+   then the element-wise minimum is used regardless of solmix. This is because averaging solref parameters in different
-   would be meaningless.
+   formats would be meaningless.
 
 .. _COverride:
 
@@ -443,12 +443,14 @@ User parameters
 ~~~~~~~~~~~~~~~
 
 A number of MJCF elements have the optional attribute :at:`user`, which defines a custom element-specific parameter
-array. This interacts with the corresponding "nuser_XXX" attribute of the :ref:`size <size>` element.
+array. This interacts with the corresponding "nuser_XXX" attribute of the :ref:`size <size>` element. If for example we
-If for example we set :at:`nuser_geom` to 5, then every geom in mjModel will have a custom array of 5 real-valued
+set :at:`nuser_geom` to 5, then every geom in mjModel will have a custom array of 5 real-valued parameters. These geom-
-parameters. These geom-specific parameters are either defined in the MJCF file via the :at:`user` attribute of
+specific parameters are either defined in the MJCF file via the :at:`user` attribute of :ref:`geom <geom>`, or set to 0
-:ref:`geom <geom>`, or set to 0 by the compiler if this attribute is omitted. MuJoCo does not use these
+by the compiler if this attribute is omitted. The default value of all "nuser_XXX" attributes is -1, which instructs the
-parameters in any internal computations; instead they are available for custom computations. The parser allows arrays
+compiler to automatically set this value to the length of the maximum associated :at:`user` attribute defined in the
-of arbitrary length in the XML, and the compiler later resizes them to length nuser_XXX.
+model. MuJoCo does not use these parameters in any internal computations; instead they are available for custom
+computations. The parser allows arrays of arbitrary length in the XML, and the compiler later resizes them to length
+nuser_XXX.
 
 Some element-specific parameters that are normally used in internal computations can also be used in custom
 computations. This is done by installing user callbacks which override parts of the simulation pipeline. For example,
@@ -483,12 +485,12 @@ well as solver statistics per iteration. We can offer the following general guid
 -  The constraint Jacobian should be dense for small models and sparse for large models. The default setting is 'auto';
    it resolves to dense when the number of degrees of freedom is up to 60, and sparse over 60. Note however that the
    threshold is better defined in terms of number of active constraints, which is model and behavior dependent.
--  The choice between pyramidal and elliptic friction cones is a modeling choice rather than an algorithmic choice, i.e.
+-  The choice between pyramidal and elliptic friction cones is a modeling choice rather than an algorithmic choice,
-   it leads to a different optimization problem solved with the same algorithms. Elliptic cones correspond more closely
+   i.e., it leads to a different optimization problem solved with the same algorithms. Elliptic cones correspond more
-   to physical reality. However pyramidal cones can improve the performance of the algorithms - but not necessarily.
+   closely to physical reality. However pyramidal cones can improve the performance of the algorithms - but not
-   While the default is pyramidal, we recommend trying the elliptic cones. When contact slip is a problem, the best way
+   necessarily. While the default is pyramidal, we recommend trying the elliptic cones. When contact slip is a problem,
-   to suppress it is to use elliptic cones, large impratio, and the Newton algorithm with very small tolerance. If that
+   the best way to suppress it is to use elliptic cones, large impratio, and the Newton algorithm with very small
-   is not sufficient, enable the Noslip solver.
+   tolerance. If that is not sufficient, enable the Noslip solver.
 -  The Newton algorithm is the best choice for most models. It has quadratic convergence near the global minimum and
    gets there in surprisingly few iterations - usually around 5, and rarely more than 20. It should be used with
    aggressive tolerance values, say 1e-10, because it is capable of achieving high accuracy without added delay (due to
@@ -635,7 +637,7 @@ the shortcut :ref:`muscle <muscle>` is more convenient. As with all other actuat
 production mechanism and the transmission are defined independently. Nevertheless, muscles only make (bio)physical
 sense when attached to tendon or joint transmissions. For concreteness we will assume a tendon transmission here.
 
-First we discuss length and length scaling. The range of feasible lengths of the transmission (i.e. MuJoCo tendon)
+First we discuss length and length scaling. The range of feasible lengths of the transmission (i.e., MuJoCo tendon)
 will play an important role; see :ref:`Length range <CLengthRange>` section above. In biomechanics, a muscle and a
 tendon are attached in series and form a muscle-tendon actuator. Our convention is somewhat different: in MuJoCo the
 entity that has spatial properties (in particular length and velocity) is the tendon, while the muscle is an abstract
@@ -684,7 +686,7 @@ force. We multiply the scaled force by a muscle-specific constant F0 to obtain t
 actuator_force = - FLV(L, V, act) \* F0
 
 The negative sign is because positive muscle activation generates pulling force. The constant F0 is the peak active
-force at zero velocity. It is related to the muscle thickness (i.e. physiological cross-sectional area or PCSA). If
+force at zero velocity. It is related to the muscle thickness (i.e., physiological cross-sectional area or PCSA). If
 known, it can be set with the attribute force of element :ref:`muscle <muscle>`. If it is not known, we
 set it to -1 which is the default. In that case we rely on the fact that larger muscles tend to act on joints that
 move more weight. The attribute scale defines this relationship as:
@@ -787,7 +789,7 @@ instability. In practice tendons are quite stiff, and their effect can be captur
 curve corresponding to the inelastic case (Zajac 89). This can be done in MuJoCo by shortening the muscle operating
 range.
 
-Pennation angle (i.e. the angle between the muscle and the line of force) is not modeled in MuJoCo and is assumed to
+Pennation angle (i.e., the angle between the muscle and the line of force) is not modeled in MuJoCo and is assumed to
 be 0. This effect can be approximated by scaling down the muscle force and also adjusting the operating range.
 
 Tendon wrapping is also more limited in MuJoCo. We allow spheres and infinite cylinders as wrapping objects, and require
@@ -1064,14 +1066,14 @@ is a grouping element for XML purposes and would violate the MJCF format if incl
 
 This functionality enables modular MJCF models; see the MPL family of models in the model library. One example of
 modularity is constructing a model of a robot (which tends to be elaborate) and then including it in multiple
-"scenes", i.e. MJCF models defining the objects in the robot's environment. Another example is creating a file with
+"scenes", i.e., MJCF models defining the objects in the robot's environment. Another example is creating a file with
 commonly used assets (say materials with carefully adjusted rgba values) and including it in multiple models which
 reference those assets.
 
 The included files are not required to be valid MJCF files on their own, but they usually are. Indeed we have designed
 this mechanism to allow MJCF models to be included in other MJCF models. To make this possible, repeated MJCF sections
 are allowed even when that does not make sense semantically in the context of a single model. For example, we allow
-the kinematic tree to have multiple roots (i.e. multiple :el:`worldbody` elements) which are merged automatically by
+the kinematic tree to have multiple roots (i.e., multiple :el:`worldbody` elements) which are merged automatically by
 the parser. Otherwise including robots into scenes would be impossible.
 
 The flexibility of repeated MCJF sections comes at a price: global settings that apply to the entire model, such as
@@ -1260,7 +1262,7 @@ in a visible way, and the energy fluctuates around the initial value instead of
 Model sizes
 ~~~~~~~~~~~
 
-MuJoCo pre-allocates all the memory needed at runtime in mjData, and does not access the C/C++ memory manager after
+MuJoCo preallocates all the memory needed at runtime in mjData, and does not access the C/C++ memory manager after
 model creation. It is therefore essential to allocate enough memory. The allocation is controlled by three size
 parameters specified in the :ref:`size <size>` element, namely the stack size :at:`nstack`, the
 maximum number of contacts :at:`nconmax`, and the maximum number of scalar constraints :at:`njmax`. The default
@@ -1287,7 +1289,7 @@ model. If you only intend to use the CG solver, you can get away with significan
 Motion capture
 ~~~~~~~~~~~~~~
 
-Mocap bodies are static children of the world (i.e. have no joints) and their :at:`mocap` attribute is set to
+Mocap bodies are static children of the world (i.e., have no joints) and their :at:`mocap` attribute is set to
 "true". They can be used to input a data stream from a motion capture device into a MuJoCo simulation. Suppose you are
 holding a VR controller, or an object instrumented with motion capture markers (e.g. Vicon), and want to have a
 simulated object moving in the same way but also interacting with other simulated objects. There is a dilemma here:

doc/overview.rst

@@ -64,7 +64,7 @@ Tendon geometry
 General actuation model
    Designing a sufficiently rich actuation model while using a model-agnostic API is challenging. MuJoCo achieves this
    goal by adopting an abstract actuation model that can have different types of transmission, force generation, and
-   internal dynamics (i.e. state variables which make the overall dynamics 3rd order). These components can be
+   internal dynamics (i.e., state variables which make the overall dynamics 3rd order). These components can be
    instantiated so as to model motors, pneumatic and hydraulic cylinders, PD controllers, biological muscles and many
    other actuators in a unified way.
 
@@ -89,7 +89,7 @@ Separation of model and data
       this is done by the user.
    -  ``mjData`` contains all dynamic variables and intermediate results. It is used as a scratch pad where all
       functions read their inputs and write their outputs -- which then become the inputs to subsequent stages in the
-      simulation pipeline. It also contains a pre-allocated and internally managed stack, so that the runtime module
+      simulation pipeline. It also contains a preallocated and internally managed stack, so that the runtime module
       does not need to call memory allocation functions after the model is initialized.
 
    ``mjModel`` is constructed by the compiler. :ref:`mjData` is constructed at runtime, given
@@ -593,6 +593,36 @@ section is to preemptively clarify the aspects that are most likely to be confus
 and a tutorial on selected topics. We will need to refer to material covered later in the documentation, but
 nevertheless the text below is as self-contained and introductory as possible.
 
+.. _Units:
+
+Units are undefined
+~~~~~~~~~~~~~~~~~~~
+
+In MuJoCo basic physical units are undefined. The user may interpret the system of units as they choose, as long as it
+is consistent. To understand this, consider an example: the dynamics of a 1 Meter spaceship that weighs 1 Kg and has a 1
+Newton thruster are the same as those of a 1 cm spaceship that weighs 1 gram and has a 1 dyn thruster. This is because
+both `MKS <https://en.wikipedia.org/wiki/MKS_system_of_units>`__ and `CGS
+<https://en.wikipedia.org/wiki/Centimetre%E2%80%93gram%E2%80%93second_system_of_units>`__ are consistent systems of
+units. This property allows the user to scale their model as they choose, which is useful when simulating very small or
+very large things, to improve the numerical properties of the simulation.
+
+That said, users are encouraged to use MKS, as there are two places where MuJoCo uses MKS-like default values:
+
+- The default value of :ref:`gravity<option>` is (0, 0, -9.81), which corresponds to Earth surface gravity in MKS.
+  Note that this does not really define system of units to be MKS, since we might be using CGS on
+  `Enceladus <https://en.wikipedia.org/wiki/Enceladus>`__.
+- The default value of :ref:`geom density<geom>` (used to infer body masses and inertias) is 1000, which corresponds to
+  the density of water in MKS.
+
+Once a consistent system of basic units (length, mass, time) is chosen, all derived units correspond to this system, as
+in `Dimensional Analysis <https://en.wikipedia.org/wiki/Dimensional_analysis>`__. For example if our model is
+interpreted as MKS, then forces and torques are in Newton and Newton-Meter, respectively.
+
+**Angles:** Although angles can be specified using degrees in MJCF (and indeed degrees are the
+:ref:`default <compiler>`), internally all angles are `Radians <https://en.wikipedia.org/wiki/Radian>`__. So e.g., if we
+are using MKS, angular velocities reported by :ref:`gyroscopes<sensor-gyro>` would be in rad/s while stiffness of hinge
+joints would be in Nm/rad.
+
 .. _NotObject:
 
 Not object-oriented
@@ -724,14 +754,14 @@ and two geoms to one body in this case.
 .. code:: XML
 
    <mujoco>
-      <worldbody>
+     <worldbody>
-         <body pos="0 0 0">
+       <body pos="0 0 0">
-            <geom type="sphere" size=".1" rgba=".9 .9 .1 1"/>
+         <geom type="sphere" size=".1" rgba=".9 .9 .1 1"/>
-            <geom type="capsule" pos="0 0 .1" size=".05 .1" rgba=".9 .9 .1 1"/>
+         <geom type="capsule" pos="0 0 .1" size=".05 .1" rgba=".9 .9 .1 1"/>
-            <site type="box" pos="0 -.1 .3" size=".02 .02 .02" rgba=".9 .1 .9 1"/>
+         <site type="box" pos="0 -.1 .3" size=".02 .02 .02" rgba=".9 .1 .9 1"/>
-            <site type="ellipsoid" pos="0 .1 .3" size=".02 .03 .04" rgba=".9 .1 .9 1"/>
+         <site type="ellipsoid" pos="0 .1 .3" size=".02 .03 .04" rgba=".9 .1 .9 1"/>
-         </body>
+       </body>
-      </worldbody>
+     </worldbody>
    </mujoco>
 
 .. figure:: images/overview/bodygeomsite.png
@@ -784,7 +814,7 @@ When working in joint coordinates, you cannot simply set the position and orient
 you want. To achieve that effect you would have to implement some form of inverse kinematics, which computes a (not
 necessarily unique) set of joint coordinates for which the forward kinematics place the body where you want it to be.
 
-The situation is different for floating bodies, i.e. bodies that are connected to the world with a free joint. The
+The situation is different for floating bodies, i.e., bodies that are connected to the world with a free joint. The
 positions and orientations as well as the linear and angular velocities of such bodies are explicitly represented in
 ``mjData.qpos`` and ``mjData.qvel``, and can therefore be manipulated directly. The general approach is to find the
 addresses in qpos and qvel where the body's data are. Of course qpos and qvel represents joints and not bodies, so you
@@ -805,7 +835,7 @@ can be obtained as:
          qveladr = m->jnt_dofadr[m->body_jntadr[bodyid]];
       }
 
-Now if everything went well (i.e. "myfloatingbody" was indeed a floating body), qposadr and qveladr are the addresses in
+Now if everything went well (i.e., "myfloatingbody" was indeed a floating body), qposadr and qveladr are the addresses
-qpos and qvel where the data for our floating body/joint lives. The position data is 7 numbers (3D position followed by
+in qpos and qvel where the data for our floating body/joint lives. The position data is 7 numbers (3D position followed
-unit quaternion) while the velocity data is 6 numbers (3D linear velocity followed by 3D angular velocity). These
+by unit quaternion) while the velocity data is 6 numbers (3D linear velocity followed by 3D angular velocity). These
 numbers can now be set to the desired pose and velocity of the body.

doc/programming.rst

@@ -129,10 +129,15 @@ mjrender.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/m
    This file defines the primitive types and structures needed by the OpenGL renderer.
 mjui.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mjui.h>`__
    This file defines the primitive types and structures needed by the UI framework.
+mjtnum.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mjtnum.h>`__
+   Defines MuJoCo's ``mjtNum`` floating-point type to be either ``double`` or ``float``. See :ref:`mjtNum`.
 mjxmacro.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mjxmacro.h>`__
    This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
    automate the mapping of mjModel and mjData into scripting languages, as well as other operations that require
    accessing all fields of mjModel and mjData. See code sample :ref:`testxml.cc <saTestXML>`.
+mjexport.h   `(source) <https://github.com/deepmind/mujoco/blob/main/include/mjexport.h>`__
+   Macros used for exporting public symbols from the MuJoCo library. This header should not be used directly by client
+   code.
 glfw3.h
    This file is optional and is not included by mujoco.h. It is the only header file needed for the GLFW library. See
    code sample :ref:`simulate.cc <saSimulate>`.
@@ -379,9 +384,10 @@ is extracted from the diagnostic fields of mjData. It is a very useful tool for
 constraint solver algorithms. The outputs of the sensors defined in the model are visualized as a bar graph.
 
 Note that the profiler shows timing information collected with high-resolution timers. On Windows, depending on the
-power settings, the OS may reduce the CPU frequency; this is because simulate.cc sleeps most of the time in order to
+power settings, the OS may reduce the CPU frequency; this is because `simulate.cc
-slow down to realtime. This results in inaccurate timings. To avoid this problem, change the Windows power plan so
+<https://github.com/deepmind/mujoco/blob/main/sample/simulate.cc>`_ sleeps most of the time in order to slow down to
-that the minimum processor state is 100%.
+realtime. This results in inaccurate timings. To avoid this problem, change the Windows power plan so that the minimum
+processor state is 100%.
 
 .. _saRecord:
 
@@ -390,11 +396,12 @@ that the minimum processor state is 100%.
 
 This code sample simulates the passive dynamics of a given model, renders it offscreen, reads the color and depth pixel
 values, and saves them into a raw data file that can then be converted into a movie file with tools such as ffmpeg. The
-rendering is simplified compared to simulate.cc because there is no user interaction, visualization options or timing;
+rendering is simplified compared to `simulate.cc <https://github.com/deepmind/mujoco/blob/main/sample/simulate.cc>`_
-instead we simply render with the default settings as fast as possible. The dimensions and number of multi-samples for
+because there is no user interaction, visualization options or timing; instead we simply render with the default
-the offscreen buffer are specified in the MuJoCo model, while the simulation duration, frames-per-second to be rendered
+settings as fast as possible. The dimensions and number of multi-samples for the offscreen buffer are specified in the
-(usually much less than the physics simulation rate), and output file name are specified as command-line arguments. For
+MuJoCo model, while the simulation duration, frames-per-second to be rendered (usually much less than the physics
-example, a 5 second animation at 60 frames per second is created with:
+simulation rate), and output file name are specified as command-line arguments. For example, a 5 second animation at 60
+frames per second is created with:
 
 .. code-block:: Shell
 
@@ -456,12 +463,12 @@ more accurate. Here ``f`` is one of the functions
 
 The code sample computes six Jacobian matrices, containing the derivative of each function with respect to its three
 arguments. The results are stored in the array ``deriv``. All six Jacobian matrices are square, with dimensionality
-equal to the number of degrees of freedom mjModel.nv. When the model configuration includes quaternion joints,
+equal to the number of degrees of freedom ``mjModel.nv``. When the model configuration includes quaternion joints,
 mjData.qpos has larger dimensionality than the other vectors, however the derivative is only defined in the tangent
-space to the configuration manifold. This is why, when differentiating with respect to the elements of mjData.qpos, we
+space to the configuration manifold. This is why, when differentiating with respect to the elements of ``mjData.qpos``,
-do not directly add ``eps`` but instead use the function :ref:`mju_quatIntegrate`
+we do not directly add ``eps`` but instead use the function :ref:`mju_quatIntegrate` to perturb the quaternion in the
-to perturb the quaternion in the tangent space, keeping it normalized. This technique should also be used in any other
+tangent space, keeping it normalized. This technique should also be used in any other situation where quaternions need
-situation where quaternions need to be perturbed.
+to be perturbed.
 
 There are some important subtleties in this code that improve speed as well as accuracy. To speed up the computation,
 we re-use intermediate results whenever possible. This relies on the skip mechanism described under :ref:`forward
@@ -475,9 +482,9 @@ Accuracy depends on the value of ``eps`` which is user-adjustable, as well as th
 of forward dynamics however, the function evaluation involves an iterative constraint solver, and this must be handled
 with care. In general, the difference between ``f(x+eps)`` and ``f(x)`` is very small, thus any noise affecting the
 two function evaluations differently can make the resulting derivatives meaningless. Different warm-starts or
-different number of solver iterations can act as such noise here. Therefore we fix the warm-start mjData.qacc to a
+different number of solver iterations can act as such noise here. Therefore we fix the warm-start ``mjData.qacc`` to a
 value pre-computed at the center point, using ``nwarmup`` extra major iterations to obtain a more accurate warm-start.
-We also fix the number of solver iterations to ``niter`` and set mjModel.opt.tolerance = 0; this disables the early
+We also fix the number of solver iterations to ``niter`` and set ``mjModel.opt.tolerance = 0``; this disables the early
 termination mechanism. Note that the original simulation options are restored in the serial code which advances the
 state.
 
@@ -485,18 +492,19 @@ We emphasize that the above subtleties are not high-order corrections that can b
 of unilateral constraints, numerical derivatives are hard to compute and there is no shortcut around it; indeed they
 would not even be defined if it wasn't for our soft-constraint model. Making the constraints softer results in more
 accurate results. This effect is so strong that in some situations it makes sense to intentionally work with the wrong
-model, i.e. a model that is softer than desired, so as to obtain more accurate derivatives.
+model, i.e., a model that is softer than desired, so as to obtain more accurate derivatives.
 
 .. _saUItools:
 
 uitools
 ~~~~~~~
 
-`(uitools.h) <https://github.com/deepmind/mujoco/blob/main/sample/uitools.h>`_
+`(uitools.h) <https://github.com/deepmind/mujoco/blob/main/sample/uitools.h>`_ `(uitools.c)
-`(uitools.c) <https://github.com/deepmind/mujoco/blob/main/sample/uitools.c>`_
+<https://github.com/deepmind/mujoco/blob/main/sample/uitools.c>`_ This is not a stand-alone code sample, but rather a
-This is not a stand-alone code sample, but rather a small utility used to hook up the new UI to GLFW. It is used in
+small utility used to hook up the new UI to GLFW. It is used in `simulate.cc
-simulate.cc and can also be used in user projects that involve the new UI. If GLFW is replaced with a different window
+<https://github.com/deepmind/mujoco/blob/main/sample/simulate.cc>`_ and can also be used in user projects that involve
-library, this is the only file that would have to be changed in order to access the UI functionality.
+the new UI. If GLFW is replaced with a different window library, this is the only file that would have to be changed in
+order to access the UI functionality.
 
 .. _Simulation:
 
@@ -514,7 +522,7 @@ be discussed later.
 
 mjModel and mjData should never be allocated directly by the user. Instead they are allocated and initialized by the
 corresponding API functions. These are very elaborate data structures, containing (arrays of) other structures,
-pre-allocated data arrays for all intermediate results, as well as an :ref:`internal stack <siStack>`. Our strategy is
+preallocated data arrays for all intermediate results, as well as an :ref:`internal stack <siStack>`. Our strategy is
 to allocate all necessary heap memory at the beginning of the simulation, and free it after the simulation is done, so
 that we never have to call the C memory allocation and deallocation functions during the simulation. This is done for
 speed, avoidance of memory fragmentation, future GPU portability, and ease of managing the state of the entire
@@ -532,7 +540,7 @@ available options are
    mjModel* m = mj_loadXML("mymodel.xml", NULL, errstr, errstr_sz);
 
    // option 2: parse and compile XML from virtual file system
-   mjModel* m = mj_loadXML("mymodel.xml", vfs, errstr_sz);
+   mjModel* m = mj_loadXML("mymodel.xml", vfs, errstr, errstr_sz);
 
    // option 3: load precompiled model from MJB file
    mjModel* m = mj_loadModel("mymodel.mjb", NULL);
@@ -613,12 +621,12 @@ The default (and recommended) way to control the system is to implement a contro
            mju_scl(d->ctrl, d->qvel, -0.1, m->nv);
    }
 
-This illustrates two concepts. First, we are checking if the number of controls mjModel.nu equals the number of dofs
+This illustrates two concepts. First, we are checking if the number of controls ``mjModel.nu`` equals the number of
-mjModel.nv. In general, the same callback may be used with multiple models depending on how the user code is structured,
+DoFs ``mjModel.nv``. In general, the same callback may be used with multiple models depending on how the user code is
-and so it is a good idea to check the model dimensions in the callback. Second, MuJoCo has a library of BLAS-like
+structured, and so it is a good idea to check the model dimensions in the callback. Second, MuJoCo has a library of
-functions that are very useful; indeed a large part of the code base consists of calling such functions internally. The
+BLAS-like functions that are very useful; indeed a large part of the code base consists of calling such functions
-:ref:`mju_scl` function above scales the velocity vector mjData.qvel by a constant feedback
+internally. The :ref:`mju_scl` function above scales the velocity vector ``mjData.qvel`` by a constant feedback
-gain and copies the result into the control vector mjData.ctrl. To install this callback, we simply assign it to the
+gain and copies the result into the control vector ``mjData.ctrl``. To install this callback, we simply assign it to the
 global control callback pointer :ref:`mjcb_control`:
 
 .. code-block:: C
@@ -631,10 +639,10 @@ signal is needed by the simulation pipeline, and as a result we will end up simu
 damping does not really do justice to the notion of control, and is better implemented as a passive joint property,
 but these are finer points).
 
-Instead of relying on a control callback, we could set the control vector mjData.ctrl directly. Alternatively we could
+Instead of relying on a control callback, we could set the control vector ``mjData.ctrl`` directly. Alternatively we
-set applied forces as explained in :ref:`state and control <siStateControl>`. If we could compute these
+could set applied forces as explained in :ref:`state and control <siStateControl>`. If we could compute these control-
-control-related quantities before mj_step is called, then the simulation loop for the controlled dynamics (without
+related quantities before mj_step is called, then the simulation loop for the controlled dynamics (without using a
-using a control callback) would become
+control callback) would become
 
 .. code-block:: C
 
@@ -646,18 +654,18 @@ using a control callback) would become
 
 Why would we not be able to compute the controls before mj_step is called? After all, isn't this what causality means?
 The answer is subtle but important, and has to do with the fact that we are simulating in discrete time. The top-level
-simulation function mj_step basically does two things: compute the :ref:`forward dynamics <siForward>` in continuous
+simulation function ``mj_step`` basically does two things: compute the :ref:`forward dynamics <siForward>` in continuous
-time, and then integrate over a time period specified by mjModel.opt.timestep. Forward dynamics computes the
+time, and then integrate over a time period specified by ``mjModel.opt.timestep``. Forward dynamics computes the
-acceleration mjData.qacc at time mjData.time, given the :ref:`state and control <siStateControl>` at time mjData.time.
+acceleration ``mjData.qacc`` at time ``mjData.time``, given the :ref:`state and control <siStateControl>` at time
-The numerical integrator then advances the state and time to mjData.time + mjModel.opt.timestep. Now, the control is
+``mjData.time``. The numerical integrator then advances the state and time to ``mjData.time + mjModel.opt.timestep``.
-required to be a function of the state at time mjData.time. However a general feedback controller can be a very complex
+Now, the control is required to be a function of the state at time ``mjData.time``. However a general feedback
-function, depending on various features of the state - in particular all the features computed by MuJoCo as intermediate
+controller can be a very complex function, depending on various features of the state - in particular all the features
-results of the simulation. These may include contacts, Jacobians, passive forces. None of these quantities are available
+computed by MuJoCo as intermediate results of the simulation. These may include contacts, Jacobians, passive forces.
-before mj_step is called (or rather, they are available but outdated by one time step). In contrast, when mj_step calls
+None of these quantities are available before ``mj_step`` is called (or rather, they are available but outdated by one
-our control callback, it does so as late in the computation as possible - namely after all the intermediate results
+time step). In contrast, when ``mj_step`` calls our control callback, it does so as late in the computation as possible
-dependent on the state but not on the control have been computed.
+- namely after all the intermediate results dependent on the state but not on the control have been computed.
 
-The same effect can be achieved without using a control callback. This is done by breaking mj_step in two parts:
+The same effect can be achieved without using a control callback. This is done by breaking ``mj_step`` in two parts:
 before the control is needed, and after the control is needed. The simulation loop now becomes
 
 .. code-block:: C
@@ -672,7 +680,7 @@ before the control is needed, and after the control is needed. The simulation lo
 There is one complication however: this only works with Euler integration. The Runge-Kutta integrator (as well as other
 advanced integrators we plan to implement) need to evaluate the entire dynamics including the feedback control law
 multiple times per step, which can only be done using a control callback. But with Euler integration, the above
-separation of mj_step into :ref:`mj_step1` and :ref:`mj_step2` is sufficient to provide the control law with the
+separation of ``mj_step`` into :ref:`mj_step1` and :ref:`mj_step2` is sufficient to provide the control law with the
 intermediate results of the computation.
 
 To make the above discussion more clear, we provide the internal implementation of mj_step, mj_step1 and mj_step2,
@@ -705,7 +713,7 @@ The control callback (if any) is called from within the forward dynamics functio
 Next we show the implementation of the two-part stepping approach, although the specifics will make sense only after
 we explain the :ref:`forward dynamics <siForward>` later. Note that the control callback is now called directly, since
 we have essentially unpacked the forward dynamics function. Note also that we always call the Euler integrator in
-mj_step2 regardless of the setting of mjModel.opt.integrator.
+mj_step2 regardless of the setting of ``mjModel.opt.integrator``.
 
 .. code-block:: C
 
@@ -767,19 +775,19 @@ The state vector in MuJoCo is:
 
 For a second-order dynamical system the state contains only position and velocity, however MuJoCo can also model
 actuators (such as cylinders and biological muscles) that have their own activation states assembled in the vector
-mjData.act. While the physics model is time-invariant, user-defined control laws may be time-varying; in particular
+``mjData.act``. While the physics model is time-invariant, user-defined control laws may be time-varying; in particular
-control laws obtained from trajectory optimizers would normally be indexed by mjData.time.
+control laws obtained from trajectory optimizers would normally be indexed by ``mjData.time``.
 
 The reason for the "official" caveat above is because user callbacks may store additional state variables that change
-over time and affect the callback outputs; indeed the field mjData.userdata exists mostly for that purpose. Other
+over time and affect the callback outputs; indeed the field ``mjData.userdata`` exists mostly for that purpose. Other
-state-like quantities that are part of mjData and are treated as inputs by forward dynamics are mjData.mocap_pos and
+state-like quantities that are part of mjData and are treated as inputs by forward dynamics are ``mjData.mocap_pos`` and
-mjData.mocap_quat. These quantities are unusual in that they are meant to change at each time step (normally driven by
+mjData.mocap_quat. These quantities are unusual in that they are meant to change at each time step (normally driven by a
-a motion capture device), however this change is implemented by the user, while the simulator treats them as
+motion capture device), however this change is implemented by the user, while the simulator treats them as constants. In
-constants. In that sense they are no different from all the constants in mjModel, or the function callback pointers
+that sense they are no different from all the constants in mjModel, or the function callback pointers set by the user:
-set by the user: such constants affect the computation, but are not part of the state vector of a dynamical system.
+such constants affect the computation, but are not part of the state vector of a dynamical system.
 
 The warm-start mechanism in the constraint solver effectively introduces another state variable. This mechanism uses
-the output of forward dynamics from the previous time step, namely the acceleration vector mjData.qacc, to estimate
+the output of forward dynamics from the previous time step, namely the acceleration vector ``mjData.qacc``, to estimate
 the current constraint forces via inverse dynamics. This estimate then initializes the optimization algorithm in the
 solver. If this algorithm runs until convergence the warm-start will affect the speed of convergence but not the final
 solution (since the underlying optimization problem is convex and does not have local minima), but in practice the
@@ -791,8 +799,8 @@ Next we turn to the controls and applied forces. The control vector in MuJoCo is
 
      u = (mjData.ctrl, mjData.qfrc_applied, mjData.xfrc_applied)
 
-These quantities specify control signals (mjData.ctrl) for the actuators defined in the model, or directly apply
+These quantities specify control signals (``mjData.ctrl``) for the actuators defined in the model, or directly apply
-forces and torques specified in joint space (mjData.qfrc_applied) or in Cartesian space (mjData.xfrc_applied).
+forces and torques specified in joint space (``mjData.qfrc_applied``) or in Cartesian space (mjData.xfrc_applied).
 
 Finally, calling mj_forward which corresponds to the abstract dynamics function ``f(t,x,u)`` computes the
 time-derivative of the state vector. The corresponding fields of mjData are
@@ -801,8 +809,8 @@ time-derivative of the state vector. The corresponding fields of mjData are
 
      dx/dt = f(t,x,u) = (1, mjData.qvel, mjData.qacc, mjData.act_dot)
 
-In the presence of quaternions (i.e. when free or ball joints are used), the position vector mjData.qpos has higher
+In the presence of quaternions (i.e., when free or ball joints are used), the position vector ``mjData.qpos`` has higher
-dimensionality than the velocity vector mjData.qvel and so this is not a simple time-derivative in the sense of
+dimensionality than the velocity vector ``mjData.qvel`` and so this is not a simple time-derivative in the sense of
 scalars, but instead takes quaternion algebra into account.
 
 To illustrate how the simulation state can be manipulated, suppose we have two mjData pointers src and dst
@@ -851,9 +859,9 @@ If the user has installed a control callback :ref:`mjcb_control` different from
 pointer), the user callback would be expected to set some of the above fields to non-zero. Note that MuJoCo will not
 clear these controls/forces at the end of the time step. This is the responsibility of the user.
 
-Also relevant in this context is the function :ref:`mj_resetData`. It sets mjData.qpos equal to the model reference
+Also relevant in this context is the function :ref:`mj_resetData`. It sets ``mjData.qpos`` equal to the model reference
-configuration mjModel.qpos0; mjData.mocap_pos and mjData.mocap_quat equal to the corresponding fixed body poses from
+configuration ``mjModel.qpos0``, ``mjData.mocap_pos`` and ``mjData.mocap_quat`` equal to the corresponding fixed body
-mjModel; and all other state and control variables to 0.
+poses from mjModel; and all other state and control variables to 0.
 
 .. _siForward:
 
@@ -861,7 +869,7 @@ Forward dynamics
 ~~~~~~~~~~~~~~~~
 
 The goal of forward dynamics is to compute the time-derivative of the state, namely the acceleration vector
-mjData.qacc and the activation time-derivative mjData.act_dot. Along the way it computes everything else needed to
+mjData.qacc and the activation time-derivative ``mjData.act_dot``. Along the way it computes everything else needed to
 simulate the dynamics, including active contacts and other constraints, joint-space inertia and its LTDL
 decomposition, constraint forces, sensor data and so on. All these intermediate results are available in mjData and
 can be used in custom computations. As illustrated in the :ref:`simulation loop <siSimulation>` section above, the
@@ -951,31 +959,31 @@ inverse dynamics (with the acceleration computed at the previous time step) to w
 solver in forward dynamics.
 
 The inputs to inverse dynamics are the same as the state vector in forward dynamics as illustrated in :ref:`state and
-control <siStateControl>`, but without mjData.act and mjData.time. Assuming no callbacks that depend on user-defined
+control <siStateControl>`, but without ``mjData.act`` and ``mjData.time``. Assuming no callbacks that depend on user-
-state variables, the inputs to inverse dynamics are the following fields of mjData:
+defined state variables, the inputs to inverse dynamics are the following fields of mjData:
 
 ::
 
      (mjData.qpos, mjData.qvel, mjData.qacc, mjData.mocap_pos, mjData.mocap_quat)
 
-The main output is mjData.qfrc_inverse. This is the force that must have acted on the system in order to achieve the
+The main output is ``mjData.qfrc_inverse``. This is the force that must have acted on the system in order to achieve the
-observed acceleration mjData.qacc. If forward dynamics were to be computed exactly, by running the iterative solver to
+observed acceleration ``mjData.qacc``. If forward dynamics were to be computed exactly, by running the iterative solver
-full convergence, we would have
+to full convergence, we would have
 
 ::
 
      mjData.qfrc_inverse = mjData.qfrc_applied + Jacobian'*mjData.xfrc_applied + mjData.qfrc_actuator
 
-where mjData.qfrc_actuator is the joint-space force produced by the actuators and the Jacobian is the mapping from
+where ``mjData.qfrc_actuator`` is the joint-space force produced by the actuators and the Jacobian is the mapping from
-joint to Cartesian space. When the "fwdinv" flag in mjModel.opt.enableflags is set, the above identity is used to
+joint to Cartesian space. When the "fwdinv" flag in ``mjModel.opt.enableflags`` is set, the above identity is used to
-monitor the quality of the forward dynamics solution. In particular, the two components of mjData.solver_fwdinv are
+monitor the quality of the forward dynamics solution. In particular, the two components of ``mjData.solver_fwdinv`` are
 set to the L2 norm of the difference between the forward and inverse solutions, in terms of joint forces and
 constraint forces respectively.
 
-Similar to forward dynamics, mj_inverse internally calls :ref:`mj_inverseSkip` with skip arguments (mjSTAGE_NONE, 0).
+Similar to forward dynamics, ``mj_inverse`` internally calls :ref:`mj_inverseSkip` with skip arguments
-The skip mechanism is the same as in forward dynamics, and can be used to speed up structured sampling. The result
+``(mjSTAGE_NONE, 0)``. The skip mechanism is the same as in forward dynamics, and can be used to speed up structured
-mjData.qfrc_inverse is obtained by using the Recursive Newton-Euler algorithm to compute the net force acting on the
+sampling. The result ``mjData.qfrc_inverse`` is obtained by using the Recursive Newton-Euler algorithm to compute the
-system, and then subtracting from it all internal forces.
+net force acting on the system, and then subtracting from it all internal forces.
 
 Inverse dynamics can be used as an analytical tool when experimental data are available. This is common in robotics as
 well as biomechanics. It can also be used to compute the joint torques needed to drive the system along a given
@@ -1073,7 +1081,7 @@ Model changes
 
 The MuJoCo model contained in mjModel is supposed to represent constant physical properties of the system, and in
 theory should not change after compilation. Of course in practice things are not that simple. It is often desirable to
-change the physics options in mjModel.opt, so as to experiment with different aspects of the physics or to create
+change the physics options in ``mjModel.opt``, so as to experiment with different aspects of the physics or to create
 custom computations. Indeed these options are designed in such a way that the user can make arbitrary changes to them
 between time steps.
 
@@ -1082,16 +1090,16 @@ because that may result in incorrect sizes or indexing. This rule does not hold
 parameters such as inertias are expected to obey certain properties. On the other hand, some structural parameters
 such as object types may be possible to change, but that depends on whether any sizes or indexes depend on them.
 Arrays of type mjtByte can be changed safely, since they are binary indicators that enable and disable certain
-features. The only exception here is mjModel.tex_rgb which is texture data represented as mjtByte.
+features. The only exception here is ``mjModel.tex_rgb`` which is texture data represented as mjtByte.
 
 When changing mjModel fields that corresponds to resources uploaded to the GPU, the user must also call the
-corresponding upload function: mjr_uploadTexture, mjr_uploadMesh, mjr_uploadHField. Otherwise the data used for
+corresponding upload function: ``mjr_uploadTexture``, ``mjr_uploadMesh``, ``mjr_uploadHField``. Otherwise the data used for
 simulation and for rendering will no longer be consistent.
 
 A related consideration has to do with changing real-valued fields of mjModel that have been used by the compiler to
 compute other real-valued fields: if we make a change, we want it to propagate. That is what the function
 :ref:`mj_setConst` does: it updates all derived fields of mjModel. These are fields whose names end with "0",
-corresponding to precomputed quantities when the model is in the reference configuration mjModel.qpos0.
+corresponding to precomputed quantities when the model is in the reference configuration ``mjModel.qpos0``.
 
 Finally, if changes are made to mjModel at runtime, it may be desirable to save them back to the XML. The function
 :ref:`mj_saveLastXML` does that in a limited sense: it copies all real-valued parameters from mjModel back to the
@@ -1123,17 +1131,17 @@ essential to keep it in mind at all times. All MuJoCo utility functions that ope
 difference between row-major and column-major formats.
 
 When possible, MuJoCo exploits sparsity. This can make all the difference between O(N) and O(N^3) scaling. The inertia
-matrix mjData.qM and its LTDL factorization mjData.qLD are always represented as sparse, using a custom indexing
+matrix ``mjData.qM`` and its LTDL factorization ``mjData.qLD`` are always represented as sparse, using a custom
-format designed for matrices that correspond to tree topology. The functions :ref:`mj_factorM`, :ref:`mj_solveM`,
+indexing format designed for matrices that correspond to tree topology. The functions :ref:`mj_factorM`,
-:ref:`mj_solveM2` and :ref:`mj_mulM` are used for sparse factorization, substitution and matrix-vector multiplication.
+:ref:`mj_solveM`, :ref:`mj_solveM2` and :ref:`mj_mulM` are used for sparse factorization, substitution and
-The user can also convert these matrices to dense format with the function :ref:`mj_fullM` although MuJoCo never does
+matrix-vector multiplication. The user can also convert these matrices to dense format with the function
-that internally.
+:ref:`mj_fullM` although MuJoCo never does that internally.
 
-The constraint Jacobian matrix mjData.efc_J is represented as sparse whenever the sparse Jacobian option is enabled.
+The constraint Jacobian matrix ``mjData.efc_J`` is represented as sparse whenever the sparse Jacobian option is
-The function :ref:`mj_isSparse` can be used to determine if sparse format is currently in use. In that case the
+enabled. The function :ref:`mj_isSparse` can be used to determine if sparse format is currently in use. In that case
-transposed Jacobian mjData.efc_JT is also computed, and the inverse constraint inertia mjData.efc_AR becomes sparse.
+the transposed Jacobian ``mjData.efc_JT`` is also computed, and the inverse constraint inertia ``mjData.efc_AR``
-Sparse matrices are stored in the compressed sparse row (CSR) format. For a generic matrix A with dimensionality
+becomes sparse. Sparse matrices are stored in the compressed sparse row (CSR) format. For a generic matrix A with
-m-by-n, this format is:
+dimensionality m-by-n, this format is:
 
 ======== ====== ============================================
 Variable Size   Meaning
@@ -1163,13 +1171,13 @@ translational component. We do not provide utility functions for working with th
 scope here. See Roy Featherstone's webpage on `Spatial Algebra <http://royfeatherstone.org/spatial/>`__. The unusual
 order (rotation before translation) is based on this material, and was apparently standard convention in the past.
 
-The data structures mjModel and mjData contain many pointers to pre-allocated buffers. The constructors of these data
+The data structures mjModel and mjData contain many pointers to preallocated buffers. The constructors of these data
-structures (mj_makeModel and mj_makeData) allocate one large buffer, namely mjModel.buffer and mjData.buffer, and then
+structures (mj_makeModel and mj_makeData) allocate one large buffer, namely ``mjModel.buffer`` and ``mjData.buffer``,
-partition it and set all the other pointers in it. mjData also contains a stack outside this main buffer, as discussed
+and then partition it and set all the other pointers in it. mjData also contains a stack outside this main buffer, as
-below. Even if two pointers appear one after the other, say mjData.qpos and mjData.qvel, do not assume that the data
+discussed below. Even if two pointers appear one after the other, say ``mjData.qpos`` and ``mjData.qvel``, do not
-arrays are contiguous and there is no gap between them. The constructors implement byte-alignment for each data array,
+assume that the data arrays are contiguous and there is no gap between them. The constructors implement byte-alignment
-and skip bytes when necessary. So if you want to copy mjData.qpos and mjData.qvel, the correct way to do it is the
+for each data array, and skip bytes when necessary. So if you want to copy ``mjData.qpos`` and ``mjData.qvel``, the
-hard way:
+correct way to do it is the hard way:
 
 .. code-block:: C
 
@@ -1180,10 +1188,10 @@ hard way:
    // DO NOT do this, there may be padding at the end of d->qpos
    mju_copy(myqposqvel, d->qpos, m->nq + m->nv);
 
-The :ref:`X Macros <tyXMacro>` defined in the optional header file mjxmacro.h can be used to automate allocation of data
+The :ref:`X Macros <tyXMacro>` defined in the optional header file ``mjxmacro.h`` can be used to automate allocation of
-structure that match mjModel and mjData, for example when writing a MuJoCo wrapper for a scripting language. In the code
+data structure that match mjModel and mjData, for example when writing a MuJoCo wrapper for a scripting language. In
-sample :ref:`testxml.cc <saTestXML>` we use these unusual macros to compare all data arrays from two instances of
+the code sample :ref:`testxml.cc <saTestXML>` we use these unusual macros to compare all data arrays from two instances
-mjModel and find the one with the largest difference. Apparently X Macros were invented in the 1960's for assembly
+of mjModel and find the one with the largest difference. Apparently X Macros were invented in the 1960's for assembly
 language, and remain a great idea.
 
 .. _siStack:
@@ -1191,23 +1199,23 @@ language, and remain a great idea.
 Internal stack
 ~~~~~~~~~~~~~~
 
-MuJoCo allocates and manages its own stack of mjtNums. mjData.stack is the pointer to the pre-allocated memory buffer.
+MuJoCo allocates and manages its own stack of mjtNums. ``mjData.stack`` is the pointer to the preallocated memory
-mjData.nstack is the maximum number of mjtNums that the stack can hold, as determined by the :at:`nstack` attribute of
+buffer. ``mjData.nstack`` is the maximum number of mjtNums that the stack can hold, as determined by the :at:`nstack`
-the :ref:`size <size>` element in MJCF. mjData.pstack is the first available address in the stack; this is our custom
+attribute of the :ref:`size <size>` element in MJCF. ``mjData.pstack`` is the first available address in the stack;
-stack pointer.
+this is our custom stack pointer.
 
 Most top-level MuJoCo functions allocate space on the stack, use it for internal computations, and then deallocate it.
 They cannot do this with the regular C stack because the allocation size is determined dynamically at runtime. And
 calling the heap memory management functions would be inefficient and result in fragmentation - thus a custom stack.
-When any MuJoCo function is called, upon return the value of mjData.pstack is the same. The only exception is the
+When any MuJoCo function is called, upon return the value of ``mjData.pstack`` is the same. The only exception is the
-function :ref:`mj_resetData` and its variants: they set mjData.pstack = 0. Note that this function is called
+function :ref:`mj_resetData` and its variants: they set ``mjData.pstack = 0``. Note that this function is called
-internally when an instability is detected in mj_step, mj_step1 and mj_step2. So if user functions take advantage of
+internally when an instability is detected in ``mj_step``, ``mj_step1`` and ``mj_step2``. So if user functions take
-the custom stack (as they should), this needs to be done in-between MuJoCo calls that have the potential to reset the
+advantage of the custom stack (as they should), this needs to be done in-between MuJoCo calls that have the potential
-simulation.
+to reset the simulation.
 
-Below is the general template for using the custom stack in user code. This assumes that mjData\* d is defined in the
+Below is the general template for using the custom stack in user code. This assumes that ``mjData\* d`` is defined in
-scope. If not, saving and restoring the stack pointer should be done manually instead of using the :ref:`mjMARKSTACK`
+the scope. If not, saving and restoring the stack pointer should be done manually instead of using the
-and :ref:`mjFREESTACK` macros.
+:ref:`mjMARKSTACK` and :ref:`mjFREESTACK` macros.
 
 .. code-block:: C
 
@@ -1289,20 +1297,20 @@ in the diagnostics section at the beginning of mjData.
 
 When the simulator encounters a situation that is not a terminal error but is nevertheless suspicious and likely to
 result in inaccurate numerical results, it triggers a warning. There are several possible warning types, indexed by
-the enum type :ref:`mjtWarning`. The array mjData.warning contains one :ref:`mjWarningStat` data structure per warning
+the enum type :ref:`mjtWarning`. The array ``mjData.warning`` contains one :ref:`mjWarningStat` data structure per
-type, indicating how many times each warning type has been triggered since the last reset and any information about
+warning type, indicating how many times each warning type has been triggered since the last reset and any information
-the warning (usually the index of the problematic model element). The counters are cleared upon reset. When a warning
+about the warning (usually the index of the problematic model element). The counters are cleared upon reset. When a
-of a given type is first triggered, the warning text is also printed by mju_warning as documented in :ref:`error and
+warning of a given type is first triggered, the warning text is also printed by mju_warning as documented in
-memory <siError>` above. All this is done by the function :ref:`mj_warning` which the simulator calls internally when
+:ref:`error and memory <siError>` above. All this is done by the function :ref:`mj_warning` which the simulator calls
-it encounters a warning. The user can also call this function directly to emulate a warning.
+internally when it encounters a warning. The user can also call this function directly to emulate a warning.
 
 When a model needs to be optimized for high-speed simulation, it is important to know where in the pipeline the CPU
 time is spent. This can in turn suggest which parts of the model to simplify or how to design the user application.
 MuJoCo provides an extensive profiling mechanism. It involves multiple timers indexed by the enum type
 :ref:`mjtTimer`. Each timer corresponds to a top-level API function, or to a component of such a function. Similar to
-warnings, timer information accumulates and is only cleared on reset. The array mjData.timer contains one
+warnings, timer information accumulates and is only cleared on reset. The array ``mjData.timer`` contains one
-:ref:`mjTimerStat` data structure per timer. The average duration per call for a given timer (corresponding to mj_step
+:ref:`mjTimerStat` data structure per timer. The average duration per call for a given timer (corresponding to
-in the example below) can be computed as:
+``mj_step`` in the example below) can be computed as:
 
 .. code-block:: C
 
@@ -1314,43 +1322,46 @@ to implement high-resolution timers in C without bringing in additional dependen
 does not need timing, and in that case there is no reason to call timing functions.
 
 One part of the simulation pipeline that needs to be monitored closely is the iterative constraint solver. The
-simplest diagnostic here is mjData.solver_iter which shows how many iterations the solver took on the last call to
+simplest diagnostic here is ``mjData.solver_iter`` which shows how many iterations the solver took on the last call to
-mj_step or mj_forward. Note that the solver has tolerance parameters for early termination, so this number is usually
+mj_step or ``mj_forward``. Note that the solver has tolerance parameters for early termination, so this number is
-smaller than the maximum number of iterations allowed. The array mjData.solver contains one :ref:`mjSolverStat` data
+usually smaller than the maximum number of iterations allowed. The array ``mjData.solver`` contains one
-structure per iteration of the constraint solver, with information about the constraint state and line search.
+:ref:`mjSolverStat` data structure per iteration of the constraint solver, with information about the constraint state
+and line search.
 
-When the option :at:`fwdinv` is enabled in mjModel.opt.enableflags, the field mjData.fwdinv is also populated. It
+When the option :at:`fwdinv` is enabled in ``mjModel.opt.enableflags``, the field ``mjData.fwdinv`` is also populated.
-contains the difference between the forward and inverse dynamics, in terms of generalized forces and constraint
+It contains the difference between the forward and inverse dynamics, in terms of generalized forces and constraint
 forces. Recall that that the inverse dynamics use analytical formulas and are always exact, thus any discrepancy is
-due to poor convergence of the iterative solver in the forward dynamics. The numbers in mjData.solver near termination
+due to poor convergence of the iterative solver in the forward dynamics. The numbers in ``mjData.solver`` near
-have similar order-of-magnitude as the numbers in mjData.fwdinv, but nevertheless these are two different diagnostics.
+termination have similar order-of-magnitude as the numbers in ``mjData.fwdinv``, but nevertheless these are two
+different diagnostics.
 
 Since MuJoCo's runtime works with compiled models, memory is preallocated when a model is compiled or loaded. Recall the
 :ref:`size <size>` element in MJCF, which has the attributes :at:`njmax`, :at:`nconmax` and :at:`nstack`. They determine
 the maximum number of scalar constraints that can be active simultaneously, the maximum number of contact points that
-can be included in mjData.contact, and the size of the internal stack. How is the user supposed to know what the
+can be included in ``mjData.contact``, and the size of the internal stack. How is the user supposed to know what the
 appropriate settings are? If there were a reliable recipe we would have implemented it in the compiler, but there isn't
 one. The theoretical worst-case, namely all geoms contacting all other geoms, calls for huge allocation which is almost
 never needed in practice. So our approach is to provide default settings in MJCF which are sufficient for most models,
 and allow the user to adjust them manually with the above attributes. If the simulator runs out of stack space at
 runtime it will trigger an error. If it runs out of space for contacts or scalar constraints, it will trigger a warning
 and omit the contacts and constraints that do not fit in the allocated buffers. When such errors or warnings are
-triggered, the user should adjust the sizes. The fields mjData.maxuse_stack, mjData.maxuse_con, mjData.maxuse_efc are
+triggered, the user should adjust the sizes. The fields ``mjData.maxuse_stack``, ``mjData.maxuse_con``,
-designed to help with this adjustment. They keep track of the maximum stack allocation, number of contacts and number of
+``mjData.maxuse_efc`` are designed to help with this adjustment. They keep track of the maximum stack allocation,
-scalar constraints respectively since the last reset. So one strategy is to make very large allocation, then monitor
+number of contacts and number of scalar constraints respectively since the last reset. So one strategy is to make very
-these maxuse_XXX statistics during typical simulations, and use them to reduce the allocation. Of course modern
+large allocation, then monitor these ``maxuse_XXX`` statistics during typical simulations, and use them to reduce the
-computers have so much memory that most users will not bother with such adjustment once they get rid of the
+allocation. Of course modern computers have so much memory that most users will not bother with such adjustment once
-out-of-memory errors and warnings, but nevertheless we provide this mechanism for the perfectionist.
+they get rid of the out-of-memory errors and warnings, but nevertheless we provide this mechanism for the
+perfectionist.
 
-The kinetic and potential energy are computed and stored in mjData.energy when the corresponding flag in
+The kinetic and potential energy are computed and stored in ``mjData.energy`` when the corresponding flag in
-mjModel.opt.enableflags is set. This can be used as another diagnostic. In general, simulation instability is
+``mjModel.opt.enableflags`` is set. This can be used as another diagnostic. In general, simulation instability is
 associated with increasing energy. In some special cases (when all unilateral constraints, actuators and dissipative
 forces are disabled) the underlying physical system is energy-conserving. In that case any temporal fluctuations in
 the total energy indicate inaccuracies in numerical integration. For such systems the Runge-Kutta integrator has much
 better performance than the default semi-implicit Euler integrator.
 
 Finally, the user can implement additional diagnostics as needed. Two examples were provided in the code samples
-testxml.cc and derivative.cc, where we computed model mismatches after save and load, and assessed the accuracy of the
+``testxml.cc`` and ``derivative.cc``, where we computed model mismatches after save and load, and assessed the accuracy of the
 numerical derivatives respectively. Key to such diagnostics is to implement two different algorithms or simulation
 paths that compute the same quantity, and compare the results numerically. This type of sanity check is essential when
 dealing with complex dynamical systems where we do not really know what the numerical output should be; if we knew
@@ -1365,30 +1376,30 @@ The derivative of any vector function with respect to its vector argument is cal
 in multi-joint kinematics and dynamics, it refers to the derivative of some spatial quantity as a function of the
 system configuration. In that case the Jacobian is also a linear map that operates on vectors in the (co)tangent space
 to the configuration manifold - such as velocities, momenta, accelerations, forces. One caveat here is that the system
-configuration encoded in mjData.qpos has dimensionality mjModel.nq, while the tangent space has dimensionality
+configuration encoded in ``mjData.qpos`` has dimensionality ``mjModel.nq``, while the tangent space has dimensionality
-mjModel.nv, and the latter is smaller when quaternion joints are present. So the size of the Jacobian matrix is
+``mjModel.nv``, and the latter is smaller when quaternion joints are present. So the size of the Jacobian matrix is
-N-by-mjModel.nv where N is the dimensionality of the spatial quantity being differentiated.
+N-by-``mjModel.nv`` where N is the dimensionality of the spatial quantity being differentiated.
 
 MuJoCo can differentiate analytically many spatial quantities. These include tendon lengths, actuator transmission
 lengths, end-effector poses, contact and other constraint violations. In the case of tendons and actuator
-transmissions the corresponding quantities are mjData.ten_moment and mjData.actuator_moment; we call them moment arms
+transmissions the corresponding quantities are ``mjData.ten_moment`` and ``mjData.actuator_moment``; we call them
-but mathematically they are Jacobians. The Jacobian matrix of all scalar constraint violations is stored in
+moment arms but mathematically they are Jacobians. The Jacobian matrix of all scalar constraint violations is stored in
-mjData.efc_J. Note that we are talking about constraint violations rather than the constraints themselves. This is
+``mjData.efc_J``. Note that we are talking about constraint violations rather than the constraints themselves. This is
-because constraint violations have units of length, i.e. they are spatial quantities that we can differentiate.
+because constraint violations have units of length, i.e., they are spatial quantities that we can differentiate.
 Constraints are more abstract entities and it is not clear what it means to differentiate them.
 
 Beyond these automatically-computed Jacobians, we provide support functions allowing the user to compute additional
 Jacobians on demand. The main function for doing this is :ref:`mj_jac`. It is given a 3D point and a MuJoCo body to
-which this point is considered to be attached. mj_jac then computes both the translational and rotational Jacobians,
+which this point is considered to be attached. ``mj_jac`` then computes both the translational and rotational
-which tell us how a spatial frame anchored at the given point will translate and rotate if we make a small change to
+Jacobians, which tell us how a spatial frame anchored at the given point will translate and rotate if we make a small
-the kinematic configuration. More precisely, the Jacobian maps joint velocities to end-effector velocities, while the
+change to the kinematic configuration. More precisely, the Jacobian maps joint velocities to end-effector velocities,
-transpose of the Jacobian maps end-effector forces to joint forces. There are also several other mj_jacXXX functions;
+while the transpose of the Jacobian maps end-effector forces to joint forces. There are also several other
-these are convenience functions that call the main mj_jac function with different points of interest - such as a body
+``mj_jacXXX`` functions; these are convenience functions that call the main ``mj_jac`` function with different points
-center of mass, geom center etc.
+of interest - such as a body center of mass, geom center etc.
 
 The ability to compute end-effector Jacobians exactly and efficiently is a key advantage of working in joint
 coordinates. Such Jacobians are the foundation of many control schemes that map end-effector errors to actuator
-commands suitable for suppressing those errors. The computation of end-effector Jacobians in MuJoCo via the mj_jac
+commands suitable for suppressing those errors. The computation of end-effector Jacobians in MuJoCo via the ``mj_jac``
 function is essentially free in terms of CPU cost; so do not hesitate to use this function.
 
 .. _siContact:
@@ -1399,7 +1410,7 @@ Contacts
 Collision detection and solving for contact forces were explained in detail in the :doc:`computation` chapter. Here we
 further clarify contact processing from a programming perspective.
 
-The collision detection stage finds contacts between geoms, and records them in the array mjData.contact of
+The collision detection stage finds contacts between geoms, and records them in the array ``mjData.contact`` of
 :ref:`mjContact` data structures. They are sorted such that multiple contacts between the same pair of bodies are
 contiguous (note that one body can have multiple geoms attached to it), and the body pairs themselves are sorted such
 that the first body acts as the major index and the second body as the minor index. Not all detected contacts are
@@ -1408,26 +1419,26 @@ mjContact.efc_address is the address in the list of active scalar constraints. R
 :at:`gap` attribute of :ref:`geom <geom>`, as well as certain kinds of internal processing that use virtual contacts
 for intermediate computations.
 
-The list mjData.contact is generated by the position stage of both forward and inverse dynamics. This is done
+The list ``mjData.contact`` is generated by the position stage of both forward and inverse dynamics. This is done
 automatically. However the user can override the internal collision detection functions, for example to implement
 non-convex mesh collisions, or to replace some of the convex collision functions we use with geom-specific primitives
 beyond the ones provided by MuJoCo. The global 2D array :ref:`mjCOLLISIONFUNC` contains the collision function pointer
 for each pair of geom types (in the upper-left triangle). To replace them, simply set these pointers to your
 functions. The collision function type is :ref:`mjfCollision`. When user collision functions detect contacts, they
 should construct an mjvContact structure for each contact and then call the function :ref:`mj_addContact` to add that
-contact to mjData.contact. The reference documentation of mj_addContact explains which fields of mjContact must be
+contact to ``mjData.contact``. The reference documentation of mj_addContact explains which fields of mjContact must be
 filled in by custom collision functions. Note that the functions we are talking about here correspond to near-phase
 collisions, and are called only after the list of candidate geom pairs has been constructed by the internal
 broad-phase collision mechanism.
 
-After the constraint forces have been computed, the vector of forces for contact i starts at:
+After the constraint forces have been computed, the vector of forces for contact ``i`` starts at:
 
 .. code-block:: C
 
    mjtNum* contactforce = d->efc_force + d->contact[i].efc_address;
 
-and similarly for all other efc_XXX vectors. Keep in mind that the contact friction cone can be pyramidal or elliptic,
+and similarly for all other ``efc_XXX`` vectors. Keep in mind that the contact friction cone can be pyramidal or
-depending on which solver is selected in mjModel.opt. The function :ref:`mj_isPyramidal`
+elliptic, depending on which solver is selected in ``mjModel.opt``. The function :ref:`mj_isPyramidal`
 can be used to determine which friction cone type is used. For pyramidal cones, the interpretation of the contact force
 (whose address we computed above) is non-trivial, because the components are forces along redundant non-orthogonal axes
 corresponding to the edges of the pyramid. The function :ref:`mj_contactForce` can be
@@ -1439,7 +1450,7 @@ columns. Here the axes are along the rows of the matrix. Thus, given that MuJoCo
 normal axis (which is the X axis of the contact frame by our convention) is in position mjContact.frame[0-2], the Y axis
 is in [3-5] and the Z axis is in [6-8]. The reason for this arrangement is because we can have frictionless contacts
 where only the normal axis is used, so it makes sense to have its coordinates in the first 3 positions of
-mjContact.frame.
+``mjContact.frame``.
 
 .. _siCoordinate:
 
@@ -1450,43 +1461,43 @@ There are multiple coordinate frames used in MuJoCo. The top-level distinction i
 Cartesian coordinates. The mapping from the vector of joints coordinates to the Cartesian positions and orientations
 of all bodies is called forward kinematics and is the first step in the physics pipeline. The opposite mapping is
 called inverse kinematics but it is not uniquely defined and is not implemented in MuJoCo. Recall that mappings
-between the tangent spaces (i.e. joint velocities and forces to Cartesian velocities and forces) are given by the body
+between the tangent spaces (i.e., joint velocities and forces to Cartesian velocities and forces) are given by the body
 Jacobians.
 
 Here we explain further subtleties and subdivisions of the coordinate frames, and summarize the available
-transformation functions. In joint coordinates, the only complication is that the position vector mjData.qpos has
+transformation functions. In joint coordinates, the only complication is that the position vector ``mjData.qpos`` has
-different dimensionality than the velocity and acceleration vectors mjData.qvel and mjData.qacc due to quaternion
+different dimensionality than the velocity and acceleration vectors ``mjData.qvel`` and ``mjData.qacc`` due to
-joints. The function :ref:`mj_differentiatePos` "subtracts" two joint position vectors and returns a velocity vector.
+quaternion joints. The function :ref:`mj_differentiatePos` "subtracts" two joint position vectors and returns a
-Conversely, the function :ref:`mj_integratePos` takes a position vector and a velocity vector, and returns a new
+velocity vector. Conversely, the function :ref:`mj_integratePos` takes a position vector and a velocity vector, and
-position vector which has been displaced by the given velocity.
+returns a new position vector which has been displaced by the given velocity.
 
 Cartesian coordinates are more complicated because there are three different coordinate frames that we use: local,
 global, and com-based. Local coordinates are used in mjModel to represent the static offsets between a parent and a
 child body, as well as the static offsets between a body and any geoms, sites, cameras and lights attached to it.
-These static offsets are applied in addition to any joint transformations. So mjModel.body_pos, mjModel.body_quat and
+These static offsets are applied in addition to any joint transformations. So ``mjModel.body_pos``,
-all other spatial quantities in mjModel are expressed in local coordinates. The job of forward kinematics is to
+``mjModel.body_quat`` and all other spatial quantities in mjModel are expressed in local coordinates. The job of
-accumulate the joint transformations and static offsets along the kinematic tree and compute all positions and
+forward kinematics is to accumulate the joint transformations and static offsets along the kinematic tree and compute
-orientations in global coordinates. The quantities in mjData that start with "x" are expressed in global coordinates.
+all positions and orientations in global coordinates. The quantities in mjData that start with "x" are expressed in
-These are mjData.xpos, mjData.geom_xpos etc. Frame orientations are usually stored as 3-by-3 matrices (xmat), except
+global coordinates. These are ``mjData.xpos``, ``mjData.geom_xpos`` etc. Frame orientations are usually stored as
-for bodies whose orientation is also stored as a unit quaternion mjData.xquat. Given this body quaternion, the
+3-by-3 matrices (xmat), except for bodies whose orientation is also stored as a unit quaternion ``mjData.xquat``. Given
-quaternions of all other objects attached to the body can be reconstructed by a quaternion multiplication. The
+this body quaternion, the quaternions of all other objects attached to the body can be reconstructed by a quaternion
-function :ref:`mj_local2Global` converts from local body coordinates to global
+multiplication. The function :ref:`mj_local2Global` converts from local body coordinates to global Cartesian
-Cartesian coordinates.
+coordinates.
 
 :ref:`mju_negPose` and :ref:`mju_trnVecPose`. A pose is a grouping of a 3D position and a unit quaternion orientation.
 There is no separate data structure; the grouping is in terms of logic. This represents a position and orientation in
 space, or in other words a spatial frame. Note that OpenGL uses 4-by-4 matrices to represent the same information,
 except here we use a quaternion for orientation. The function mju_mulPose multiplies two poses, meaning that it
-transforms the first pose by the second pose (the order is important). mju_negPose constructs the opposite pose, while
+transforms the first pose by the second pose (the order is important). ``mju_negPose`` constructs the opposite pose,
-mju_trnVecPose transforms a 3D vector by a pose, mapping it from local coordinates to global coordinates if we think
+while ``mju_trnVecPose`` transforms a 3D vector by a pose, mapping it from local coordinates to global coordinates if
-of the pose as a coordinate frame. If we want to manipulate only the orientation part, we can do that with the
+we think of the pose as a coordinate frame. If we want to manipulate only the orientation part, we can do that with the
 analogous quaternion utility functions :ref:`mju_mulQuat`, :ref:`mju_negQuat` and :ref:`mju_rotVecQuat`.
 
 Finally, there is the com-based frame. This is used to represent 6D spatial vectors containing a 3D angular velocity
 or acceleration or torque, followed by a 3D linear velocity or acceleration or force. Note the backwards order:
-rotation followed by translation. mjData.cdof and mjData.cacc are example of such vectors; the names start with "c".
+rotation followed by translation. ``mjData.cdof`` and ``mjData.cacc`` are example of such vectors; the names start with
-These vectors play a key role in the multi-joint dynamics computation. Explaining this is beyond our scope here; see
+"c". These vectors play a key role in the multi-joint dynamics computation. Explaining this is beyond our scope here;
-Featherstone's excellent `slides <http://royfeatherstone.org/spatial>`__ on the subject. In general, the user should
+see Featherstone's excellent `slides <http://royfeatherstone.org/spatial>`__ on the subject. In general, the user should
 avoid working with such quantities directly. Instead use the functions :ref:`mj_objectVelocity`,
 :ref:`mj_objectAcceleration` and the low-level :ref:`mju_transformSpatial` to obtain linear and angular velocities,
 accelerations and forces for a given body. Still, for the interested reader, we summarize the most unusual aspect of
@@ -1646,7 +1657,7 @@ mjCAMERA_FIXED
    targeting mode, it will move.
 mjCAMERA_USER
    This means that the abstract camera is ignored during an update and the low-level OpenGL cameras are not changed. It
-   is equivalent to not specifying an abstract camera at all, i.e. passing a NULL pointer to mjvCamera in the update
+   is equivalent to not specifying an abstract camera at all, i.e., passing a NULL pointer to mjvCamera in the update
    functions explained below.
 
 The low-level mjvGLCamera is what determines the actual rendering. There are two such cameras embedded in mjvScene, one
@@ -1655,9 +1666,9 @@ frame, while up corresponds to the positive Y axis. There is also a frustum in t
 average of the left and right frustum edges and then during rendering compute the actual edges from the viewport aspect
 ratio assuming 1:1 pixel aspect ratio. The distance between the two camera positions corresponds to the inter-pupilary
 distance (ipd). When the low-level camera parameters are computed automatically from an abstract camera, the ipd as well
-as vertical field of view (fovy) are taken from mjModel.vis.global.ipd/fovy for free and tracking cameras, and from the
+as vertical field of view (fovy) are taken from ``mjModel.vis.global.ipd``/``fovy`` for free and tracking cameras, and
-camera-specific mjModel.cam_ipd/fovy for cameras defined in the model. When stereoscopic mode is not enabled, as
+from the camera-specific ``mjModel.cam_ipd/fovy`` for cameras defined in the model. When stereoscopic mode is not
-determined by mjvScene.stereo, the camera data for the two eyes are internally averaged during rendering.
+enabled, as determined by mjvScene.stereo, the camera data for the two eyes are internally averaged during rendering.
 
 .. _viSelect:
 
@@ -1686,24 +1697,24 @@ Perturbations
 
 Interactive perturbations have proven very useful in exploring the model dynamics as well as probing closed-loop
 control systems. The user is free to implement any perturbation mechanism of their choice by setting
-mjData.qfrc_applied or mjData.xfrc_applied to suitable forces (in generalized and Cartesian coordinates respectively).
+``mjData.qfrc_applied`` or ``mjData.xfrc_applied`` to suitable forces (in generalized and Cartesian coordinates respectively).
 
 Prior to MuJoCo version 1.40, user code had to maintain a collection of objects in order to implement perturbations.
 All these objects are now grouped into the data structure :ref:`mjvPerturb`. Its use is illustrated in
 :ref:`simulate.cc <saSimulate>`.
-The idea is to select a MuJoCo body of interest, and provide a reference pose (i.e. a 3D position and quaternion
+The idea is to select a MuJoCo body of interest, and provide a reference pose (i.e., a 3D position and quaternion
 orientation) for that body. These are stored in mjPerturb.respos/quat. The function :ref:`mjv_movePerturb` is a mouse
 hook for controlling the reference pose with the mouse. The function :ref:`mjv_initPerturb` is used to set the
 reference pose equal to the selected body pose at the onset of perturbation, so as to avoid jumps.
 
-This perturbation object can then be used to move the selected body directly (when the simulation is paused or when
+This perturbation object can then be used to move the selected body directly (when the simulation is paused or when the
-the selected body is a mocap body), or to apply forces and torques to the body. This is done with the functions
+selected body is a mocap body), or to apply forces and torques to the body. This is done with the functions
 :ref:`mjv_applyPerturbPose` and :ref:`mjv_applyPerturbForce` respectively. The latter function writes the external
-perturbation force to mjData.xfrc_applied for the selected body. However it does not clear mjData.xfrc_applied for the
+perturbation force to ``mjData.xfrc_applied`` for the selected body. However it does not clear ``mjData.xfrc_applied``
-remaining bodies, thus it is recommended to clear it in user code, in case the selected body changed and some
+for the remaining bodies, thus it is recommended to clear it in user code, in case the selected body changed and some
 perturbation force was left over from a previous time step. If there is more than one device that can apply
-perturbations or user code needs to add perturbations from other sources, the user must implement the necessary logic
+perturbations or user code needs to add perturbations from other sources, the user must implement the necessary logic so
-so that only the desired perturbations are present in mjData.xfrc_applied and any old perturbations are cleared.
+that only the desired perturbations are present in ``mjData.xfrc_applied`` and any old perturbations are cleared.
 
 In addition to affecting the simulation, the perturbation object is recognized by the abstract visualizer and can be
 rendered. This is done by adding a visual string to denote the positional difference, and a rotating cube to denote
@@ -1728,13 +1739,13 @@ else needed for specify how rendering should be done. mjvScene also contains up
 copied from the model, as well as a headlight which is in light position 0 when present.
 
 The above procedure is the most common approach, and it updates the entire scene at each frame. In addition, we
-provide two functions for finer control. :ref:`mjv_updateCamera` updates only the camera (i.e. maps the abstract
+provide two functions for finer control. :ref:`mjv_updateCamera` updates only the camera (i.e., maps the abstract
 mjvCamera to the low-level mjvGLCamera) but does not touch the geoms or lights. This is useful when the user is moving
 the camera rapidly but the simulation state has not changed - in that case there is no point in re-creating the lists
 of geoms and lights.
 
 More advanced rendering effects can be achieved by manipulating the list of abstract geoms. For example, the user can
-add custom geoms at the end of the list. Sometimes it is desirable to render a sequence of simulation states (i.e. a
+add custom geoms at the end of the list. Sometimes it is desirable to render a sequence of simulation states (i.e., a
 trajectory) and not just the current state. For this purpose, we have provided the function :ref:`mjv_addGeoms` which
 adds the geoms corresponding to the current simulation state to the list already in mjvScene. It does not change the
 list of lights, because lighting is additive and having too many lights will make the scene too bright. Importantly,
@@ -1809,7 +1820,7 @@ yet the user is expected to move them programmatically at each simulation step.
 through contacts, or better yet, through soft equality constraints to regular bodies which in turn make contacts. The
 latter approach is illustrated in the MPL models available on the Forum. It provides effective dynamic filtering and
 avoids contacts involving bodies that behave as if they are infinitely heavy (which is what a fixed body is). Note
-that the time-varying positions and orientations of the mocap bodies are stored in mjData.mocap_pos/quat, as opposed
+that the time-varying positions and orientations of the mocap bodies are stored in ``mjData.mocap_pos/quat``, as opposed
 to storing them in mjModel. This is because mjModel is supposed to remain constant. The fixed mocap body pose stored
 in mjModel is only used at initialization and reset, when user code has not yet had a chance to update
 mjData.mocap_pos/quat.
@@ -1864,11 +1875,11 @@ mjrContext is specific to MuJoCo. After creation, it contains references (called
 resources that were uploaded to the GPU by mjr_makeContext. These include model-specific resources such as meshes and
 textures, as well as generic resources such as font bitmaps for the specified font scale, framebuffer objects for
 shadow mapping and offscreen rendering, and associated renderbuffers. It also contains OpenGL-related options copied
-from mjModel.vis, capabilities of the default window framebuffer that are discovered automatically, and the currently
+from ``mjModel.vis``, capabilities of the default window framebuffer that are discovered automatically, and the
-active buffer for rendering; see :ref:`buffers <reBuffer>` below. Note that even though MuJoCo uses fixed-function
+currently active buffer for rendering; see :ref:`buffers <reBuffer>` below. Note that even though MuJoCo uses
-OpenGL, it avoids immediate mode rendering and instead uploads all resources to the GPU upfront. This makes it as
+fixed-function OpenGL, it avoids immediate mode rendering and instead uploads all resources to the GPU upfront. This
-efficient as a modern shader, and possibly more efficient, because fixed-function OpenGL is now implemented via
+makes it as efficient as a modern shader, and possibly more efficient, because fixed-function OpenGL is now implemented
-internal shaders that have been written by the video driver developers and tuned extensively.
+via internal shaders that have been written by the video driver developers and tuned extensively.
 
 Most of the fields of mjrContext remain constant after the call to mjr_makeContext. The only exception is
 mjrContext.currentBuffer which changes whenever the active buffer changes. Some of the GPU resources may also change
@@ -1918,7 +1929,7 @@ code samples. All OpenGL can do is detect these properties; we do this in mjr_ma
 various window capabilities fields of mjrContext. This is why such properties are not part of the MuJoCo model; they are
 session/software-specific and not model-specific. In contrast, the offscreen framebuffer is managed entirely by OpenGL,
 and so we can create that buffer with whatever properties we want, namely with the resolution and multi-sample
-properties specified in mjModel.vis.
+properties specified in ``mjModel.vis``.
 
 The user can directly access the pixels in the two buffers. This is done with the functions :ref:`mjr_readPixels`,
 :ref:`mjr_drawPixels` and :ref:`mjr_blitBuffer`. Read/draw transfer pixels from/to the active buffer to/from the CPU.

doc/python.rst

@@ -0,0 +1,338 @@
+===============
+Python Bindings
+===============
+
+Introduction
+------------
+
+Starting with version 2.1.2, MuJoCo comes with native Python bindings that are developed in C++ using
+`pybind11 <https://pybind11.readthedocs.io/>`__. Unlike previous Python bindings, these are officially supported by the
+MuJoCo development team and will be kept up-to-date with the latest developments in MuJoCo itself.
+
+The Python bindings are distributed as the ``mujoco`` package on `PyPI <https://pypi.org/project/mujoco>`__. These are
+low-level bindings that are meant to give as close to a direct access to the MuJoCo library as possible. However, in
+order to provide an API and semantics that developers would expect in a typical Python library, the bindings
+deliberately diverge from the raw MuJoCo API in a number of places, which are documented throughout this page.
+
+DeepMind’s `dm_control <https://github.com/deepmind/dm_control>`__ reinforcement learning library (which prior to
+version 1.0.0 implemented its own MuJoCo bindings based on ``ctypes``) has been updated to depend on the ``mujoco``
+package and continues to be supported by DeepMind. Changes in dm_control should be largely transparent to users of
+previous versions, however code that depended directly on its low-level API may need to be updated. Consult the
+`migration guide <https://github.com/deepmind/dm_control/blob/main/migration_guide_1.0.md>`__ for detail.
+
+For mujoco-py users, we include :ref:`notes <PyMjpy_migration>` below to aid migration.
+
+Installation
+------------
+
+The package can be installed from `PyPI <https://pypi.org/project/mujoco/>`__ via
+
+.. code-block:: shell
+
+   pip install mujoco
+
+A copy of the MuJoCo library is provided as part of the package and does **not** need to be downloaded or installed
+separately.
+
+Building from source
+--------------------
+
+Source code for the Python bindings are available in the ``python`` top-level directory in MuJoCo's
+`GitHub repository <https://github.com/deepmind/mujoco>`__. Developers wishing to build the bindings from source should
+work with a full clone of the Git repository, run the ``make_sdist.sh`` script to generate a
+`source distribution (sdist) <https://packaging.python.org/en/latest/glossary/#term-Source-Distribution-or-sdist>`__
+tarball, then run ``pip wheel name_of_sdist.tar.gz`` to build the libraries and generate a
+`wheel <https://packaging.python.org/en/latest/glossary/#term-Built-Distribution>`__. The ``make_sdist.sh`` script
+generates additional C++ header files that are needed to build the bindings, and also pulls in other required files from
+elsewhere in the repository outside the ``python`` directory into the sdist.
+
+CMake and a C++17 compiler are needed to build the bindings from source.
+
+Basic usage
+-----------
+
+Once installed, the package can be imported via ``import mujoco``. Structs, functions, constants, and enums are
+available directly from the top-level ``mujoco`` module.
+
+.. _PyStructs:
+
+Structs
+=======
+
+  MuJoCo data structures are exposed as Python classes. In order to conform to
+  `PEP 8 <https://peps.python.org/pep-0008/>`__ naming guidelines, struct names begin with a capital letter, for example
+  ``mjData`` becomes ``mujoco.MjData`` in Python.
+
+  All structs other than ``mjModel`` have constructors in Python. For structs that have an ``mj_defaultFoo``-style
+  initialization function, the Python constructor calls the default initializer automatically, so for example
+  ``mujoco.MjOption()`` creates a new ``mjOption`` instance that is pre-initialized with :ref:`mj_defaultOption`.
+  Otherwise, the Python constructor zero-initializes the underlying C struct.
+
+  Structs with a ``mj_makeFoo``-style initialization function have corresponding constructor overloads in Python,
+  for example ``mujoco.MjvScene(model, maxgeom=10)`` in Python creates a new ``mjvScene`` instance that is
+  initialized with ``mjv_makeScene(model, [the new mjvScene instance], 10)`` in C. When this form of initialization is
+  used, the corresponding deallocation function ``mj_freeFoo/mj_deleteFoo`` is automatically called when the Python
+  object is deleted. The user does not need to manually free resources.
+
+  The ``mujoco.MjModel`` class does not a have Python constructor. Instead, we provide three static factory functions
+  that create a new ``mjModel`` instance: ``mujoco.MjModel.from_xml_string``, ``mujoco.MjModel.from_xml_path``, and
+  ``mujoco.MjModel.from_binary_path``. The first function accepts a model XML as a string, while the latter two
+  functions accept the path to either an XML or MJB model file. All three functions optionally accept a Python
+  dictionary which is converted into a MuJoCo :ref:`Virtualfilesystem` for use during model compilation.
+
+Functions
+=========
+
+  MuJoCo functions are exposed as Python functions of the same name. Unlike with structs, we do not attempt to make
+  the function names `PEP 8 <https://peps.python.org/pep-0008/>`__-compliant, as MuJoCo uses both underscores and
+  CamelCases. In most cases, function arguments appear exactly as they do in C, and keyword arguments are supported
+  with the same names as declared in :ref:`mujoco.h<inHeader>`. Python bindings to C functions that accept array input
+  arguments expect NumPy arrays or iterable objects that are convertible to NumPy arrays (e.g. lists). Output
+  arguments (i.e. array arguments that MuJoCo expect to write values back to the caller) must always be writeable
+  NumPy arrays.
+
+  In the C API, functions that take dynamically-sized arrays as inputs expect a pointer argument to the array along with
+  an integer argument that specifies the array's size. In Python, the size arguments are omitted since we can
+  automatically (and indeed, more safely) deduce it from the NumPy array. When calling these functions, pass all
+  arguments other than array sizes in the same order as they appear in :ref:`mujoco.h<inHeader>`, or use keyword
+  arguments. For example, :ref:`mj_jac` should be called as ``mujoco.mj_jac(m, d, jacp, jacr, point, body)`` in Python.
+
+  The bindings **releases the Python Global Interpreter Lock (GIL)** before calling the underlying MuJoCo function.
+  This allows for some thread-based parallelism, however users should bear in mind that the GIL is only released for the
+  duration of the MuJoCo C function itself, and not during the execution of any other Python code.
+
+Enums and constants
+===================
+
+  MuJoCo enums are available as ``mujoco.mjtEnumType.ENUM_VALUE``, for example ``mujoco.mjtObj.mjOBJ_SITE``. MuJoCo
+  constants are available with the same name directly under the ``mujoco`` module, for example ``mujoco.mjVISSTRING``.
+
+Minimal example
+---------------
+
+  .. code-block:: python
+
+     import mujoco
+
+     XML=r"""
+     <mujoco>
+       <asset>
+         <mesh file="gizmo.stl"/>
+       </asset>
+       <worldbody>
+         <body>
+           <freejoint/>
+           <geom type="mesh" name="gizmo" mesh="gizmo"/>
+         </body>
+       </worldbody>
+     </mujoco>
+     """
+
+     ASSETS=dict()
+     with open('/path/to/gizmo.stl', 'rb') as f:
+       ASSETS['gizmo.stl'] = f.read()
+
+     model = mujoco.MjModel.from_xml_string(XML, ASSETS)
+     data = mujoco.MjData(model)
+     while data.time < 1:
+       mujoco.mj_step(model, data)
+       print(data.geom_xpos)
+
+.. _PyNamed:
+
+Named access
+------------
+
+Most well-designed MuJoCo models assign names to objects (joints, geoms, bodies, etc.) of interest. When the model is
+compiled down to an ``mjModel`` instance, these names become associated with numeric IDs that are used to index into the
+various array members. For convenience and code readability, the Python bindings provide "named access" API on
+``MjModel`` and ``MjData``. Each ``name_fooadr`` field in the ``mjModel`` struct defines a name category ``foo``.
+
+For each name category ``foo``, ``mujoco.MjModel`` and ``mujoco.MjData`` objects provide a method ``foo`` that takes
+a single string argument, and returns an accessor object for all arrays corresponding to the entity ``foo`` of the
+given name. The accessor object contains attributes whose names correspond to the fields of either ``mujoco.MjModel`` or
+``mujoco.MjData`` but with the part before the underscore removed. For example:
+
+- ``m.geom('gizmo')`` returns an accessor for arrays in the ``MjModel`` object ``m`` associated with the geom named
+  "gizmo".
+- ``m.geom('gizmo').rgba`` is a NumPy array view of length 4 that specifies the RGBA color for the geom.
+  Specifically, it corresponds to the portion of ``m.geom_rgba[4*i:4*i+4]`` where
+  ``i = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'gizmo')``.
+
+Additionally, the Python API define a number of aliases for some name categories corresponding to the XML element name
+in the MJCF schema that defines an entity of that category. For example, ``m.joint('foo')`` is the same as
+``m.jnt('foo')``. A complete list of these aliases are provided below.
+
+The accessor for joints is somewhat different that of the other categories. Some ``mjModel`` and ``mjData`` fields
+(those of size size ``nq`` or ``nv``) are associated with degrees of freedom (DoFs) rather than joints. This is because
+different types of joints have different numbers of DoFs. We nevertheless associate these fields to their corresponding
+joints, for example through ``d.joint('foo').qpos`` and ``d.joint('foo').qvel``, however the size of these arrays would
+differ between accessors depending on the joint's type.
+
+Named access is guaranteed to be O(1) in the number of entities in the model. In other words, the time it takes to
+access an entity by name does not grow with the number of names or entities in the model. (This is currently **not** the
+case for the :ref:`mj_name2id` function, which performs a linear scan.)
+
+For completeness, we provide here a complete list of all name categories in MuJoCo, along with their corresponding
+aliases defined in the Python API.
+
+- ``body``
+- ``jnt`` or ``joint``
+- ``geom``
+- ``site``
+- ``cam`` or ``camera``
+- ``light``
+- ``mesh``
+- ``skin``
+- ``hfield``
+- ``tex`` or ``texture``
+- ``mat`` or ``material``
+- ``pair``
+- ``exclude``
+- ``eq`` or ``equality``
+- ``tendon`` or ``ten``
+- ``actuator``
+- ``sensor``
+- ``numeric``
+- ``text``
+- ``tuple``
+- ``key`` or ``keyframe``
+
+Rendering
+---------
+
+MuJoCo itself expects users to set up a working OpenGL context before calling any of its ``mjr_`` rendering routine.
+The Python bindings provide a basic class ``mujoco.GLContext`` that helps users set up such a context for offscreen
+rendering. To create a context, call ``ctx = mujoco.GLContext(max_width, max_height)``. Once the context is created,
+it must be made current before MuJoCo rendering functions can be called, which you can do so via ``ctx.make_current()``.
+Note that a context can only be made current on one thread at any given time, and all subsequent rendering calls must be
+made on the same thread.
+
+The context is freed automatically when the ``ctx`` object is deleted, but in some multi-threaded scenario it may be
+necessary to explicitly free the underlying OpenGL context. To do so, call ``ctx.free()``, after which point it is the
+user's responsibility to ensure that no further rendering calls are made on the context.
+
+Once the context is created, users can follow MuJoCo's standard rendering, for example as documented in the
+:ref:`Visualization` section.
+
+Error handling
+--------------
+
+MuJoCo reports irrecoverable errors via the :ref:`mju_error` mechanism, which immediately terminates the entire process.
+Users are permitted to install a custom error handler via the :ref:`mju_user_error` callback, but it too is expected
+to terminate the process, otherwise the behavior of MuJoCo after the callback returns is undefined. In actuality, it is
+sufficient to ensure that error callbacks do not return *to MuJoCo*, but it is permitted to use
+`longjmp <https://en.cppreference.com/w/c/program/longjmp>`__ to skip MuJoCo's call stack back to the external callsite.
+
+The Python bindings utilises longjmp to allow it to convert irrecoverable MuJoCo errors into Python exceptions of type
+``mujoco.FatalError`` that can be caught and processed in the usual Pythonic way. Furthermore, it installs its error
+callback in a thread-local manner using a currently private API, thus allowing for concurrent calls into MuJoCo from
+multiple threads.
+
+Callbacks
+---------
+
+MuJoCo allows users to install custom callback functions to modify certain parts of its computation pipeline.
+For example, :ref:`mjcb_sensor` can be used to implement custom sensors, and :ref:`mjcb_control` can be used to
+implement custom actuators. Callbacks are exposed through the function pointers prefixed ``mjcb_`` in
+:ref:`mujoco.h<inHeader>`.
+
+For each callback ``mjcb_foo``, users can set it to a Python callable via ``mujoco.set_mjcb_foo(some_callable)``. To
+reset it, call ``mujoco.set_mjcb_foo(None)``. To retrieve the currently installed callback, call
+``mujoco.get_mjcb_foo()``. (The getter **should not** be used if the callback is not installed via the Python bindings.)
+The bindings automatically acquire the GIL each time the callback is entered, and release it before reentering MuJoCo.
+This is likely to incur a severe performance impact as callbacks are triggered several times throughout MuJoCo's
+computation pipeline and is unlikely to be suitable for "production" use case. However, it is expected that this feature
+will be useful for prototyping complex models.
+
+Alternatively, if a callback is implemented in a native dynamic library, users can use
+`ctypes <https://docs.python.org/3/library/ctypes.html>`__ to obtain a Python handle to the C function pointer and pass
+it to ``mujoco.set_mjcb_foo``. The bindings will then retrieve the underlying function pointer and assign it directly to
+the raw callback pointer, and the GIL will **not** be acquired each time the callback is entered.
+
+.. _PyMjpy_migration:
+
+Migration Notes for mujoco-py
+-----------------------------
+
+In mujoco-py, the main entry point is the `MjSim <https://github.com/openai/mujoco-py/blob/master/mujoco_py/mjsim.pyx>`_
+class.  Users constuct a stateful ``MjSim`` instance from an MJCF model (similar to ``dm_control.Physics``), and this
+instance holds references to an ``mjModel`` instance and its associated ``mjData``.  In contrast, the MuJoCo Python
+bindings (``mujoco``) take a more low-level approach, as explained above: following the design principle of the C
+library, the ``mujoco`` module itself is stateless, and merely wraps the underlying native structs and functions.
+
+While a complete survey of mujoco-py is beyond the scope of this document, we offer below implementation notes for a
+non-exhaustive list of specific mujoco-py features:
+
+``mujoco_py.load_model_from_xml(bstring)``
+===========================================
+
+This factory function constructs a stateful ``MjSim`` instance.  When using ``mujoco``, the user should call the factory
+function ``mujoco.MjModel.from_xml_*`` as described :ref:`above <PyStructs>`. The user is then responsible for holding
+the resulting ``MjModel`` struct instance and explicitly generating the corresponding ``MjData`` by calling
+``mujoco.MjData(model)``.
+
+``sim.reset()``, ``sim.forward()``, ``sim.step()``
+==================================================
+
+Here as above, ``mujoco`` users needs to call the underlying library functions, passing instances of ``MjModel`` and
+``MjData``: :ref:`mujoco.mj_resetData(model, data) <mj_resetData>`, :ref:`mujoco.mj_forward(model, data) <mj_forward>`,
+and :ref:`mujoco.mj_step(model, data) <mj_step>`.
+
+``sim.get_state()``, ``sim.set_state(state)``, ``sim.get_flattened_state()``, ``sim.set_state_from_flattened(state)``
+=====================================================================================================================
+
+The MuJoCo library’s computation is deterministic given a specific input, as explained in the :ref:`Programming section
+<Simulation>`.  mujoco-py implements methods for getting and setting some of the relevant fields (and similarly
+``dm_control.Physics`` offers methods that correspond to the flattened case).  ``mujoco`` do not offer such abstraction,
+and the user is expected to get/set the values of the relevant fields explicitly.
+
+``sim.model.get_joint_qvel_addr(joint_name)``
+=============================================
+
+This is a convenience method in mujoco-py that returns a list of contiguous indices corresponding to this joint. The
+list starts from ``model.jnt_qposadr[joint_index]``, and its length depends on the joint type.  ``mujoco`` doesn't offer
+this functionality, but this list can be easily constructed using ``model.jnt_qposadr[joint_index]`` and ``xrange``.
+
+``sim.model.*_name2id(name)``
+=============================
+
+mujoco-py creates dicts in ``MjSim`` that allow for efficient lookup of indices for objects of different types:
+``site_name2id``, ``body_name2id`` etc.  These functions replace the function :ref:`mujoco.mj_name2id(model, type_enum,
+name) <mj_name2id>` whose current implementation is inefficient.  ``mujoco`` offers a different
+approach for using entity names – :ref:`named access <PyNamed>`, as well as access to the native :ref:`mj_name2id`.
+
+``sim.save(fstream, format_name)``
+==================================
+
+This is the one context in which the MuJoCo library (and therefore also ``mujoco``) is stateful: it holds a copy in
+memory of the last XML that was compiled, which is used in :ref:`mujoco.mj_saveLastXML(fname) <mj_saveLastXML>`. Note
+that mujoco-py’s implementation has a convenient extra feature, whereby the pose (as determined by ``sim.data``’s
+state) is transformed to a keyframe that’s added to the model before saving.  This extra feature is not currently
+available in ``mujoco``.
+
+Code Sample: open-loop rollout
+------------------------------
+
+We include a code sample showing how to add additional C/C++ functionality, exposed as a Python module via pybind11. The
+sample, implemented in ``rollout.cc`` and wrapped in ``rollout.py``, implements a common use case where tight loops
+implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling ``mj_step()`` in a loop), given an
+intial state and sequence of controls, and returning subsequent states and sensor values. The canonical usage form is
+
+  .. code-block:: python
+
+     state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+``initial_state`` is a ``nstate x nqva`` array, with ``nstate`` initial states of length ``nqva``, where ``nqva =
+model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities
+(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control
+sequences.
+
+The ``rollout`` function is designed to be completely stateless, so all inputs of the stepping pipeline are set and any
+values already present in the given ``MjData`` instance will have no effect on the output. In order to facilitate this,
+all inputs including ``time`` and ``qacc_warmstart`` are set to default values, as are auxillary controls
+(``qfrc_applied``, ``xfrc_applied`` and ``mocap_{pos,quat}``). These can also be optionally set by the user.
+
+Since the Global Interpreter Lock can be released, this function can be efficiently threaded using Python threads. See
+the ``test_threading`` function in ``rollout_test.py`` for an example of threaded operation.
+

doc/requirements.txt

@@ -8,3 +8,5 @@ pygments==2.7.4
 jq==1.1.1
 Jinja2==2.11.3
 wheel
+# see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe
+markupsafe==2.0.1

doc/unity.rst

@@ -16,8 +16,8 @@ Installation instructions
 
 The plug-in directory (available at https://github.com/deepmind/mujoco/tree/main/unity) includes a ``package.json``
 file.  Unity's package manager recognizes this file and will import the plug-in's C# codebase to your project.  In
-addition, Unity also needs the native MuJoCo library, which can be found in the specific platfomr archive at
+addition, Unity also needs the native MuJoCo library, which can be found in the specific platform archive at
-https://github.com/deepmind/mujoco/release.
+https://github.com/deepmind/mujoco/releases.
 
 On Unity version 2020.2 and later, the Package Manager will look for the native library file and copy it to the package
 directory when the package is imported. Alternatively, you can manually copy the native library to the package directory
@@ -170,7 +170,7 @@ components, you can call ``MjScene.CreateScene()`` when the initialization phase
 
 Scene recreation maintains continuity of physics and state in the following way:
 
-1. The position and velocity of joints is cached.
+1. The position and velocity of joints are cached.
 2. MuJoCo’s state is reset (to ``qpos0``) and Unity transforms are synchronized.
 3. A new XML is generated, creating a model that has the same ``qpos0`` as the previous one for the joints that
    persisted.

include/mjdata.h

@@ -15,7 +15,8 @@
 #ifndef MUJOCO_MJDATA_H_
 #define MUJOCO_MJDATA_H_
 
-#include "mjmodel.h"
+#include <mjtnum.h>
+#include <mjmodel.h>
 
 //---------------------------------- primitive types (mjt) -----------------------------------------
 
@@ -169,19 +170,19 @@ struct mjData_ {
   mjtNum*   qfrc_applied;         // applied generalized force                (nv x 1)
   mjtNum*   xfrc_applied;         // applied Cartesian force/torque           (nbody x 6)
 
+  // mocap data
+  mjtNum*   mocap_pos;            // positions of mocap bodies                (nmocap x 3)
+  mjtNum*   mocap_quat;           // orientations of mocap bodies             (nmocap x 4)
+
   // dynamics
   mjtNum*   qacc;                 // acceleration                             (nv x 1)
   mjtNum*   act_dot;              // time-derivative of actuator activation   (na x 1)
 
-  // mocap data
-  mjtNum*  mocap_pos;             // positions of mocap bodies                (nmocap x 3)
-  mjtNum*  mocap_quat;            // orientations of mocap bodies             (nmocap x 4)
-
   // user data
-  mjtNum*  userdata;              // user data, not touched by engine         (nuserdata x 1)
+  mjtNum*   userdata;             // user data, not touched by engine         (nuserdata x 1)
 
   // sensors
-  mjtNum*  sensordata;            // sensor data array                        (nsensordata x 1)
+  mjtNum*   sensordata;           // sensor data array                        (nsensordata x 1)
 
   //-------------------------------- POSITION dependent
 

include/mjexport.h

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MUJOCO_MUJOCO_EXPORT_H_
+#ifndef MUJOCO_MJEXPORT_H_
-#define MUJOCO_MUJOCO_EXPORT_H_
+#define MUJOCO_MJEXPORT_H_
 
 #if defined _WIN32 || defined __CYGWIN__
   #define MUJOCO_HELPER_DLL_IMPORT __declspec(dllimport)
@@ -44,4 +44,4 @@
   #define MJLOCAL MUJOCO_HELPER_DLL_LOCAL
 #endif
 
-#endif  // MUJOCO_MUJOCO_EXPORT_H_
+#endif  // MUJOCO_MJEXPORT_H_

include/mjmodel.h

@@ -15,21 +15,7 @@
 #ifndef MUJOCO_MJMODEL_H_
 #define MUJOCO_MJMODEL_H_
 
-//---------------------------------- floating-point definitions ------------------------------------
+#include <mjtnum.h>
-
-// compile-time configuration options
-#define mjUSEDOUBLE               // single or double precision for mjtNum
-
-
-// floating point data type and minval
-#ifdef mjUSEDOUBLE
-  typedef double mjtNum;
-  #define mjMINVAL    1E-15       // minimum value in any denominator
-#else
-  typedef float mjtNum;
-  #define mjMINVAL    1E-15f
-#endif
-
 
 // global constants
 #define mjPI            3.14159265358979323846
@@ -48,6 +34,7 @@
 #define mjNDYN          10        // number of actuator dynamics parameters
 #define mjNGAIN         10        // number of actuator gain parameters
 #define mjNBIAS         10        // number of actuator bias parameters
+#define mjNFLUID        12        // number of fluid interaction parameters
 #define mjNREF          2         // number of solver reference parameters
 #define mjNIMP          5         // number of solver impedance parameters
 #define mjNSOLVER       1000      // size of mjData.solver_XXX arrays
@@ -677,6 +664,7 @@ struct mjModel_ {
   mjtNum*   geom_friction;        // friction for (slide, spin, roll)         (ngeom x 3)
   mjtNum*   geom_margin;          // detect contact if dist<margin            (ngeom x 1)
   mjtNum*   geom_gap;             // include in solver if dist<margin-gap     (ngeom x 1)
+  mjtNum*   geom_fluid;           // fluid interaction parameters             (ngeom x mjNFLUID)
   mjtNum*   geom_user;            // user data                                (ngeom x nuser_geom)
   float*    geom_rgba;            // rgba when material is omitted            (ngeom x 4)
 
@@ -860,6 +848,8 @@ struct mjModel_ {
   int*      sensor_needstage;     // required compute stage (mjtStage)        (nsensor x 1)
   int*      sensor_objtype;       // type of sensorized object (mjtObj)       (nsensor x 1)
   int*      sensor_objid;         // id of sensorized object                  (nsensor x 1)
+  int*      sensor_reftype;       // type of reference frame (mjtObj)         (nsensor x 1)
+  int*      sensor_refid;         // id of reference frame; -1: global frame  (nsensor x 1)
   int*      sensor_dim;           // number of scalar outputs                 (nsensor x 1)
   int*      sensor_adr;           // address in sensor array                  (nsensor x 1)
   mjtNum*   sensor_cutoff;        // cutoff for real and positive; 0: ignore  (nsensor x 1)

include/mjrender.h

@@ -15,7 +15,11 @@
 #ifndef MUJOCO_MJRENDER_H_
 #define MUJOCO_MJRENDER_H_
 
-#include "mjmodel.h"
+#include <mjmodel.h>
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
 
 #define mjNAUX          10        // number of auxiliary buffers
 #define mjMAXTEXTURE    1000      // maximum number of textures
@@ -148,4 +152,7 @@ struct mjrContext_ {              // custom OpenGL context
 };
 typedef struct mjrContext_ mjrContext;
 
+#if defined(__cplusplus)
+}
+#endif
 #endif  // MUJOCO_MJRENDER_H_

include/mjtnum.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_INCLUDE_MJTNUM_H_
+#define MUJOCO_INCLUDE_MJTNUM_H_
+
+//---------------------------------- floating-point definitions ------------------------------------
+
+// compile-time configuration options
+#define mjUSEDOUBLE               // single or double precision for mjtNum
+
+
+// floating point data type and minval
+#ifdef mjUSEDOUBLE
+  typedef double mjtNum;
+  #define mjMINVAL    1E-15       // minimum value in any denominator
+#else
+  typedef float mjtNum;
+  #define mjMINVAL    1E-15f
+#endif
+
+
+#endif  // MUJOCO_INCLUDE_MJTNUM_H_

include/mjui.h

@@ -15,7 +15,7 @@
 #ifndef MUJOCO_MJUI_H_
 #define MUJOCO_MJUI_H_
 
-#include "mjrender.h"
+#include <mjrender.h>
 
 #define mjMAXUISECT     10      // maximum number of sections
 #define mjMAXUIITEM     80      // maximum number of items per section

include/mjvisualize.h

@@ -15,9 +15,11 @@
 #ifndef MUJOCO_MJVISUALIZE_H_
 #define MUJOCO_MJVISUALIZE_H_
 
-#include "mjmodel.h"
+#include <mjtnum.h>
+#include <mjmodel.h>
 
 #define mjNGROUP        6         // number of geom, site, joint groups with visflags
+#define mjMAXLIGHT      100       // maximum number of lights in a scene
 #define mjMAXOVERLAY    500       // maximum number of characters in overlay text
 #define mjMAXLINE       100       // maximum number of lines per plot
 #define mjMAXLINEPNT    1000      // maximum number points per line
@@ -278,7 +280,7 @@ struct mjvScene_ {                // abstract scene passed to OpenGL renderer
 
   // OpenGL lights
   int      nlight;                // number of lights currently in buffer
-  mjvLight lights[8];             // buffer for lights
+  mjvLight lights[mjMAXLIGHT];    // buffer for lights
 
   // OpenGL cameras
   mjvGLCamera camera[2];          // left and right camera

include/mjxmacro.h

@@ -15,50 +15,50 @@
 #ifndef MUJOCO_MJXMACRO_H_
 #define MUJOCO_MJXMACRO_H_
 
-//-------------------------------- mjOption ---------------------------------------------
+
+//-------------------------------- mjOption --------------------------------------------------------
 
 // scalar fields of mjOption
-#define MJOPTION_FLOATS           \
+#define MJOPTION_FLOATS             \
-    X( mjtNum,  timestep        ) \
+    X( mjtNum,  timestep        )   \
-    X( mjtNum,  apirate         ) \
+    X( mjtNum,  apirate         )   \
-    X( mjtNum,  impratio        ) \
+    X( mjtNum,  impratio        )   \
-    X( mjtNum,  tolerance       ) \
+    X( mjtNum,  tolerance       )   \
-    X( mjtNum,  noslip_tolerance ) \
+    X( mjtNum,  noslip_tolerance )  \
-    X( mjtNum,  mpr_tolerance   ) \
+    X( mjtNum,  mpr_tolerance   )   \
-    X( mjtNum,  density         ) \
+    X( mjtNum,  density         )   \
-    X( mjtNum,  viscosity       ) \
+    X( mjtNum,  viscosity       )   \
-    X( mjtNum,  o_margin        ) \
+    X( mjtNum,  o_margin        )   \
 
 
-#define MJOPTION_INTS             \
+#define MJOPTION_INTS               \
-    X( int,     integrator      ) \
+    X( int,     integrator      )   \
-    X( int,     collision       ) \
+    X( int,     collision       )   \
-    X( int,     cone            ) \
+    X( int,     cone            )   \
-    X( int,     jacobian        ) \
+    X( int,     jacobian        )   \
-    X( int,     solver          ) \
+    X( int,     solver          )   \
-    X( int,     iterations      ) \
+    X( int,     iterations      )   \
     X( int,     noslip_iterations ) \
-    X( int,     mpr_iterations  ) \
+    X( int,     mpr_iterations  )   \
-    X( int,     disableflags    ) \
+    X( int,     disableflags    )   \
     X( int,     enableflags     )
 
 
-#define MJOPTION_SCALARS          \
+#define MJOPTION_SCALARS            \
-    MJOPTION_FLOATS               \
+    MJOPTION_FLOATS                 \
     MJOPTION_INTS
 
 
 // vector fields of mjOption
-#define MJOPTION_VECTORS          \
+#define MJOPTION_VECTORS            \
-    X( gravity,         3       ) \
+    X( gravity,         3       )   \
-    X( wind,            3       ) \
+    X( wind,            3       )   \
-    X( magnetic,        3       ) \
+    X( magnetic,        3       )   \
-    X( o_solref,        mjNREF  ) \
+    X( o_solref,        mjNREF  )   \
     X( o_solimp,        mjNIMP  )
 
 
-
+//-------------------------------- mjModel ---------------------------------------------------------
-//-------------------------------- mjModel ----------------------------------------------
 
 // int fields of mjModel
 #define MJMODEL_INTS        \
@@ -122,396 +122,402 @@
 
 
 // define symbols needed in MJMODEL_POINTERS (corresponding to number of columns)
-#define MJMODEL_POINTERS_PREAMBLE( m ) \
+#define MJMODEL_POINTERS_PREAMBLE( m )      \
-    int nuser_body = m->nuser_body; \
+    int nuser_body = m->nuser_body;         \
-    int nuser_jnt = m->nuser_jnt; \
+    int nuser_jnt = m->nuser_jnt;           \
-    int nuser_geom = m->nuser_geom; \
+    int nuser_geom = m->nuser_geom;         \
-    int nuser_site = m->nuser_site; \
+    int nuser_site = m->nuser_site;         \
-    int nuser_cam = m->nuser_cam; \
+    int nuser_cam = m->nuser_cam;           \
-    int nuser_tendon = m->nuser_tendon; \
+    int nuser_tendon = m->nuser_tendon;     \
     int nuser_actuator = m->nuser_actuator; \
-    int nuser_sensor = m->nuser_sensor; \
+    int nuser_sensor = m->nuser_sensor;     \
-    int nq = m->nq; \
+    int nq = m->nq;                         \
-    int nv = m->nv; \
+    int nv = m->nv;                         \
-    int na = m->na; \
+    int na = m->na;                         \
-    int nmocap3 = 3*m->nmocap; \
+    int nmocap = m->nmocap;
-    int nmocap4 = 4*m->nmocap;
+
+
+// macro for annotating that an array size in an X macro is a member of mjModel
+// by default this macro does nothing, but users can redefine it as necessary
+#define MJ_M(n) n
 
 
 // pointer fields of mjModel
-#define MJMODEL_POINTERS                                          \
+#define MJMODEL_POINTERS                                                     \
-    X( mjtNum,  qpos0,                   nq,        1           ) \
+    X( mjtNum,  qpos0,                 nq,            1                    ) \
-    X( mjtNum,  qpos_spring,             nq,        1           ) \
+    X( mjtNum,  qpos_spring,           nq,            1                    ) \
-    X( int,     body_parentid,           nbody,     1           ) \
+    X( int,     body_parentid,         nbody,         1                    ) \
-    X( int,     body_rootid,             nbody,     1           ) \
+    X( int,     body_rootid,           nbody,         1                    ) \
-    X( int,     body_weldid,             nbody,     1           ) \
+    X( int,     body_weldid,           nbody,         1                    ) \
-    X( int,     body_mocapid,            nbody,     1           ) \
+    X( int,     body_mocapid,          nbody,         1                    ) \
-    X( int,     body_jntnum,             nbody,     1           ) \
+    X( int,     body_jntnum,           nbody,         1                    ) \
-    X( int,     body_jntadr,             nbody,     1           ) \
+    X( int,     body_jntadr,           nbody,         1                    ) \
-    X( int,     body_dofnum,             nbody,     1           ) \
+    X( int,     body_dofnum,           nbody,         1                    ) \
-    X( int,     body_dofadr,             nbody,     1           ) \
+    X( int,     body_dofadr,           nbody,         1                    ) \
-    X( int,     body_geomnum,            nbody,     1           ) \
+    X( int,     body_geomnum,          nbody,         1                    ) \
-    X( int,     body_geomadr,            nbody,     1           ) \
+    X( int,     body_geomadr,          nbody,         1                    ) \
-    X( mjtByte, body_simple,             nbody,     1           ) \
+    X( mjtByte, body_simple,           nbody,         1                    ) \
-    X( mjtByte, body_sameframe,          nbody,     1           ) \
+    X( mjtByte, body_sameframe,        nbody,         1                    ) \
-    X( mjtNum,  body_pos,                nbody,     3           ) \
+    X( mjtNum,  body_pos,              nbody,         3                    ) \
-    X( mjtNum,  body_quat,               nbody,     4           ) \
+    X( mjtNum,  body_quat,             nbody,         4                    ) \
-    X( mjtNum,  body_ipos,               nbody,     3           ) \
+    X( mjtNum,  body_ipos,             nbody,         3                    ) \
-    X( mjtNum,  body_iquat,              nbody,     4           ) \
+    X( mjtNum,  body_iquat,            nbody,         4                    ) \
-    X( mjtNum,  body_mass,               nbody,     1           ) \
+    X( mjtNum,  body_mass,             nbody,         1                    ) \
-    X( mjtNum,  body_subtreemass,        nbody,     1           ) \
+    X( mjtNum,  body_subtreemass,      nbody,         1                    ) \
-    X( mjtNum,  body_inertia,            nbody,     3           ) \
+    X( mjtNum,  body_inertia,          nbody,         3                    ) \
-    X( mjtNum,  body_invweight0,         nbody,     2           ) \
+    X( mjtNum,  body_invweight0,       nbody,         2                    ) \
-    X( mjtNum,  body_user,               nbody,     nuser_body  ) \
+    X( mjtNum,  body_user,             nbody,         MJ_M(nuser_body)     ) \
-    X( int,     jnt_type,                njnt,      1           ) \
+    X( int,     jnt_type,              njnt,          1                    ) \
-    X( int,     jnt_qposadr,             njnt,      1           ) \
+    X( int,     jnt_qposadr,           njnt,          1                    ) \
-    X( int,     jnt_dofadr,              njnt,      1           ) \
+    X( int,     jnt_dofadr,            njnt,          1                    ) \
-    X( int,     jnt_bodyid,              njnt,      1           ) \
+    X( int,     jnt_bodyid,            njnt,          1                    ) \
-    X( int,     jnt_group,               njnt,      1           ) \
+    X( int,     jnt_group,             njnt,          1                    ) \
-    X( mjtByte, jnt_limited,             njnt,      1           ) \
+    X( mjtByte, jnt_limited,           njnt,          1                    ) \
-    X( mjtNum,  jnt_solref,              njnt,      mjNREF      ) \
+    X( mjtNum,  jnt_solref,            njnt,          mjNREF               ) \
-    X( mjtNum,  jnt_solimp,              njnt,      mjNIMP      ) \
+    X( mjtNum,  jnt_solimp,            njnt,          mjNIMP               ) \
-    X( mjtNum,  jnt_pos,                 njnt,      3           ) \
+    X( mjtNum,  jnt_pos,               njnt,          3                    ) \
-    X( mjtNum,  jnt_axis,                njnt,      3           ) \
+    X( mjtNum,  jnt_axis,              njnt,          3                    ) \
-    X( mjtNum,  jnt_stiffness,           njnt,      1           ) \
+    X( mjtNum,  jnt_stiffness,         njnt,          1                    ) \
-    X( mjtNum,  jnt_range,               njnt,      2           ) \
+    X( mjtNum,  jnt_range,             njnt,          2                    ) \
-    X( mjtNum,  jnt_margin,              njnt,      1           ) \
+    X( mjtNum,  jnt_margin,            njnt,          1                    ) \
-    X( mjtNum,  jnt_user,                njnt,      nuser_jnt   ) \
+    X( mjtNum,  jnt_user,              njnt,          MJ_M(nuser_jnt)      ) \
-    X( int,     dof_bodyid,              nv,        1           ) \
+    X( int,     dof_bodyid,            nv,            1                    ) \
-    X( int,     dof_jntid,               nv,        1           ) \
+    X( int,     dof_jntid,             nv,            1                    ) \
-    X( int,     dof_parentid,            nv,        1           ) \
+    X( int,     dof_parentid,          nv,            1                    ) \
-    X( int,     dof_Madr,                nv,        1           ) \
+    X( int,     dof_Madr,              nv,            1                    ) \
-    X( int,     dof_simplenum,           nv,        1           ) \
+    X( int,     dof_simplenum,         nv,            1                    ) \
-    X( mjtNum,  dof_solref,              nv,        mjNREF      ) \
+    X( mjtNum,  dof_solref,            nv,            mjNREF               ) \
-    X( mjtNum,  dof_solimp,              nv,        mjNIMP      ) \
+    X( mjtNum,  dof_solimp,            nv,            mjNIMP               ) \
-    X( mjtNum,  dof_frictionloss,        nv,        1           ) \
+    X( mjtNum,  dof_frictionloss,      nv,            1                    ) \
-    X( mjtNum,  dof_armature,            nv,        1           ) \
+    X( mjtNum,  dof_armature,          nv,            1                    ) \
-    X( mjtNum,  dof_damping,             nv,        1           ) \
+    X( mjtNum,  dof_damping,           nv,            1                    ) \
-    X( mjtNum,  dof_invweight0,          nv,        1           ) \
+    X( mjtNum,  dof_invweight0,        nv,            1                    ) \
-    X( mjtNum,  dof_M0,                  nv,        1           ) \
+    X( mjtNum,  dof_M0,                nv,            1                    ) \
-    X( int,     geom_type,               ngeom,     1           ) \
+    X( int,     geom_type,             ngeom,         1                    ) \
-    X( int,     geom_contype,            ngeom,     1           ) \
+    X( int,     geom_contype,          ngeom,         1                    ) \
-    X( int,     geom_conaffinity,        ngeom,     1           ) \
+    X( int,     geom_conaffinity,      ngeom,         1                    ) \
-    X( int,     geom_condim,             ngeom,     1           ) \
+    X( int,     geom_condim,           ngeom,         1                    ) \
-    X( int,     geom_bodyid,             ngeom,     1           ) \
+    X( int,     geom_bodyid,           ngeom,         1                    ) \
-    X( int,     geom_dataid,             ngeom,     1           ) \
+    X( int,     geom_dataid,           ngeom,         1                    ) \
-    X( int,     geom_matid,              ngeom,     1           ) \
+    X( int,     geom_matid,            ngeom,         1                    ) \
-    X( int,     geom_group,              ngeom,     1           ) \
+    X( int,     geom_group,            ngeom,         1                    ) \
-    X( int,     geom_priority,           ngeom,     1           ) \
+    X( int,     geom_priority,         ngeom,         1                    ) \
-    X( mjtByte, geom_sameframe,          ngeom,     1           ) \
+    X( mjtByte, geom_sameframe,        ngeom,         1                    ) \
-    X( mjtNum,  geom_solmix,             ngeom,     1           ) \
+    X( mjtNum,  geom_solmix,           ngeom,         1                    ) \
-    X( mjtNum,  geom_solref,             ngeom,     mjNREF      ) \
+    X( mjtNum,  geom_solref,           ngeom,         mjNREF               ) \
-    X( mjtNum,  geom_solimp,             ngeom,     mjNIMP      ) \
+    X( mjtNum,  geom_solimp,           ngeom,         mjNIMP               ) \
-    X( mjtNum,  geom_size,               ngeom,     3           ) \
+    X( mjtNum,  geom_size,             ngeom,         3                    ) \
-    X( mjtNum,  geom_rbound,             ngeom,     1           ) \
+    X( mjtNum,  geom_rbound,           ngeom,         1                    ) \
-    X( mjtNum,  geom_pos,                ngeom,     3           ) \
+    X( mjtNum,  geom_pos,              ngeom,         3                    ) \
-    X( mjtNum,  geom_quat,               ngeom,     4           ) \
+    X( mjtNum,  geom_quat,             ngeom,         4                    ) \
-    X( mjtNum,  geom_friction,           ngeom,     3           ) \
+    X( mjtNum,  geom_friction,         ngeom,         3                    ) \
-    X( mjtNum,  geom_margin,             ngeom,     1           ) \
+    X( mjtNum,  geom_margin,           ngeom,         1                    ) \
-    X( mjtNum,  geom_gap,                ngeom,     1           ) \
+    X( mjtNum,  geom_gap,              ngeom,         1                    ) \
-    X( mjtNum,  geom_user,               ngeom,     nuser_geom  ) \
+    X( mjtNum,  geom_fluid,            ngeom,         mjNFLUID             ) \
-    X( float,   geom_rgba,               ngeom,     4           ) \
+    X( mjtNum,  geom_user,             ngeom,         MJ_M(nuser_geom)     ) \
-    X( int,     site_type,               nsite,     1           ) \
+    X( float,   geom_rgba,             ngeom,         4                    ) \
-    X( int,     site_bodyid,             nsite,     1           ) \
+    X( int,     site_type,             nsite,         1                    ) \
-    X( int,     site_matid,              nsite,     1           ) \
+    X( int,     site_bodyid,           nsite,         1                    ) \
-    X( int,     site_group,              nsite,     1           ) \
+    X( int,     site_matid,            nsite,         1                    ) \
-    X( mjtByte, site_sameframe,          nsite,     1           ) \
+    X( int,     site_group,            nsite,         1                    ) \
-    X( mjtNum,  site_size,               nsite,     3           ) \
+    X( mjtByte, site_sameframe,        nsite,         1                    ) \
-    X( mjtNum,  site_pos,                nsite,     3           ) \
+    X( mjtNum,  site_size,             nsite,         3                    ) \
-    X( mjtNum,  site_quat,               nsite,     4           ) \
+    X( mjtNum,  site_pos,              nsite,         3                    ) \
-    X( mjtNum,  site_user,               nsite,     nuser_site  ) \
+    X( mjtNum,  site_quat,             nsite,         4                    ) \
-    X( float,   site_rgba,               nsite,     4           ) \
+    X( mjtNum,  site_user,             nsite,         MJ_M(nuser_site)     ) \
-    X( int,     cam_mode,                ncam,      1           ) \
+    X( float,   site_rgba,             nsite,         4                    ) \
-    X( int,     cam_bodyid,              ncam,      1           ) \
+    X( int,     cam_mode,              ncam,          1                    ) \
-    X( int,     cam_targetbodyid,        ncam,      1           ) \
+    X( int,     cam_bodyid,            ncam,          1                    ) \
-    X( mjtNum,  cam_pos,                 ncam,      3           ) \
+    X( int,     cam_targetbodyid,      ncam,          1                    ) \
-    X( mjtNum,  cam_quat,                ncam,      4           ) \
+    X( mjtNum,  cam_pos,               ncam,          3                    ) \
-    X( mjtNum,  cam_poscom0,             ncam,      3           ) \
+    X( mjtNum,  cam_quat,              ncam,          4                    ) \
-    X( mjtNum,  cam_pos0,                ncam,      3           ) \
+    X( mjtNum,  cam_poscom0,           ncam,          3                    ) \
-    X( mjtNum,  cam_mat0,                ncam,      9           ) \
+    X( mjtNum,  cam_pos0,              ncam,          3                    ) \
-    X( mjtNum,  cam_fovy,                ncam,      1           ) \
+    X( mjtNum,  cam_mat0,              ncam,          9                    ) \
-    X( mjtNum,  cam_ipd,                 ncam,      1           ) \
+    X( mjtNum,  cam_fovy,              ncam,          1                    ) \
-    X( mjtNum,  cam_user,                ncam,      nuser_cam   ) \
+    X( mjtNum,  cam_ipd,               ncam,          1                    ) \
-    X( int,     light_mode,              nlight,    1           ) \
+    X( mjtNum,  cam_user,              ncam,          MJ_M(nuser_cam)      ) \
-    X( int,     light_bodyid,            nlight,    1           ) \
+    X( int,     light_mode,            nlight,        1                    ) \
-    X( int,     light_targetbodyid,      nlight,    1           ) \
+    X( int,     light_bodyid,          nlight,        1                    ) \
-    X( mjtByte, light_directional,       nlight,    1           ) \
+    X( int,     light_targetbodyid,    nlight,        1                    ) \
-    X( mjtByte, light_castshadow,        nlight,    1           ) \
+    X( mjtByte, light_directional,     nlight,        1                    ) \
-    X( mjtByte, light_active,            nlight,    1           ) \
+    X( mjtByte, light_castshadow,      nlight,        1                    ) \
-    X( mjtNum,  light_pos,               nlight,    3           ) \
+    X( mjtByte, light_active,          nlight,        1                    ) \
-    X( mjtNum,  light_dir,               nlight,    3           ) \
+    X( mjtNum,  light_pos,             nlight,        3                    ) \
-    X( mjtNum,  light_poscom0,           nlight,    3           ) \
+    X( mjtNum,  light_dir,             nlight,        3                    ) \
-    X( mjtNum,  light_pos0,              nlight,    3           ) \
+    X( mjtNum,  light_poscom0,         nlight,        3                    ) \
-    X( mjtNum,  light_dir0,              nlight,    3           ) \
+    X( mjtNum,  light_pos0,            nlight,        3                    ) \
-    X( float,   light_attenuation,       nlight,    3           ) \
+    X( mjtNum,  light_dir0,            nlight,        3                    ) \
-    X( float,   light_cutoff,            nlight,    1           ) \
+    X( float,   light_attenuation,     nlight,        3                    ) \
-    X( float,   light_exponent,          nlight,    1           ) \
+    X( float,   light_cutoff,          nlight,        1                    ) \
-    X( float,   light_ambient,           nlight,    3           ) \
+    X( float,   light_exponent,        nlight,        1                    ) \
-    X( float,   light_diffuse,           nlight,    3           ) \
+    X( float,   light_ambient,         nlight,        3                    ) \
-    X( float,   light_specular,          nlight,    3           ) \
+    X( float,   light_diffuse,         nlight,        3                    ) \
-    X( int,     mesh_vertadr,            nmesh,     1           ) \
+    X( float,   light_specular,        nlight,        3                    ) \
-    X( int,     mesh_vertnum,            nmesh,     1           ) \
+    X( int,     mesh_vertadr,          nmesh,         1                    ) \
-    X( int,     mesh_texcoordadr,        nmesh,     1           ) \
+    X( int,     mesh_vertnum,          nmesh,         1                    ) \
-    X( int,     mesh_faceadr,            nmesh,     1           ) \
+    X( int,     mesh_texcoordadr,      nmesh,         1                    ) \
-    X( int,     mesh_facenum,            nmesh,     1           ) \
+    X( int,     mesh_faceadr,          nmesh,         1                    ) \
-    X( int,     mesh_graphadr,           nmesh,     1           ) \
+    X( int,     mesh_facenum,          nmesh,         1                    ) \
-    X( float,   mesh_vert,               nmeshvert, 3           ) \
+    X( int,     mesh_graphadr,         nmesh,         1                    ) \
-    X( float,   mesh_normal,             nmeshvert, 3           ) \
+    X( float,   mesh_vert,             nmeshvert,     3                    ) \
-    X( float,   mesh_texcoord,           nmeshtexvert, 2        ) \
+    X( float,   mesh_normal,           nmeshvert,     3                    ) \
-    X( int,     mesh_face,               nmeshface, 3           ) \
+    X( float,   mesh_texcoord,         nmeshtexvert,  2                    ) \
-    X( int,     mesh_graph,              nmeshgraph,1           ) \
+    X( int,     mesh_face,             nmeshface,     3                    ) \
-    X( int,     skin_matid,              nskin,     1           ) \
+    X( int,     mesh_graph,            nmeshgraph,    1                    ) \
-    X( float,   skin_rgba,               nskin,     4           ) \
+    X( int,     skin_matid,            nskin,         1                    ) \
-    X( float,   skin_inflate,            nskin,     1           ) \
+    X( float,   skin_rgba,             nskin,         4                    ) \
-    X( int,     skin_vertadr,            nskin,     1           ) \
+    X( float,   skin_inflate,          nskin,         1                    ) \
-    X( int,     skin_vertnum,            nskin,     1           ) \
+    X( int,     skin_vertadr,          nskin,         1                    ) \
-    X( int,     skin_texcoordadr,        nskin,     1           ) \
+    X( int,     skin_vertnum,          nskin,         1                    ) \
-    X( int,     skin_faceadr,            nskin,     1           ) \
+    X( int,     skin_texcoordadr,      nskin,         1                    ) \
-    X( int,     skin_facenum,            nskin,     1           ) \
+    X( int,     skin_faceadr,          nskin,         1                    ) \
-    X( int,     skin_boneadr,            nskin,     1           ) \
+    X( int,     skin_facenum,          nskin,         1                    ) \
-    X( int,     skin_bonenum,            nskin,     1           ) \
+    X( int,     skin_boneadr,          nskin,         1                    ) \
-    X( float,   skin_vert,               nskinvert, 3           ) \
+    X( int,     skin_bonenum,          nskin,         1                    ) \
-    X( float,   skin_texcoord,           nskintexvert, 2        ) \
+    X( float,   skin_vert,             nskinvert,     3                    ) \
-    X( int,     skin_face,               nskinface, 3           ) \
+    X( float,   skin_texcoord,         nskintexvert,  2                    ) \
-    X( int,     skin_bonevertadr,        nskinbone, 1           ) \
+    X( int,     skin_face,             nskinface,     3                    ) \
-    X( int,     skin_bonevertnum,        nskinbone, 1           ) \
+    X( int,     skin_bonevertadr,      nskinbone,     1                    ) \
-    X( float,   skin_bonebindpos,        nskinbone, 3           ) \
+    X( int,     skin_bonevertnum,      nskinbone,     1                    ) \
-    X( float,   skin_bonebindquat,       nskinbone, 4           ) \
+    X( float,   skin_bonebindpos,      nskinbone,     3                    ) \
-    X( int,     skin_bonebodyid,         nskinbone, 1           ) \
+    X( float,   skin_bonebindquat,     nskinbone,     4                    ) \
-    X( int,     skin_bonevertid,         nskinbonevert, 1       ) \
+    X( int,     skin_bonebodyid,       nskinbone,     1                    ) \
-    X( float,   skin_bonevertweight,     nskinbonevert, 1       ) \
+    X( int,     skin_bonevertid,       nskinbonevert, 1                    ) \
-    X( mjtNum,  hfield_size,             nhfield,   4           ) \
+    X( float,   skin_bonevertweight,   nskinbonevert, 1                    ) \
-    X( int,     hfield_nrow,             nhfield,   1           ) \
+    X( mjtNum,  hfield_size,           nhfield,       4                    ) \
-    X( int,     hfield_ncol,             nhfield,   1           ) \
+    X( int,     hfield_nrow,           nhfield,       1                    ) \
-    X( int,     hfield_adr,              nhfield,   1           ) \
+    X( int,     hfield_ncol,           nhfield,       1                    ) \
-    X( float,   hfield_data,             nhfielddata, 1         ) \
+    X( int,     hfield_adr,            nhfield,       1                    ) \
-    X( int,     tex_type,                ntex,      1           ) \
+    X( float,   hfield_data,           nhfielddata,   1                    ) \
-    X( int,     tex_height,              ntex,      1           ) \
+    X( int,     tex_type,              ntex,          1                    ) \
-    X( int,     tex_width,               ntex,      1           ) \
+    X( int,     tex_height,            ntex,          1                    ) \
-    X( int,     tex_adr,                 ntex,      1           ) \
+    X( int,     tex_width,             ntex,          1                    ) \
-    X( mjtByte, tex_rgb,                 ntexdata,  1           ) \
+    X( int,     tex_adr,               ntex,          1                    ) \
-    X( int,     mat_texid,               nmat,      1           ) \
+    X( mjtByte, tex_rgb,               ntexdata,      1                    ) \
-    X( mjtByte, mat_texuniform,          nmat,      1           ) \
+    X( int,     mat_texid,             nmat,          1                    ) \
-    X( float,   mat_texrepeat,           nmat,      2           ) \
+    X( mjtByte, mat_texuniform,        nmat,          1                    ) \
-    X( float,   mat_emission,            nmat,      1           ) \
+    X( float,   mat_texrepeat,         nmat,          2                    ) \
-    X( float,   mat_specular,            nmat,      1           ) \
+    X( float,   mat_emission,          nmat,          1                    ) \
-    X( float,   mat_shininess,           nmat,      1           ) \
+    X( float,   mat_specular,          nmat,          1                    ) \
-    X( float,   mat_reflectance,         nmat,      1           ) \
+    X( float,   mat_shininess,         nmat,          1                    ) \
-    X( float,   mat_rgba,                nmat,      4           ) \
+    X( float,   mat_reflectance,       nmat,          1                    ) \
-    X( int,     pair_dim,                npair,     1           ) \
+    X( float,   mat_rgba,              nmat,          4                    ) \
-    X( int,     pair_geom1,              npair,     1           ) \
+    X( int,     pair_dim,              npair,         1                    ) \
-    X( int,     pair_geom2,              npair,     1           ) \
+    X( int,     pair_geom1,            npair,         1                    ) \
-    X( int,     pair_signature,          npair,     1           ) \
+    X( int,     pair_geom2,            npair,         1                    ) \
-    X( mjtNum,  pair_solref,             npair,     mjNREF      ) \
+    X( int,     pair_signature,        npair,         1                    ) \
-    X( mjtNum,  pair_solimp,             npair,     mjNIMP      ) \
+    X( mjtNum,  pair_solref,           npair,         mjNREF               ) \
-    X( mjtNum,  pair_margin,             npair,     1           ) \
+    X( mjtNum,  pair_solimp,           npair,         mjNIMP               ) \
-    X( mjtNum,  pair_gap,                npair,     1           ) \
+    X( mjtNum,  pair_margin,           npair,         1                    ) \
-    X( mjtNum,  pair_friction,           npair,     5           ) \
+    X( mjtNum,  pair_gap,              npair,         1                    ) \
-    X( int,     exclude_signature,       nexclude,  1           ) \
+    X( mjtNum,  pair_friction,         npair,         5                    ) \
-    X( int,     eq_type,                 neq,       1           ) \
+    X( int,     exclude_signature,     nexclude,      1                    ) \
-    X( int,     eq_obj1id,               neq,       1           ) \
+    X( int,     eq_type,               neq,           1                    ) \
-    X( int,     eq_obj2id,               neq,       1           ) \
+    X( int,     eq_obj1id,             neq,           1                    ) \
-    X( mjtByte, eq_active,               neq,       1           ) \
+    X( int,     eq_obj2id,             neq,           1                    ) \
-    X( mjtNum,  eq_solref,               neq,       mjNREF      ) \
+    X( mjtByte, eq_active,             neq,           1                    ) \
-    X( mjtNum,  eq_solimp,               neq,       mjNIMP      ) \
+    X( mjtNum,  eq_solref,             neq,           mjNREF               ) \
-    X( mjtNum,  eq_data,                 neq,       mjNEQDATA   ) \
+    X( mjtNum,  eq_solimp,             neq,           mjNIMP               ) \
-    X( int,     tendon_adr,              ntendon,   1           ) \
+    X( mjtNum,  eq_data,               neq,           mjNEQDATA            ) \
-    X( int,     tendon_num,              ntendon,   1           ) \
+    X( int,     tendon_adr,            ntendon,       1                    ) \
-    X( int,     tendon_matid,            ntendon,   1           ) \
+    X( int,     tendon_num,            ntendon,       1                    ) \
-    X( int,     tendon_group,            ntendon,   1           ) \
+    X( int,     tendon_matid,          ntendon,       1                    ) \
-    X( mjtByte, tendon_limited,          ntendon,   1           ) \
+    X( int,     tendon_group,          ntendon,       1                    ) \
-    X( mjtNum,  tendon_width,            ntendon,   1           ) \
+    X( mjtByte, tendon_limited,        ntendon,       1                    ) \
-    X( mjtNum,  tendon_solref_lim,       ntendon,   mjNREF      ) \
+    X( mjtNum,  tendon_width,          ntendon,       1                    ) \
-    X( mjtNum,  tendon_solimp_lim,       ntendon,   mjNIMP      ) \
+    X( mjtNum,  tendon_solref_lim,     ntendon,       mjNREF               ) \
-    X( mjtNum,  tendon_solref_fri,       ntendon,   mjNREF      ) \
+    X( mjtNum,  tendon_solimp_lim,     ntendon,       mjNIMP               ) \
-    X( mjtNum,  tendon_solimp_fri,       ntendon,   mjNIMP      ) \
+    X( mjtNum,  tendon_solref_fri,     ntendon,       mjNREF               ) \
-    X( mjtNum,  tendon_range,            ntendon,   2           ) \
+    X( mjtNum,  tendon_solimp_fri,     ntendon,       mjNIMP               ) \
-    X( mjtNum,  tendon_margin,           ntendon,   1           ) \
+    X( mjtNum,  tendon_range,          ntendon,       2                    ) \
-    X( mjtNum,  tendon_stiffness,        ntendon,   1           ) \
+    X( mjtNum,  tendon_margin,         ntendon,       1                    ) \
-    X( mjtNum,  tendon_damping,          ntendon,   1           ) \
+    X( mjtNum,  tendon_stiffness,      ntendon,       1                    ) \
-    X( mjtNum,  tendon_frictionloss,     ntendon,   1           ) \
+    X( mjtNum,  tendon_damping,        ntendon,       1                    ) \
-    X( mjtNum,  tendon_lengthspring,     ntendon,   1           ) \
+    X( mjtNum,  tendon_frictionloss,   ntendon,       1                    ) \
-    X( mjtNum,  tendon_length0,          ntendon,   1           ) \
+    X( mjtNum,  tendon_lengthspring,   ntendon,       1                    ) \
-    X( mjtNum,  tendon_invweight0,       ntendon,   1           ) \
+    X( mjtNum,  tendon_length0,        ntendon,       1                    ) \
-    X( mjtNum,  tendon_user,             ntendon,   nuser_tendon) \
+    X( mjtNum,  tendon_invweight0,     ntendon,       1                    ) \
-    X( float,   tendon_rgba,             ntendon,   4           ) \
+    X( mjtNum,  tendon_user,           ntendon,       MJ_M(nuser_tendon)   ) \
-    X( int,     wrap_type,               nwrap,     1           ) \
+    X( float,   tendon_rgba,           ntendon,       4                    ) \
-    X( int,     wrap_objid,              nwrap,     1           ) \
+    X( int,     wrap_type,             nwrap,         1                    ) \
-    X( mjtNum,  wrap_prm,                nwrap,     1           ) \
+    X( int,     wrap_objid,            nwrap,         1                    ) \
-    X( int,     actuator_trntype,        nu,        1           ) \
+    X( mjtNum,  wrap_prm,              nwrap,         1                    ) \
-    X( int,     actuator_dyntype,        nu,        1           ) \
+    X( int,     actuator_trntype,      nu,            1                    ) \
-    X( int,     actuator_gaintype,       nu,        1           ) \
+    X( int,     actuator_dyntype,      nu,            1                    ) \
-    X( int,     actuator_biastype,       nu,        1           ) \
+    X( int,     actuator_gaintype,     nu,            1                    ) \
-    X( int,     actuator_trnid,          nu,        2           ) \
+    X( int,     actuator_biastype,     nu,            1                    ) \
-    X( int,     actuator_group,          nu,        1           ) \
+    X( int,     actuator_trnid,        nu,            2                    ) \
-    X( mjtByte, actuator_ctrllimited,    nu,        1           ) \
+    X( int,     actuator_group,        nu,            1                    ) \
-    X( mjtByte, actuator_forcelimited,   nu,        1           ) \
+    X( mjtByte, actuator_ctrllimited,  nu,            1                    ) \
-    X( mjtNum,  actuator_dynprm,         nu,        mjNDYN      ) \
+    X( mjtByte, actuator_forcelimited, nu,            1                    ) \
-    X( mjtNum,  actuator_gainprm,        nu,        mjNGAIN     ) \
+    X( mjtNum,  actuator_dynprm,       nu,            mjNDYN               ) \
-    X( mjtNum,  actuator_biasprm,        nu,        mjNBIAS     ) \
+    X( mjtNum,  actuator_gainprm,      nu,            mjNGAIN              ) \
-    X( mjtNum,  actuator_ctrlrange,      nu,        2           ) \
+    X( mjtNum,  actuator_biasprm,      nu,            mjNBIAS              ) \
-    X( mjtNum,  actuator_forcerange,     nu,        2           ) \
+    X( mjtNum,  actuator_ctrlrange,    nu,            2                    ) \
-    X( mjtNum,  actuator_gear,           nu,        6           ) \
+    X( mjtNum,  actuator_forcerange,   nu,            2                    ) \
-    X( mjtNum,  actuator_cranklength,    nu,        1           ) \
+    X( mjtNum,  actuator_gear,         nu,            6                    ) \
-    X( mjtNum,  actuator_acc0,           nu,        1           ) \
+    X( mjtNum,  actuator_cranklength,  nu,            1                    ) \
-    X( mjtNum,  actuator_length0,        nu,        1           ) \
+    X( mjtNum,  actuator_acc0,         nu,            1                    ) \
-    X( mjtNum,  actuator_lengthrange,    nu,        2           ) \
+    X( mjtNum,  actuator_length0,      nu,            1                    ) \
-    X( mjtNum,  actuator_user,           nu,        nuser_actuator) \
+    X( mjtNum,  actuator_lengthrange,  nu,            2                    ) \
-    X( int,     sensor_type,             nsensor,   1           ) \
+    X( mjtNum,  actuator_user,         nu,            MJ_M(nuser_actuator) ) \
-    X( int,     sensor_datatype,         nsensor,   1           ) \
+    X( int,     sensor_type,           nsensor,       1                    ) \
-    X( int,     sensor_needstage,        nsensor,   1           ) \
+    X( int,     sensor_datatype,       nsensor,       1                    ) \
-    X( int,     sensor_objtype,          nsensor,   1           ) \
+    X( int,     sensor_needstage,      nsensor,       1                    ) \
-    X( int,     sensor_objid,            nsensor,   1           ) \
+    X( int,     sensor_objtype,        nsensor,       1                    ) \
-    X( int,     sensor_dim,              nsensor,   1           ) \
+    X( int,     sensor_objid,          nsensor,       1                    ) \
-    X( int,     sensor_adr,              nsensor,   1           ) \
+    X( int,     sensor_reftype,        nsensor,       1                    ) \
-    X( mjtNum,  sensor_cutoff,           nsensor,   1           ) \
+    X( int,     sensor_refid,          nsensor,       1                    ) \
-    X( mjtNum,  sensor_noise,            nsensor,   1           ) \
+    X( int,     sensor_dim,            nsensor,       1                    ) \
-    X( mjtNum,  sensor_user,             nsensor,   nuser_sensor) \
+    X( int,     sensor_adr,            nsensor,       1                    ) \
-    X( int,     numeric_adr,             nnumeric,  1           ) \
+    X( mjtNum,  sensor_cutoff,         nsensor,       1                    ) \
-    X( int,     numeric_size,            nnumeric,  1           ) \
+    X( mjtNum,  sensor_noise,          nsensor,       1                    ) \
-    X( mjtNum,  numeric_data,            nnumericdata, 1        ) \
+    X( mjtNum,  sensor_user,           nsensor,       MJ_M(nuser_sensor)   ) \
-    X( int,     text_adr,                ntext,     1           ) \
+    X( int,     numeric_adr,           nnumeric,      1                    ) \
-    X( int,     text_size,               ntext,     1           ) \
+    X( int,     numeric_size,          nnumeric,      1                    ) \
-    X( char,    text_data,               ntextdata, 1           ) \
+    X( mjtNum,  numeric_data,          nnumericdata,  1                    ) \
-    X( int,     tuple_adr,               ntuple,    1           ) \
+    X( int,     text_adr,              ntext,         1                    ) \
-    X( int,     tuple_size,              ntuple,    1           ) \
+    X( int,     text_size,             ntext,         1                    ) \
-    X( int,     tuple_objtype,           ntupledata, 1          ) \
+    X( char,    text_data,             ntextdata,     1                    ) \
-    X( int,     tuple_objid,             ntupledata, 1          ) \
+    X( int,     tuple_adr,             ntuple,        1                    ) \
-    X( mjtNum,  tuple_objprm,            ntupledata, 1          ) \
+    X( int,     tuple_size,            ntuple,        1                    ) \
-    X( mjtNum,  key_time,                nkey,      1           ) \
+    X( int,     tuple_objtype,         ntupledata,    1                    ) \
-    X( mjtNum,  key_qpos,                nkey,      nq          ) \
+    X( int,     tuple_objid,           ntupledata,    1                    ) \
-    X( mjtNum,  key_qvel,                nkey,      nv          ) \
+    X( mjtNum,  tuple_objprm,          ntupledata,    1                    ) \
-    X( mjtNum,  key_act,                 nkey,      na          ) \
+    X( mjtNum,  key_time,              nkey,          1                    ) \
-    X( mjtNum,  key_mpos,                nkey,      nmocap3     ) \
+    X( mjtNum,  key_qpos,              nkey,          MJ_M(nq)             ) \
-    X( mjtNum,  key_mquat,               nkey,      nmocap4     ) \
+    X( mjtNum,  key_qvel,              nkey,          MJ_M(nv)             ) \
-    X( int,     name_bodyadr,            nbody,     1           ) \
+    X( mjtNum,  key_act,               nkey,          MJ_M(na)             ) \
-    X( int,     name_jntadr,             njnt,      1           ) \
+    X( mjtNum,  key_mpos,              nkey,          MJ_M(nmocap)*3       ) \
-    X( int,     name_geomadr,            ngeom,     1           ) \
+    X( mjtNum,  key_mquat,             nkey,          MJ_M(nmocap)*4       ) \
-    X( int,     name_siteadr,            nsite,     1           ) \
+    X( int,     name_bodyadr,          nbody,         1                    ) \
-    X( int,     name_camadr,             ncam,      1           ) \
+    X( int,     name_jntadr,           njnt,          1                    ) \
-    X( int,     name_lightadr,           nlight,    1           ) \
+    X( int,     name_geomadr,          ngeom,         1                    ) \
-    X( int,     name_meshadr,            nmesh,     1           ) \
+    X( int,     name_siteadr,          nsite,         1                    ) \
-    X( int,     name_skinadr,            nskin,     1           ) \
+    X( int,     name_camadr,           ncam,          1                    ) \
-    X( int,     name_hfieldadr,          nhfield,   1           ) \
+    X( int,     name_lightadr,         nlight,        1                    ) \
-    X( int,     name_texadr,             ntex,      1           ) \
+    X( int,     name_meshadr,          nmesh,         1                    ) \
-    X( int,     name_matadr,             nmat,      1           ) \
+    X( int,     name_skinadr,          nskin,         1                    ) \
-    X( int,     name_pairadr,            npair,     1           ) \
+    X( int,     name_hfieldadr,        nhfield,       1                    ) \
-    X( int,     name_excludeadr,         nexclude,  1           ) \
+    X( int,     name_texadr,           ntex,          1                    ) \
-    X( int,     name_eqadr,              neq,       1           ) \
+    X( int,     name_matadr,           nmat,          1                    ) \
-    X( int,     name_tendonadr,          ntendon,   1           ) \
+    X( int,     name_pairadr,          npair,         1                    ) \
-    X( int,     name_actuatoradr,        nu,        1           ) \
+    X( int,     name_excludeadr,       nexclude,      1                    ) \
-    X( int,     name_sensoradr,          nsensor,   1           ) \
+    X( int,     name_eqadr,            neq,           1                    ) \
-    X( int,     name_numericadr,         nnumeric,  1           ) \
+    X( int,     name_tendonadr,        ntendon,       1                    ) \
-    X( int,     name_textadr,            ntext,     1           ) \
+    X( int,     name_actuatoradr,      nu,            1                    ) \
-    X( int,     name_tupleadr,           ntuple,    1           ) \
+    X( int,     name_sensoradr,        nsensor,       1                    ) \
-    X( int,     name_keyadr,             nkey,      1           ) \
+    X( int,     name_numericadr,       nnumeric,      1                    ) \
-    X( char,    names,                   nnames,    1           )
+    X( int,     name_textadr,          ntext,         1                    ) \
+    X( int,     name_tupleadr,         ntuple,        1                    ) \
+    X( int,     name_keyadr,           nkey,          1                    ) \
+    X( char,    names,                 nnames,        1                    )
 
 
-
+//-------------------------------- mjData ----------------------------------------------------------
-//-------------------------------- mjData -----------------------------------------------
-
 
 // define symbols needed in MJDATA_POINTERS (corresponding to number of columns)
 #define MJDATA_POINTERS_PREAMBLE( m ) \
-    int nv = m->nv; \
+    int nv = m->nv;                   \
     int njmax = m->njmax;
 
+
 // pointer fields of mjData
-#define MJDATA_POINTERS                                   \
+#define MJDATA_POINTERS                                         \
-    X( mjtNum,  qpos,               nq,         1       ) \
+    X( mjtNum,    qpos,              nq,          1           ) \
-    X( mjtNum,  qvel,               nv,         1       ) \
+    X( mjtNum,    qvel,              nv,          1           ) \
-    X( mjtNum,  act,                na,         1       ) \
+    X( mjtNum,    act,               na,          1           ) \
-    X( mjtNum,  qacc_warmstart,     nv,         1       ) \
+    X( mjtNum,    qacc_warmstart,    nv,          1           ) \
-    X( mjtNum,  ctrl,               nu,         1       ) \
+    X( mjtNum,    ctrl,              nu,          1           ) \
-    X( mjtNum,  qfrc_applied,       nv,         1       ) \
+    X( mjtNum,    qfrc_applied,      nv,          1           ) \
-    X( mjtNum,  xfrc_applied,       nbody,      6       ) \
+    X( mjtNum,    xfrc_applied,      nbody,       6           ) \
-    X( mjtNum,  qacc,               nv,         1       ) \
+    X( mjtNum,    mocap_pos,         nmocap,      3           ) \
-    X( mjtNum,  act_dot,            na,         1       ) \
+    X( mjtNum,    mocap_quat,        nmocap,      4           ) \
-    X( mjtNum,  mocap_pos,          nmocap,     3       ) \
+    X( mjtNum,    qacc,              nv,          1           ) \
-    X( mjtNum,  mocap_quat,         nmocap,     4       ) \
+    X( mjtNum,    act_dot,           na,          1           ) \
-    X( mjtNum,  userdata,           nuserdata,  1       ) \
+    X( mjtNum,    userdata,          nuserdata,   1           ) \
-    X( mjtNum,  sensordata,         nsensordata,1       ) \
+    X( mjtNum,    sensordata,        nsensordata, 1           ) \
-    X( mjtNum,  xpos,               nbody,      3       ) \
+    X( mjtNum,    xpos,              nbody,       3           ) \
-    X( mjtNum,  xquat,              nbody,      4       ) \
+    X( mjtNum,    xquat,             nbody,       4           ) \
-    X( mjtNum,  xmat,               nbody,      9       ) \
+    X( mjtNum,    xmat,              nbody,       9           ) \
-    X( mjtNum,  xipos,              nbody,      3       ) \
+    X( mjtNum,    xipos,             nbody,       3           ) \
-    X( mjtNum,  ximat,              nbody,      9       ) \
+    X( mjtNum,    ximat,             nbody,       9           ) \
-    X( mjtNum,  xanchor,            njnt,       3       ) \
+    X( mjtNum,    xanchor,           njnt,        3           ) \
-    X( mjtNum,  xaxis,              njnt,       3       ) \
+    X( mjtNum,    xaxis,             njnt,        3           ) \
-    X( mjtNum,  geom_xpos,          ngeom,      3       ) \
+    X( mjtNum,    geom_xpos,         ngeom,       3           ) \
-    X( mjtNum,  geom_xmat,          ngeom,      9       ) \
+    X( mjtNum,    geom_xmat,         ngeom,       9           ) \
-    X( mjtNum,  site_xpos,          nsite,      3       ) \
+    X( mjtNum,    site_xpos,         nsite,       3           ) \
-    X( mjtNum,  site_xmat,          nsite,      9       ) \
+    X( mjtNum,    site_xmat,         nsite,       9           ) \
-    X( mjtNum,  cam_xpos,           ncam,       3       ) \
+    X( mjtNum,    cam_xpos,          ncam,        3           ) \
-    X( mjtNum,  cam_xmat,           ncam,       9       ) \
+    X( mjtNum,    cam_xmat,          ncam,        9           ) \
-    X( mjtNum,  light_xpos,         nlight,     3       ) \
+    X( mjtNum,    light_xpos,        nlight,      3           ) \
-    X( mjtNum,  light_xdir,         nlight,     3       ) \
+    X( mjtNum,    light_xdir,        nlight,      3           ) \
-    X( mjtNum,  subtree_com,        nbody,      3       ) \
+    X( mjtNum,    subtree_com,       nbody,       3           ) \
-    X( mjtNum,  cdof,               nv,         6       ) \
+    X( mjtNum,    cdof,              nv,          6           ) \
-    X( mjtNum,  cinert,             nbody,      10      ) \
+    X( mjtNum,    cinert,            nbody,       10          ) \
-    X( int,     ten_wrapadr,        ntendon,    1       ) \
+    X( int,       ten_wrapadr,       ntendon,     1           ) \
-    X( int,     ten_wrapnum,        ntendon,    1       ) \
+    X( int,       ten_wrapnum,       ntendon,     1           ) \
-    X( int,     ten_J_rownnz,       ntendon,    1       ) \
+    X( int,       ten_J_rownnz,      ntendon,     1           ) \
-    X( int,     ten_J_rowadr,       ntendon,    1       ) \
+    X( int,       ten_J_rowadr,      ntendon,     1           ) \
-    X( int,     ten_J_colind,       ntendon,    nv      ) \
+    X( int,       ten_J_colind,      ntendon,     MJ_M(nv)    ) \
-    X( mjtNum,  ten_length,         ntendon,    1       ) \
+    X( mjtNum,    ten_length,        ntendon,     1           ) \
-    X( mjtNum,  ten_J,              ntendon,    nv      ) \
+    X( mjtNum,    ten_J,             ntendon,     MJ_M(nv)    ) \
-    X( int,     wrap_obj,           nwrap,      2       ) \
+    X( int,       wrap_obj,          nwrap,       2           ) \
-    X( mjtNum,  wrap_xpos,          nwrap,      6       ) \
+    X( mjtNum,    wrap_xpos,         nwrap,       6           ) \
-    X( mjtNum,  actuator_length,    nu,         1       ) \
+    X( mjtNum,    actuator_length,   nu,          1           ) \
-    X( mjtNum,  actuator_moment,    nu,         nv      ) \
+    X( mjtNum,    actuator_moment,   nu,          MJ_M(nv)    ) \
-    X( mjtNum,  crb,                nbody,      10      ) \
+    X( mjtNum,    crb,               nbody,       10          ) \
-    X( mjtNum,  qM,                 nM,         1       ) \
+    X( mjtNum,    qM,                nM,          1           ) \
-    X( mjtNum,  qLD,                nM,         1       ) \
+    X( mjtNum,    qLD,               nM,          1           ) \
-    X( mjtNum,  qLDiagInv,          nv,         1       ) \
+    X( mjtNum,    qLDiagInv,         nv,          1           ) \
-    X( mjtNum,  qLDiagSqrtInv,      nv,         1       ) \
+    X( mjtNum,    qLDiagSqrtInv,     nv,          1           ) \
-    X( mjContact, contact,          nconmax,    1       ) \
+    X( mjContact, contact,           nconmax,     1           ) \
-    X( int,     efc_type,           njmax,      1       ) \
+    X( int,       efc_type,          njmax,       1           ) \
-    X( int,     efc_id,             njmax,      1       ) \
+    X( int,       efc_id,            njmax,       1           ) \
-    X( int,     efc_J_rownnz,       njmax,      1       ) \
+    X( int,       efc_J_rownnz,      njmax,       1           ) \
-    X( int,     efc_J_rowadr,       njmax,      1       ) \
+    X( int,       efc_J_rowadr,      njmax,       1           ) \
-    X( int,     efc_J_rowsuper,     njmax,      1       ) \
+    X( int,       efc_J_rowsuper,    njmax,       1           ) \
-    X( int,     efc_J_colind,       njmax,      nv      ) \
+    X( int,       efc_J_colind,      njmax,       MJ_M(nv)    ) \
-    X( int,     efc_JT_rownnz,      nv,         1       ) \
+    X( int,       efc_JT_rownnz,     nv,          1           ) \
-    X( int,     efc_JT_rowadr,      nv,         1       ) \
+    X( int,       efc_JT_rowadr,     nv,          1           ) \
-    X( int,     efc_JT_rowsuper,    nv,         1       ) \
+    X( int,       efc_JT_rowsuper,   nv,          1           ) \
-    X( int,     efc_JT_colind,      nv,         njmax   ) \
+    X( int,       efc_JT_colind,     nv,          MJ_M(njmax) ) \
-    X( mjtNum,  efc_J,              njmax,      nv      ) \
+    X( mjtNum,    efc_J,             njmax,       MJ_M(nv)    ) \
-    X( mjtNum,  efc_JT,             nv,         njmax   ) \
+    X( mjtNum,    efc_JT,            nv,          MJ_M(njmax) ) \
-    X( mjtNum,  efc_pos,            njmax,      1       ) \
+    X( mjtNum,    efc_pos,           njmax,       1           ) \
-    X( mjtNum,  efc_margin,         njmax,      1       ) \
+    X( mjtNum,    efc_margin,        njmax,       1           ) \
-    X( mjtNum,  efc_frictionloss,   njmax,      1       ) \
+    X( mjtNum,    efc_frictionloss,  njmax,       1           ) \
-    X( mjtNum,  efc_diagApprox,     njmax,      1       ) \
+    X( mjtNum,    efc_diagApprox,    njmax,       1           ) \
-    X( mjtNum,  efc_KBIP,           njmax,      4       ) \
+    X( mjtNum,    efc_KBIP,          njmax,       4           ) \
-    X( mjtNum,  efc_D,              njmax,      1       ) \
+    X( mjtNum,    efc_D,             njmax,       1           ) \
-    X( mjtNum,  efc_R,              njmax,      1       ) \
+    X( mjtNum,    efc_R,             njmax,       1           ) \
-    X( int,     efc_AR_rownnz,      njmax,      1       ) \
+    X( int,       efc_AR_rownnz,     njmax,       1           ) \
-    X( int,     efc_AR_rowadr,      njmax,      1       ) \
+    X( int,       efc_AR_rowadr,     njmax,       1           ) \
-    X( int,     efc_AR_colind,      njmax,      njmax   ) \
+    X( int,       efc_AR_colind,     njmax,       MJ_M(njmax) ) \
-    X( mjtNum,  efc_AR,             njmax,      njmax   ) \
+    X( mjtNum,    efc_AR,            njmax,       MJ_M(njmax) ) \
-    X( mjtNum,  ten_velocity,       ntendon,    1       ) \
+    X( mjtNum,    ten_velocity,      ntendon,     1           ) \
-    X( mjtNum,  actuator_velocity,  nu,         1       ) \
+    X( mjtNum,    actuator_velocity, nu,          1           ) \
-    X( mjtNum,  cvel,               nbody,      6       ) \
+    X( mjtNum,    cvel,              nbody,       6           ) \
-    X( mjtNum,  cdof_dot,           nv,         6       ) \
+    X( mjtNum,    cdof_dot,          nv,          6           ) \
-    X( mjtNum,  qfrc_bias,          nv,         1       ) \
+    X( mjtNum,    qfrc_bias,         nv,          1           ) \
-    X( mjtNum,  qfrc_passive,       nv,         1       ) \
+    X( mjtNum,    qfrc_passive,      nv,          1           ) \
-    X( mjtNum,  efc_vel,            njmax,      1       ) \
+    X( mjtNum,    efc_vel,           njmax,       1           ) \
-    X( mjtNum,  efc_aref,           njmax,      1       ) \
+    X( mjtNum,    efc_aref,          njmax,       1           ) \
-    X( mjtNum,  subtree_linvel,     nbody,      3       ) \
+    X( mjtNum,    subtree_linvel,    nbody,       3           ) \
-    X( mjtNum,  subtree_angmom,     nbody,      3       ) \
+    X( mjtNum,    subtree_angmom,    nbody,       3           ) \
-    X( mjtNum,  actuator_force,     nu,         1       ) \
+    X( mjtNum,    actuator_force,    nu,          1           ) \
-    X( mjtNum,  qfrc_actuator,      nv,         1       ) \
+    X( mjtNum,    qfrc_actuator,     nv,          1           ) \
-    X( mjtNum,  qfrc_unc,           nv,         1       ) \
+    X( mjtNum,    qfrc_unc,          nv,          1           ) \
-    X( mjtNum,  qacc_unc,           nv,         1       ) \
+    X( mjtNum,    qacc_unc,          nv,          1           ) \
-    X( mjtNum,  efc_b,              njmax,      1       ) \
+    X( mjtNum,    efc_b,             njmax,       1           ) \
-    X( mjtNum,  efc_force,          njmax,      1       ) \
+    X( mjtNum,    efc_force,         njmax,       1           ) \
-    X( int,     efc_state,          njmax,      1       ) \
+    X( int,       efc_state,         njmax,       1           ) \
-    X( mjtNum,  qfrc_constraint,    nv,         1       ) \
+    X( mjtNum,    qfrc_constraint,   nv,          1           ) \
-    X( mjtNum,  qfrc_inverse,       nv,         1       ) \
+    X( mjtNum,    qfrc_inverse,      nv,          1           ) \
-    X( mjtNum,  cacc,               nbody,      6       ) \
+    X( mjtNum,    cacc,              nbody,       6           ) \
-    X( mjtNum,  cfrc_int,           nbody,      6       ) \
+    X( mjtNum,    cfrc_int,          nbody,       6           ) \
-    X( mjtNum,  cfrc_ext,           nbody,      6       )
+    X( mjtNum,    cfrc_ext,          nbody,       6           )
 
 
 // scalar fields of mjData
@@ -539,4 +545,5 @@
     X( mjtNum,         solver_fwdinv,  2,           1 ) \
     X( mjtNum,         energy,         2,           1 )
 
+
 #endif  // MUJOCO_MJXMACRO_H_

include/mujoco.h

@@ -15,7 +15,7 @@
 #ifndef MUJOCO_MUJOCO_H_
 #define MUJOCO_MUJOCO_H_
 
-#include "mujoco_export.h"
+#include <mjexport.h>
 
 
 // this is a C-API
@@ -24,7 +24,7 @@ extern "C" {
 #endif
 
 // header version; should match the library version as returned by mj_version()
-#define mjVERSION_HEADER 211
+#define mjVERSION_HEADER 212
 
 // needed to define size_t, fabs and log10
 #include "stdlib.h"
@@ -32,11 +32,12 @@ extern "C" {
 
 
 // type definitions
-#include "mjmodel.h"
+#include <mjdata.h>
-#include "mjdata.h"
+#include <mjmodel.h>
-#include "mjvisualize.h"
+#include <mjrender.h>
-#include "mjrender.h"
+#include <mjtnum.h>
-#include "mjui.h"
+#include <mjui.h>
+#include <mjvisualize.h>
 
 
 // macros
@@ -186,9 +187,11 @@ MJAPI void mj_deleteModel(mjModel* m);
 MJAPI int mj_sizeModel(const mjModel* m);
 
 // 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);
 
 // Reset data to defaults.
@@ -220,13 +223,18 @@ MJAPI int mj_setLengthRange(mjModel* m, mjData* d, int index,
 //---------------------------------- Printing ------------------------------------------------------
 
 // Print mjModel to text file, specifying format.
-// float_format_str must be a valid printf-style format string for a single float value.
+// 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_str);
+                                  const char* float_format);
 
 // Print model 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);
 
@@ -358,7 +366,7 @@ MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d);
 // Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians.
 // If cost is not NULL, set *cost = s(jar) where jar = Jac*qacc-aref.
 MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
-                               mjtNum* cost, int flg_coneHessian);
+                               mjtNum cost[1], int flg_coneHessian);
 
 
 //---------------------------------- Support -------------------------------------------------------
@@ -429,19 +437,19 @@ MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
 
 // Apply cartesian force and torque (outside xfrc_applied mechanism).
 MJAPI void mj_applyFT(const mjModel* m, mjData* d,
-                      const mjtNum* force, const mjtNum* torque,
+                      const mjtNum force[3], const mjtNum torque[3],
-                      const mjtNum* point, int body, mjtNum* qfrc_target);
+                      const mjtNum point[3], int body, mjtNum* qfrc_target);
 
 // 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, int flg_local);
+                             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, int flg_local);
+                                 int objtype, int objid, mjtNum res[6], int flg_local);
 
 // Extract 6D force:torque given contact id, in the contact frame.
-MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum* result);
+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,
@@ -454,8 +462,8 @@ MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, m
 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, mjtNum* xmat,
+MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
-                           const mjtNum* pos, const mjtNum* quat,
+                           const mjtNum pos[3], const mjtNum quat[4],
                            int body, mjtByte sameframe);
 
 // Sum all body masses.
@@ -474,27 +482,28 @@ MJAPI const char* mj_versionString();
 //---------------------------------- Ray collisions ------------------------------------------------
 
 // Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude.
-// Return geomid and distance (x) to nearest surface, or -1 if no intersection.
+// Return distance (x) to nearest surface, or -1 if no intersection and output geomid.
 // 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,
+MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
                     const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
-                    int* geomid);
+                    int geomid[1]);
 
 // Interect ray with hfield, return nearest distance or -1 if no intersection.
 MJAPI mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
-                          const mjtNum* pnt, const mjtNum* vec);
+                          const mjtNum pnt[3], const mjtNum vec[3]);
 
 // Interect ray with mesh, return nearest distance or -1 if no intersection.
 MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
-                        const mjtNum* pnt, const mjtNum* vec);
+                        const mjtNum pnt[3], const mjtNum vec[3]);
 
 // Interect ray with pure geom, return nearest distance or -1 if no intersection.
-MJAPI mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
-                         const mjtNum* pnt, const mjtNum* vec, int geomtype);
+                         const mjtNum pnt[3], const mjtNum vec[3], int geomtype);
 
-// Interect ray with skin, return nearest vertex id.
+// Interect ray with skin, return nearest distance or -1 if no intersection,
+// and also output 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);
+                         const mjtNum pnt[3], const mjtNum vec[3], int vertid[1]);
 
 
 //---------------------------------- Interaction ---------------------------------------------------
@@ -506,26 +515,26 @@ MJAPI void mjv_defaultCamera(mjvCamera* cam);
 MJAPI void mjv_defaultPerturb(mjvPerturb* pert);
 
 // Transform pose from room to model space.
-MJAPI void mjv_room2model(mjtNum* modelpos, mjtNum* modelquat, const mjtNum* roompos,
+MJAPI void mjv_room2model(mjtNum modelpos[3], mjtNum modelquat[4], const mjtNum roompos[3],
-                          const mjtNum* roomquat, const mjvScene* scn);
+                          const mjtNum roomquat[4], const mjvScene* scn);
 
 // Transform pose from model to room space.
-MJAPI void mjv_model2room(mjtNum* roompos, mjtNum* roomquat, const mjtNum* modelpos,
+MJAPI void mjv_model2room(mjtNum roompos[3], mjtNum roomquat[4], const mjtNum modelpos[3],
-                          const mjtNum* modelquat, const mjvScene* scn);
+                          const mjtNum modelquat[4], 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,
+MJAPI void mjv_cameraInModel(mjtNum headpos[3], mjtNum forward[3], mjtNum up[3],
                              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,
+MJAPI void mjv_cameraInRoom(mjtNum headpos[3], mjtNum forward[3], mjtNum up[3],
                             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);
+MJAPI void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum forward[3]);
 
 // Move camera with mouse; action is mjtMouse.
 MJAPI void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
@@ -537,7 +546,7 @@ MJAPI void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum
 
 // Move model with mouse; action is mjtMouse.
 MJAPI void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
-                         const mjtNum* roomup, mjvScene* scn);
+                         const mjtNum roomup[3], mjvScene* scn);
 
 // Copy perturb pos,quat from selected body; set scale for perturbation.
 MJAPI void mjv_initPerturb(const mjModel* m, const mjData* d,
@@ -557,7 +566,7 @@ MJAPI mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera*
 // 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, int* geomid, int* skinid);
+                     const mjvScene* scn, mjtNum selpnt[3], int geomid[1], int skinid[1]);
 
 
 //---------------------------------- Visualization -------------------------------------------------
@@ -569,8 +578,8 @@ MJAPI void mjv_defaultOption(mjvOption* opt);
 MJAPI void mjv_defaultFigure(mjvFigure* fig);
 
 // Initialize given geom fields when not NULL, set the rest to their default values.
-MJAPI void mjv_initGeom(mjvGeom* geom, int type, const mjtNum* size,
+MJAPI void mjv_initGeom(mjvGeom* geom, int type, const mjtNum size[3],
-                        const mjtNum* pos, const mjtNum* mat, const float* rgba);
+                        const mjtNum pos[3], const mjtNum mat[9], const float rgba[4]);
 
 // Set (type, size, pos, mat) for connector-type geom between given points.
 // Assume that mjv_initGeom was already called to set all other properties.
@@ -994,7 +1003,7 @@ MJAPI void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[
 
 //--------------------------------- Decompositions -------------------------------------------------
 
-// Cholesky decomposition: mat = L*L'; return rank.
+// Cholesky decomposition: mat = L*L'; return rank, decomposition performed in-place into mat.
 MJAPI int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
 
 // Solve mat * res = vec, where mat is Cholesky-factorized
@@ -1004,7 +1013,7 @@ MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
 
 // Eigenvalue decomposition of symmetric 3x3 matrix.
-MJAPI int mju_eig3(mjtNum* eigval, mjtNum* eigvec, mjtNum* quat, const mjtNum* mat);
+MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
 
 
 //---------------------- Miscellaneous --------------------------------------------------

introspect/ast_nodes.py

@@ -0,0 +1,220 @@
+# 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.
+# ==============================================================================
+"""Classes that roughly correspond to Clang AST node types."""
+
+import collections
+import dataclasses
+import re
+from typing import Dict, Optional, Sequence, Tuple, Union
+
+# We are relying on Clang to do the actual source parsing and are only doing
+# a little bit of extra parsing of function parameter type declarations here.
+# These patterns are here for sanity checking rather than actual parsing.
+VALID_TYPE_NAME_PATTERN = re.compile('[A-Za-z_][A-Za-z0-9_]*')
+C_INVALID_TYPE_NAMES = frozenset([
+    'auto', 'break', 'case', 'const', 'continue', 'default', 'do', 'else',
+    'enum', 'extern', 'for', 'goto', 'if', 'inline', 'register', 'restrict',
+    'return', 'sizeof', 'static', 'struct', 'switch', 'typedef', 'union',
+    'volatile', 'while', '_Alignas', '_Atomic', '_Generic', '_Imaginary',
+    '_Noreturn', '_Static_assert', '_Thread_local', '__attribute__', '_Pragma'])
+
+
+def _is_valid_integral_type(type_str: str):
+  """Checks if a string is a valid integral type."""
+  parts = re.split(r'\s+', type_str)
+  counter = collections.defaultdict(lambda: 0)
+  wildcard_counter = 0
+  for part in parts:
+    if part in ('signed', 'unsigned', 'short', 'long', 'int', 'char'):
+      counter[part] += 1
+    elif VALID_TYPE_NAME_PATTERN.fullmatch(part):
+      # a non-keyword can be a typedef for int
+      wildcard_counter += 1
+    else:
+      return False
+
+  if (counter['signed'] + counter['unsigned'] > 1 or
+      counter['short'] > 1 or counter['long'] > 2 or
+      (counter['short'] and counter['long']) or
+      ((counter['short'] or counter['long']) and counter['char']) or
+      counter['char'] + counter['int'] + wildcard_counter > 1):
+    return False
+  else:
+    return True
+
+
+@dataclasses.dataclass
+class ValueType:
+  """Represents a C type that is neither a pointer type nor an array type."""
+
+  name: str
+  is_const: bool = False
+  is_volatile: bool = False
+
+  def __init__(self, name: str, is_const: bool = False,
+               is_volatile: bool = False):
+    is_valid_type_name = (
+        VALID_TYPE_NAME_PATTERN.fullmatch(name) or
+        _is_valid_integral_type(name)) and name not in C_INVALID_TYPE_NAMES
+    if not is_valid_type_name:
+      raise ValueError(f'{name!r} is not a valid value type name')
+    self.name = name
+    self.is_const = is_const
+    self.is_volatile = is_volatile
+
+  def decl(self, name_or_decl: Optional[str] = None) -> str:
+    parts = []
+    if self.is_const:
+      parts.append('const')
+    if self.is_volatile:
+      parts.append('volatile')
+    parts.append(self.name)
+    if name_or_decl:
+      parts.append(name_or_decl)
+    return ' '.join(parts)
+
+  def __str__(self):
+    return self.decl()
+
+
+@dataclasses.dataclass
+class ArrayType:
+  """Represents a C array type."""
+
+  inner_type: Union[ValueType, 'PointerType']
+  extents: Tuple[int]
+
+  def __init__(self, inner_type: Union[ValueType, 'PointerType'],
+               extents: Sequence[int]):
+    self.inner_type = inner_type
+    self.extents = tuple(extents)
+
+  @property
+  def _extents_str(self) -> str:
+    return ''.join(f'[{n}]' for n in self.extents)
+
+  def decl(self, name_or_decl: Optional[str] = None) -> str:
+    name_or_decl = name_or_decl or ''
+    return self.inner_type.decl(f'{name_or_decl}{self._extents_str}')
+
+  def __str__(self):
+    return self.decl()
+
+
+@dataclasses.dataclass
+class PointerType:
+  """Represents a C pointer type."""
+
+  inner_type: Union[ValueType, ArrayType, 'PointerType']
+  is_const: bool = False
+  is_volatile: bool = False
+  is_restrict: bool = False
+
+  def decl(self, name_or_decl: Optional[str] = None) -> str:
+    """Creates a string that declares an object of this type."""
+    parts = ['*']
+    if self.is_const:
+      parts.append('const')
+    if self.is_volatile:
+      parts.append('volatile')
+    if self.is_restrict:
+      parts.append('restrict')
+    if name_or_decl:
+      parts.append(name_or_decl)
+    ptr_decl = ' '.join(parts)
+    if isinstance(self.inner_type, ArrayType):
+      ptr_decl = f'({ptr_decl})'
+    return self.inner_type.decl(ptr_decl)
+
+  def __str__(self):
+    return self.decl()
+
+
+@dataclasses.dataclass
+class FunctionParameterDecl:
+  """Represents a parameter in a function declaration.
+
+  Note that according to the C language rule, a function parameter of array
+  type undergoes array-to-pointer decay, and therefore appears as a pointer
+  parameter in an actual C AST. We retain the arrayness of a parameter here
+  since the array's extents are informative.
+  """
+
+  name: str
+  type: Union[ValueType, ArrayType, PointerType]
+
+  def __str__(self):
+    return self.type.decl(self.name)
+
+  @property
+  def decltype(self) -> str:
+    return self.type.decl()
+
+
+@dataclasses.dataclass
+class FunctionDecl:
+  """Represents a function declaration."""
+
+  name: str
+  return_type: Union[ValueType, ArrayType, PointerType]
+  parameters: Tuple[FunctionParameterDecl]
+  doc: str
+
+  def __init__(self, name: str,
+               return_type: Union[ValueType, ArrayType, PointerType],
+               parameters: Sequence[FunctionParameterDecl],
+               doc: str):
+    self.name = name
+    self.return_type = return_type
+    self.parameters = tuple(parameters)
+    self.doc = doc
+
+  def __str__(self):
+    param_str = ', '.join(str(p) for p in self.parameters)
+    return f'{self.return_type} {self.name}({param_str})'
+
+  @property
+  def decltype(self) -> str:
+    param_str = ', '.join(str(p.decltype) for p in self.parameters)
+    return f'{self.return_type} ({param_str})'
+
+
+class _EnumDeclValues(Dict[str, int]):
+  """A dict with modified stringified representation.
+
+  The __repr__ method of this class adds a trailing comma to the list of values.
+  This is done as a hint for code formatters to place one item per line when
+  the stringified OrderedDict is used in generated Python code.
+  """
+
+  def __repr__(self):
+    out = super().__repr__()
+    if self:
+      out = re.sub(r'\(\[(.+)\]\)\Z', r'([\1,])', out)
+    return re.sub(r'\A_EnumDeclValues', 'dict', out)
+
+
+@dataclasses.dataclass
+class EnumDecl:
+  """Represents an enum declaration."""
+
+  name: str
+  declname: str
+  values: Dict[str, int]
+
+  def __init__(self, name: str, declname: str, values: Dict[str, int]):
+    self.name = name
+    self.declname = declname
+    self.values = _EnumDeclValues(values)

introspect/ast_nodes_test.py

@@ -0,0 +1,123 @@
+# 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.
+# ==============================================================================
+"""Tests for ast_nodes.py."""
+
+from absl.testing import absltest
+
+from google3.third_party.mujoco.introspect import ast_nodes
+
+
+class AstNodesTest(absltest.TestCase):
+
+  def test_value_type(self):
+    value_type = ast_nodes.ValueType('int')
+    self.assertEqual(str(value_type), 'int')
+    self.assertEqual(value_type.decl('var'), 'int var')
+
+    const_value_type = ast_nodes.ValueType('double', is_const=True)
+    self.assertEqual(str(const_value_type), 'const double')
+    self.assertEqual(const_value_type.decl('var2'), 'const double var2')
+
+  def test_pointer_type(self):
+    pointer_type = ast_nodes.PointerType(ast_nodes.ValueType('int'))
+    self.assertEqual(str(pointer_type), 'int *')
+    self.assertEqual(pointer_type.decl('var'), 'int * var')
+
+    const_pointer_type = ast_nodes.PointerType(
+        ast_nodes.ValueType('double'), is_const=True)
+    self.assertEqual(str(const_pointer_type), 'double * const')
+    self.assertEqual(const_pointer_type.decl('var2'), 'double * const var2')
+
+    pointer_to_const_type = ast_nodes.PointerType(
+        ast_nodes.ValueType('float', is_const=True))
+    self.assertEqual(str(pointer_to_const_type), 'const float *')
+    self.assertEqual(pointer_to_const_type.decl('var3'), 'const float * var3')
+
+    restrict_volatile_pointer_to_const_type = ast_nodes.PointerType(
+        ast_nodes.ValueType('char', is_const=True),
+        is_volatile=True, is_restrict=True)
+    self.assertEqual(str(restrict_volatile_pointer_to_const_type),
+                     'const char * volatile restrict')
+    self.assertEqual(
+        restrict_volatile_pointer_to_const_type.decl('var4'),
+        'const char * volatile restrict var4')
+
+    pointer_to_array_type = ast_nodes.PointerType(
+        ast_nodes.ArrayType(ast_nodes.ValueType('long'), (3,)))
+    self.assertEqual(str(pointer_to_array_type), 'long (*)[3]')
+    self.assertEqual(pointer_to_array_type.decl('var5'), 'long (* var5)[3]')
+
+    const_pointer_to_array_type = ast_nodes.PointerType(
+        ast_nodes.ArrayType(ast_nodes.ValueType('unsigned int'), (4,)),
+        is_const=True)
+    self.assertEqual(
+        str(const_pointer_to_array_type), 'unsigned int (* const)[4]')
+    self.assertEqual(
+        const_pointer_to_array_type.decl('var6'),
+        'unsigned int (* const var6)[4]')
+
+  def test_array_type(self):
+    array_type = ast_nodes.ArrayType(ast_nodes.ValueType('int'), (4,))
+    self.assertEqual(str(array_type), 'int [4]')
+    self.assertEqual(array_type.decl('var'), 'int var[4]')
+
+    array_2d_type = ast_nodes.ArrayType(
+        ast_nodes.ValueType('double', is_const=True), (2, 3))
+    self.assertEqual(str(array_2d_type), 'const double [2][3]')
+    self.assertEqual(array_2d_type.decl('var2'), 'const double var2[2][3]')
+
+    array_to_pointer_type = ast_nodes.ArrayType(
+        ast_nodes.PointerType(ast_nodes.ValueType('char', is_const=True)), (5,))
+    self.assertEqual(str(array_to_pointer_type), 'const char * [5]')
+    self.assertEqual(array_to_pointer_type.decl('var3'), 'const char * var3[5]')
+
+    array_to_const_pointer_type = ast_nodes.ArrayType(
+        ast_nodes.PointerType(ast_nodes.ValueType('float'), is_const=True),
+        (7,))
+    self.assertEqual(str(array_to_const_pointer_type), 'float * const [7]')
+    self.assertEqual(
+        array_to_const_pointer_type.decl('var4'), 'float * const var4[7]')
+
+  def test_complex_type(self):
+    complex_type = ast_nodes.ArrayType(
+        extents=[9],
+        inner_type=ast_nodes.PointerType(
+            ast_nodes.PointerType(
+                is_const=True,
+                inner_type=ast_nodes.ArrayType(
+                    extents=[7],
+                    inner_type=ast_nodes.PointerType(
+                        is_const=True,
+                        inner_type=ast_nodes.PointerType(
+                            ast_nodes.ArrayType(
+                                extents=(3, 4),
+                                inner_type=ast_nodes.ValueType(
+                                    'unsigned int', is_const=True)
+                                )
+                            )
+                        )
+                    )
+                )
+            )
+        )
+    self.assertEqual(str(complex_type),
+                     'const unsigned int (* * const (* const * [9])[7])[3][4]')
+    self.assertEqual(
+        complex_type.decl('var'),
+        'const unsigned int (* * const (* const * var[9])[7])[3][4]')
+
+
+if __name__ == '__main__':
+  absltest.main()

introspect/codegen/formatter.py

@@ -0,0 +1,149 @@
+# 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.
+# ==============================================================================
+"""Utility for formatting AST node as Python code."""
+
+import contextlib
+import dataclasses
+from typing import Any, Iterable, Mapping, Sequence
+
+INDENT_WIDTH = 4
+MAX_LINE_WIDTH = 80
+SIMPLE_TYPES = frozenset([int, float, str, bool, bytes, type(None)])
+
+
+def format_as_python_code(obj: Any) -> str:
+  """Formats an AST node object as well-indented Python code."""
+  formatter = _Formatter()
+  formatter.add(obj)
+  return str(formatter)
+
+
+def _is_all_simple(seq: Iterable[Any]) -> bool:
+  return all(type(obj) in SIMPLE_TYPES for obj in seq)
+
+
+class _Formatter:
+  """A helper for pretty-printing AST nodes as Python code."""
+
+  def __init__(self):
+    self._line_prefix = ''
+    self._lines = []
+    self._add_to_last_line = False
+
+  @contextlib.contextmanager
+  def _indent(self, width: int = INDENT_WIDTH):
+    self._line_prefix += ' ' * width
+    yield
+    self._line_prefix = self._line_prefix[:-width]
+
+  @contextlib.contextmanager
+  def _append_at_end(self, s):
+    yield
+    self._lines[-1] += s
+
+  def _add_line(self, line: str, no_break: bool = False):
+    if self._add_to_last_line:
+      self._lines[-1] += line
+    else:
+      self._lines.append(self._line_prefix + line)
+    self._add_to_last_line = no_break
+
+  def _add_dict(self, obj: Mapping[Any, Any]):
+    """Adds a dict to the formatted output."""
+    self._add_line('dict([')
+    with self._indent():
+      for k, v in obj.items():
+
+        # Try to fit everything into a single line first.
+        if _is_all_simple((k, v)):
+          single_line = f'({k!r}, {v!r}),'
+          if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
+            self._add_line(single_line)
+            continue
+
+        self._add_line(f"('{k}',")
+        with self._append_at_end('),'):
+          with self._indent(1):
+            self.add(v)
+
+    self._add_line('])')
+
+  def _add_dataclass(self, obj: Any):
+    """Adds a dataclass object to the formatted output."""
+    # Filter out default values.
+    kv_pairs = []
+    for k in dataclasses.fields(obj):
+      v = getattr(obj, k.name)
+      if v != k.default:
+        kv_pairs.append((k, v))
+
+    # Try to fit everything into a single line first.
+    if _is_all_simple(v for _, v in kv_pairs):
+      single_line = ', '.join(f'{k.name}={v!r}' for k, v in kv_pairs)
+      single_line = f'{obj.__class__.__name__}({single_line})'
+      if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
+        self._add_line(single_line)
+        return
+
+    self._add_line(obj.__class__.__name__ + '(')
+    with self._indent():
+      for k, v in kv_pairs:
+        self._add_line(k.name + '=', no_break=True)
+        with self._append_at_end(','):
+          self.add(v)
+    self._add_line(')')
+
+  def _add_sequence(self, obj: Sequence[Any]) -> None:
+    """Adds a sequence to the formatted output."""
+    default_str = repr(obj)
+    open_token, close_token = default_str[0], default_str[-1]
+    # Try to fit everything into a single line first.
+    if _is_all_simple(obj):
+      single_line = (
+          f"{open_token}{', '.join(repr(o) for o in obj)}{close_token}")
+      if close_token == ')' and len(obj) == 1:
+        single_line = f'{single_line[:-1]},)'
+      if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
+        self._add_line(single_line)
+        return
+
+    self._add_line(open_token)
+    with self._indent():
+      for v in obj:
+        with self._append_at_end(','):
+          self.add(v)
+    self._add_line(close_token)
+
+  def add(self, obj: Any) -> None:
+    """Adds an object to the formatted output."""
+    if _is_all_simple((obj,)):
+      self._add_line(repr(obj))
+    elif dataclasses.is_dataclass(obj):
+      self._add_dataclass(obj)
+    elif isinstance(obj, Mapping):
+      self._add_dict(obj)
+    elif isinstance(obj, Sequence):
+      self._add_sequence(obj)
+    else:
+      raise NotImplementedError
+
+  def __str__(self):
+    lines = []
+    for line in self._lines:
+      if len(line) > MAX_LINE_WIDTH:
+        lines.append(f'{line}  # pylint: disable=line-too-long')
+      else:
+        lines.append(line)
+    return '\n'.join(lines)

introspect/codegen/generate_enums.py

@@ -0,0 +1,126 @@
+# 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.
+# ==============================================================================
+"""Generates enums.py.
+
+The JSON input can be generated via:
+    clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments mujoco.h
+"""
+
+import json
+from typing import Any, Mapping, Sequence
+
+from absl import app
+from absl import flags
+
+from google3.third_party.mujoco.introspect import ast_nodes
+from google3.third_party.mujoco.introspect.codegen import formatter
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    'json_path', None,
+    'Path to the JSON file representing the Clang AST for mujoco.h')
+
+ClangJsonNode = Mapping[str, Any]
+
+
+def traverse(node, visitor):
+  visitor.visit(node)
+  children = node.get('inner', [])
+  for child in children:
+    traverse(child, visitor)
+
+
+class MjEnumVisitor:
+  """A Clang AST JSON node visitor for MuJoCo API enum declarations."""
+
+  def __init__(self):
+    self._enums = {}
+    self._typedefs = {}
+
+  def _make_enum(self, node: ClangJsonNode) -> ast_nodes.EnumDecl:
+    """Makes a EnumDecl from a Clang AST EnumDecl node."""
+    name = f"enum {node['name']}"
+    values = []
+    for child in node['inner']:
+      child_kind = child.get('kind')
+      if child_kind == 'EnumConstantDecl':
+        next_idx = values[-1][1] + 1 if values else 0
+        value = int(child['inner'][0].get('value', next_idx))
+        values.append((child['name'], value))
+    return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values))
+
+  def visit(self, node: ClangJsonNode) -> None:
+    if (node.get('kind') == 'EnumDecl' and
+        node.get('name', '').startswith('mj')):
+      enum_decl = self._make_enum(node)
+      self._enums[enum_decl.name] = enum_decl
+    elif (node.get('kind') == 'TypedefDecl' and
+          node['type']['qualType'].startswith('enum mj')):
+      enum = self._enums[node['type']['qualType']]
+      self._typedefs[node['name']] = ast_nodes.EnumDecl(
+          name=node['name'], declname=enum.declname, values=dict(enum.values))
+
+  @property
+  def enums(self) -> Mapping[str, ast_nodes.EnumDecl]:
+    return self._enums
+
+  @property
+  def typedefs(self) -> Mapping[str, ast_nodes.EnumDecl]:
+    return self._typedefs
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  with open(FLAGS.json_path, 'r', encoding='utf-8') as f:
+    root = json.load(f)
+
+  visitor = MjEnumVisitor()
+
+  traverse(root, visitor)
+
+  enums_str = formatter.format_as_python_code(visitor.typedefs)
+
+  print(f'''
+# 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.
+# ==============================================================================
+"""Provides information about MuJoCo API enums.
+
+DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+"""
+
+from typing import Mapping
+
+from google3.third_party.mujoco.introspect.ast_nodes import EnumDecl
+
+ENUMS: Mapping[str, EnumDecl] = {enums_str}
+'''.strip())  # `print` adds a trailing newline
+
+
+if __name__ == '__main__':
+  app.run(main)

introspect/codegen/generate_functions.py

@@ -0,0 +1,159 @@
+# 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.
+# ==============================================================================
+"""Generates functions.py.
+
+The JSON input can be generated via:
+    clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments mujoco.h
+"""
+
+import json
+from typing import Any, Mapping, Sequence
+
+from absl import app
+from absl import flags
+
+from google3.third_party.mujoco.introspect import ast_nodes
+from google3.third_party.mujoco.introspect import type_parsing
+from google3.third_party.mujoco.introspect.codegen import formatter
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('header_path', None, 'Path to the original mujoco.h')
+flags.DEFINE_string(
+    'json_path', None,
+    'Path to the JSON file representing the Clang AST for mujoco.h')
+
+ClangJsonNode = Mapping[str, Any]
+
+
+def traverse(node, visitor):
+  visitor.visit(node)
+  children = node.get('inner', [])
+  for child in children:
+    traverse(child, visitor)
+
+
+class MjFunctionVisitor:
+  """A Clang AST JSON node visitor for MuJoCo API function declarations."""
+
+  def __init__(self, raw_header):
+    self._raw_header = raw_header
+    self._functions = {}
+
+  def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl:
+    """Makes a FunctionDecl from a Clang AST FunctionDecl node."""
+    name = node['name']
+    return_type = type_parsing.parse_function_return_type(
+        node['type']['qualType'])
+    parameters = []
+    comments = []
+    for child in node['inner']:
+      child_kind = child.get('kind')
+      if child_kind == 'ParmVarDecl':
+        parameters.append(self._make_parameter(child))
+      if child_kind == 'FullComment':
+        comments.append(self._make_comment(child))
+    comment = ' '.join(comments).strip()
+    return ast_nodes.FunctionDecl(
+        name=name, return_type=return_type, parameters=parameters, doc=comment)
+
+  def _make_parameter(
+      self, node: ClangJsonNode) -> ast_nodes.FunctionParameterDecl:
+    """Makes a ParameterDecl from a Clang AST ParmVarDecl node."""
+    name = node['name']
+    type_name = node['type']['qualType']
+
+    # For a pointer parameters, look up in the original header to see if
+    # n array extent was declared there.
+    if type_name.endswith('*'):
+      decl_begin = node['range']['begin']['offset']
+      decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen']
+      decl = self._raw_header[decl_begin:decl_end]
+      name_begin = node['loc']['offset'] - decl_begin
+      name_end = name_begin + node['loc']['tokLen']
+      type_name = decl[:name_begin] + decl[name_end:]
+
+    return ast_nodes.FunctionParameterDecl(
+        name=name, type=type_parsing.parse_type(type_name))
+
+  def _make_comment(self, node: ClangJsonNode) -> str:
+    """Makes a comment string from a Clang AST FullComment node."""
+    kind = node.get('kind')
+    if kind == 'TextComment':
+      return node['text'].replace('\N{NO-BREAK SPACE}', ' ')
+    else:
+      strings = []
+      for child in node['inner']:
+        strings.append(self._make_comment(child))
+      return ''.join(strings)
+
+  def visit(self, node: ClangJsonNode) -> None:
+    if (node.get('kind') == 'FunctionDecl' and
+        node.get('name', '').startswith('mj')):
+      func_decl = self._make_function(node)
+      self._functions[func_decl.name] = func_decl
+
+  @property
+  def functions(self) -> Mapping[str, ast_nodes.FunctionDecl]:
+    return self._functions
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  with open(FLAGS.json_path, 'r', encoding='utf-8') as f:
+    root = json.load(f)
+
+  with open(FLAGS.header_path, 'r') as f:
+    visitor = MjFunctionVisitor(f.read())
+
+  traverse(root, visitor)
+
+  functions_str = formatter.format_as_python_code(visitor.functions)
+
+  print(f'''
+# 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.
+# ==============================================================================
+"""Provides information about MuJoCo API functions.
+
+DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+"""
+
+from typing import Mapping
+
+from google3.third_party.mujoco.introspect.ast_nodes import ArrayType
+from google3.third_party.mujoco.introspect.ast_nodes import FunctionDecl
+from google3.third_party.mujoco.introspect.ast_nodes import FunctionParameterDecl
+from google3.third_party.mujoco.introspect.ast_nodes import PointerType
+from google3.third_party.mujoco.introspect.ast_nodes import ValueType
+
+FUNCTIONS: Mapping[str, FunctionDecl] = {functions_str}
+'''.strip())  # `print` adds a trailing newline
+
+
+if __name__ == '__main__':
+  app.run(main)

introspect/enums.py

@@ -0,0 +1,633 @@
+# 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.
+# ==============================================================================
+"""Provides information about MuJoCo API enums.
+
+DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+"""
+
+from typing import Mapping
+
+from .ast_nodes import EnumDecl
+
+ENUMS: Mapping[str, EnumDecl] = dict([
+    ('mjtDisableBit',
+     EnumDecl(
+         name='mjtDisableBit',
+         declname='enum mjtDisableBit_',
+         values=dict([
+             ('mjDSBL_CONSTRAINT', 1),
+             ('mjDSBL_EQUALITY', 2),
+             ('mjDSBL_FRICTIONLOSS', 4),
+             ('mjDSBL_LIMIT', 8),
+             ('mjDSBL_CONTACT', 16),
+             ('mjDSBL_PASSIVE', 32),
+             ('mjDSBL_GRAVITY', 64),
+             ('mjDSBL_CLAMPCTRL', 128),
+             ('mjDSBL_WARMSTART', 256),
+             ('mjDSBL_FILTERPARENT', 512),
+             ('mjDSBL_ACTUATION', 1024),
+             ('mjDSBL_REFSAFE', 2048),
+             ('mjNDISABLE', 12),
+         ]),
+     )),
+    ('mjtEnableBit',
+     EnumDecl(
+         name='mjtEnableBit',
+         declname='enum mjtEnableBit_',
+         values=dict([
+             ('mjENBL_OVERRIDE', 1),
+             ('mjENBL_ENERGY', 2),
+             ('mjENBL_FWDINV', 4),
+             ('mjENBL_SENSORNOISE', 8),
+             ('mjNENABLE', 4),
+         ]),
+     )),
+    ('mjtJoint',
+     EnumDecl(
+         name='mjtJoint',
+         declname='enum mjtJoint_',
+         values=dict([
+             ('mjJNT_FREE', 0),
+             ('mjJNT_BALL', 1),
+             ('mjJNT_SLIDE', 2),
+             ('mjJNT_HINGE', 3),
+         ]),
+     )),
+    ('mjtGeom',
+     EnumDecl(
+         name='mjtGeom',
+         declname='enum mjtGeom_',
+         values=dict([
+             ('mjGEOM_PLANE', 0),
+             ('mjGEOM_HFIELD', 1),
+             ('mjGEOM_SPHERE', 2),
+             ('mjGEOM_CAPSULE', 3),
+             ('mjGEOM_ELLIPSOID', 4),
+             ('mjGEOM_CYLINDER', 5),
+             ('mjGEOM_BOX', 6),
+             ('mjGEOM_MESH', 7),
+             ('mjNGEOMTYPES', 8),
+             ('mjGEOM_ARROW', 100),
+             ('mjGEOM_ARROW1', 101),
+             ('mjGEOM_ARROW2', 102),
+             ('mjGEOM_LINE', 103),
+             ('mjGEOM_SKIN', 104),
+             ('mjGEOM_LABEL', 105),
+             ('mjGEOM_NONE', 1001),
+         ]),
+     )),
+    ('mjtCamLight',
+     EnumDecl(
+         name='mjtCamLight',
+         declname='enum mjtCamLight_',
+         values=dict([
+             ('mjCAMLIGHT_FIXED', 0),
+             ('mjCAMLIGHT_TRACK', 1),
+             ('mjCAMLIGHT_TRACKCOM', 2),
+             ('mjCAMLIGHT_TARGETBODY', 3),
+             ('mjCAMLIGHT_TARGETBODYCOM', 4),
+         ]),
+     )),
+    ('mjtTexture',
+     EnumDecl(
+         name='mjtTexture',
+         declname='enum mjtTexture_',
+         values=dict([
+             ('mjTEXTURE_2D', 0),
+             ('mjTEXTURE_CUBE', 1),
+             ('mjTEXTURE_SKYBOX', 2),
+         ]),
+     )),
+    ('mjtIntegrator',
+     EnumDecl(
+         name='mjtIntegrator',
+         declname='enum mjtIntegrator_',
+         values=dict([
+             ('mjINT_EULER', 0),
+             ('mjINT_RK4', 1),
+         ]),
+     )),
+    ('mjtCollision',
+     EnumDecl(
+         name='mjtCollision',
+         declname='enum mjtCollision_',
+         values=dict([
+             ('mjCOL_ALL', 0),
+             ('mjCOL_PAIR', 1),
+             ('mjCOL_DYNAMIC', 2),
+         ]),
+     )),
+    ('mjtCone',
+     EnumDecl(
+         name='mjtCone',
+         declname='enum mjtCone_',
+         values=dict([
+             ('mjCONE_PYRAMIDAL', 0),
+             ('mjCONE_ELLIPTIC', 1),
+         ]),
+     )),
+    ('mjtJacobian',
+     EnumDecl(
+         name='mjtJacobian',
+         declname='enum mjtJacobian_',
+         values=dict([
+             ('mjJAC_DENSE', 0),
+             ('mjJAC_SPARSE', 1),
+             ('mjJAC_AUTO', 2),
+         ]),
+     )),
+    ('mjtSolver',
+     EnumDecl(
+         name='mjtSolver',
+         declname='enum mjtSolver_',
+         values=dict([
+             ('mjSOL_PGS', 0),
+             ('mjSOL_CG', 1),
+             ('mjSOL_NEWTON', 2),
+         ]),
+     )),
+    ('mjtEq',
+     EnumDecl(
+         name='mjtEq',
+         declname='enum mjtEq_',
+         values=dict([
+             ('mjEQ_CONNECT', 0),
+             ('mjEQ_WELD', 1),
+             ('mjEQ_JOINT', 2),
+             ('mjEQ_TENDON', 3),
+             ('mjEQ_DISTANCE', 4),
+         ]),
+     )),
+    ('mjtWrap',
+     EnumDecl(
+         name='mjtWrap',
+         declname='enum mjtWrap_',
+         values=dict([
+             ('mjWRAP_NONE', 0),
+             ('mjWRAP_JOINT', 1),
+             ('mjWRAP_PULLEY', 2),
+             ('mjWRAP_SITE', 3),
+             ('mjWRAP_SPHERE', 4),
+             ('mjWRAP_CYLINDER', 5),
+         ]),
+     )),
+    ('mjtTrn',
+     EnumDecl(
+         name='mjtTrn',
+         declname='enum mjtTrn_',
+         values=dict([
+             ('mjTRN_JOINT', 0),
+             ('mjTRN_JOINTINPARENT', 1),
+             ('mjTRN_SLIDERCRANK', 2),
+             ('mjTRN_TENDON', 3),
+             ('mjTRN_SITE', 4),
+             ('mjTRN_UNDEFINED', 1000),
+         ]),
+     )),
+    ('mjtDyn',
+     EnumDecl(
+         name='mjtDyn',
+         declname='enum mjtDyn_',
+         values=dict([
+             ('mjDYN_NONE', 0),
+             ('mjDYN_INTEGRATOR', 1),
+             ('mjDYN_FILTER', 2),
+             ('mjDYN_MUSCLE', 3),
+             ('mjDYN_USER', 4),
+         ]),
+     )),
+    ('mjtGain',
+     EnumDecl(
+         name='mjtGain',
+         declname='enum mjtGain_',
+         values=dict([
+             ('mjGAIN_FIXED', 0),
+             ('mjGAIN_MUSCLE', 1),
+             ('mjGAIN_USER', 2),
+         ]),
+     )),
+    ('mjtBias',
+     EnumDecl(
+         name='mjtBias',
+         declname='enum mjtBias_',
+         values=dict([
+             ('mjBIAS_NONE', 0),
+             ('mjBIAS_AFFINE', 1),
+             ('mjBIAS_MUSCLE', 2),
+             ('mjBIAS_USER', 3),
+         ]),
+     )),
+    ('mjtObj',
+     EnumDecl(
+         name='mjtObj',
+         declname='enum mjtObj_',
+         values=dict([
+             ('mjOBJ_UNKNOWN', 0),
+             ('mjOBJ_BODY', 1),
+             ('mjOBJ_XBODY', 2),
+             ('mjOBJ_JOINT', 3),
+             ('mjOBJ_DOF', 4),
+             ('mjOBJ_GEOM', 5),
+             ('mjOBJ_SITE', 6),
+             ('mjOBJ_CAMERA', 7),
+             ('mjOBJ_LIGHT', 8),
+             ('mjOBJ_MESH', 9),
+             ('mjOBJ_SKIN', 10),
+             ('mjOBJ_HFIELD', 11),
+             ('mjOBJ_TEXTURE', 12),
+             ('mjOBJ_MATERIAL', 13),
+             ('mjOBJ_PAIR', 14),
+             ('mjOBJ_EXCLUDE', 15),
+             ('mjOBJ_EQUALITY', 16),
+             ('mjOBJ_TENDON', 17),
+             ('mjOBJ_ACTUATOR', 18),
+             ('mjOBJ_SENSOR', 19),
+             ('mjOBJ_NUMERIC', 20),
+             ('mjOBJ_TEXT', 21),
+             ('mjOBJ_TUPLE', 22),
+             ('mjOBJ_KEY', 23),
+         ]),
+     )),
+    ('mjtConstraint',
+     EnumDecl(
+         name='mjtConstraint',
+         declname='enum mjtConstraint_',
+         values=dict([
+             ('mjCNSTR_EQUALITY', 0),
+             ('mjCNSTR_FRICTION_DOF', 1),
+             ('mjCNSTR_FRICTION_TENDON', 2),
+             ('mjCNSTR_LIMIT_JOINT', 3),
+             ('mjCNSTR_LIMIT_TENDON', 4),
+             ('mjCNSTR_CONTACT_FRICTIONLESS', 5),
+             ('mjCNSTR_CONTACT_PYRAMIDAL', 6),
+             ('mjCNSTR_CONTACT_ELLIPTIC', 7),
+         ]),
+     )),
+    ('mjtConstraintState',
+     EnumDecl(
+         name='mjtConstraintState',
+         declname='enum mjtConstraintState_',
+         values=dict([
+             ('mjCNSTRSTATE_SATISFIED', 0),
+             ('mjCNSTRSTATE_QUADRATIC', 1),
+             ('mjCNSTRSTATE_LINEARNEG', 2),
+             ('mjCNSTRSTATE_LINEARPOS', 3),
+             ('mjCNSTRSTATE_CONE', 4),
+         ]),
+     )),
+    ('mjtSensor',
+     EnumDecl(
+         name='mjtSensor',
+         declname='enum mjtSensor_',
+         values=dict([
+             ('mjSENS_TOUCH', 0),
+             ('mjSENS_ACCELEROMETER', 1),
+             ('mjSENS_VELOCIMETER', 2),
+             ('mjSENS_GYRO', 3),
+             ('mjSENS_FORCE', 4),
+             ('mjSENS_TORQUE', 5),
+             ('mjSENS_MAGNETOMETER', 6),
+             ('mjSENS_RANGEFINDER', 7),
+             ('mjSENS_JOINTPOS', 8),
+             ('mjSENS_JOINTVEL', 9),
+             ('mjSENS_TENDONPOS', 10),
+             ('mjSENS_TENDONVEL', 11),
+             ('mjSENS_ACTUATORPOS', 12),
+             ('mjSENS_ACTUATORVEL', 13),
+             ('mjSENS_ACTUATORFRC', 14),
+             ('mjSENS_BALLQUAT', 15),
+             ('mjSENS_BALLANGVEL', 16),
+             ('mjSENS_JOINTLIMITPOS', 17),
+             ('mjSENS_JOINTLIMITVEL', 18),
+             ('mjSENS_JOINTLIMITFRC', 19),
+             ('mjSENS_TENDONLIMITPOS', 20),
+             ('mjSENS_TENDONLIMITVEL', 21),
+             ('mjSENS_TENDONLIMITFRC', 22),
+             ('mjSENS_FRAMEPOS', 23),
+             ('mjSENS_FRAMEQUAT', 24),
+             ('mjSENS_FRAMEXAXIS', 25),
+             ('mjSENS_FRAMEYAXIS', 26),
+             ('mjSENS_FRAMEZAXIS', 27),
+             ('mjSENS_FRAMELINVEL', 28),
+             ('mjSENS_FRAMEANGVEL', 29),
+             ('mjSENS_FRAMELINACC', 30),
+             ('mjSENS_FRAMEANGACC', 31),
+             ('mjSENS_SUBTREECOM', 32),
+             ('mjSENS_SUBTREELINVEL', 33),
+             ('mjSENS_SUBTREEANGMOM', 34),
+             ('mjSENS_USER', 35),
+         ]),
+     )),
+    ('mjtStage',
+     EnumDecl(
+         name='mjtStage',
+         declname='enum mjtStage_',
+         values=dict([
+             ('mjSTAGE_NONE', 0),
+             ('mjSTAGE_POS', 1),
+             ('mjSTAGE_VEL', 2),
+             ('mjSTAGE_ACC', 3),
+         ]),
+     )),
+    ('mjtDataType',
+     EnumDecl(
+         name='mjtDataType',
+         declname='enum mjtDataType_',
+         values=dict([
+             ('mjDATATYPE_REAL', 0),
+             ('mjDATATYPE_POSITIVE', 1),
+             ('mjDATATYPE_AXIS', 2),
+             ('mjDATATYPE_QUATERNION', 3),
+         ]),
+     )),
+    ('mjtLRMode',
+     EnumDecl(
+         name='mjtLRMode',
+         declname='enum mjtLRMode_',
+         values=dict([
+             ('mjLRMODE_NONE', 0),
+             ('mjLRMODE_MUSCLE', 1),
+             ('mjLRMODE_MUSCLEUSER', 2),
+             ('mjLRMODE_ALL', 3),
+         ]),
+     )),
+    ('mjtWarning',
+     EnumDecl(
+         name='mjtWarning',
+         declname='enum mjtWarning_',
+         values=dict([
+             ('mjWARN_INERTIA', 0),
+             ('mjWARN_CONTACTFULL', 1),
+             ('mjWARN_CNSTRFULL', 2),
+             ('mjWARN_VGEOMFULL', 3),
+             ('mjWARN_BADQPOS', 4),
+             ('mjWARN_BADQVEL', 5),
+             ('mjWARN_BADQACC', 6),
+             ('mjWARN_BADCTRL', 7),
+             ('mjNWARNING', 8),
+         ]),
+     )),
+    ('mjtTimer',
+     EnumDecl(
+         name='mjtTimer',
+         declname='enum mjtTimer_',
+         values=dict([
+             ('mjTIMER_STEP', 0),
+             ('mjTIMER_FORWARD', 1),
+             ('mjTIMER_INVERSE', 2),
+             ('mjTIMER_POSITION', 3),
+             ('mjTIMER_VELOCITY', 4),
+             ('mjTIMER_ACTUATION', 5),
+             ('mjTIMER_ACCELERATION', 6),
+             ('mjTIMER_CONSTRAINT', 7),
+             ('mjTIMER_POS_KINEMATICS', 8),
+             ('mjTIMER_POS_INERTIA', 9),
+             ('mjTIMER_POS_COLLISION', 10),
+             ('mjTIMER_POS_MAKE', 11),
+             ('mjTIMER_POS_PROJECT', 12),
+             ('mjNTIMER', 13),
+         ]),
+     )),
+    ('mjtGridPos',
+     EnumDecl(
+         name='mjtGridPos',
+         declname='enum mjtGridPos_',
+         values=dict([
+             ('mjGRID_TOPLEFT', 0),
+             ('mjGRID_TOPRIGHT', 1),
+             ('mjGRID_BOTTOMLEFT', 2),
+             ('mjGRID_BOTTOMRIGHT', 3),
+         ]),
+     )),
+    ('mjtFramebuffer',
+     EnumDecl(
+         name='mjtFramebuffer',
+         declname='enum mjtFramebuffer_',
+         values=dict([
+             ('mjFB_WINDOW', 0),
+             ('mjFB_OFFSCREEN', 1),
+         ]),
+     )),
+    ('mjtFontScale',
+     EnumDecl(
+         name='mjtFontScale',
+         declname='enum mjtFontScale_',
+         values=dict([
+             ('mjFONTSCALE_50', 50),
+             ('mjFONTSCALE_100', 100),
+             ('mjFONTSCALE_150', 150),
+             ('mjFONTSCALE_200', 200),
+             ('mjFONTSCALE_250', 250),
+             ('mjFONTSCALE_300', 300),
+         ]),
+     )),
+    ('mjtFont',
+     EnumDecl(
+         name='mjtFont',
+         declname='enum mjtFont_',
+         values=dict([
+             ('mjFONT_NORMAL', 0),
+             ('mjFONT_SHADOW', 1),
+             ('mjFONT_BIG', 2),
+         ]),
+     )),
+    ('mjtButton',
+     EnumDecl(
+         name='mjtButton',
+         declname='enum mjtButton_',
+         values=dict([
+             ('mjBUTTON_NONE', 0),
+             ('mjBUTTON_LEFT', 1),
+             ('mjBUTTON_RIGHT', 2),
+             ('mjBUTTON_MIDDLE', 3),
+         ]),
+     )),
+    ('mjtEvent',
+     EnumDecl(
+         name='mjtEvent',
+         declname='enum mjtEvent_',
+         values=dict([
+             ('mjEVENT_NONE', 0),
+             ('mjEVENT_MOVE', 1),
+             ('mjEVENT_PRESS', 2),
+             ('mjEVENT_RELEASE', 3),
+             ('mjEVENT_SCROLL', 4),
+             ('mjEVENT_KEY', 5),
+             ('mjEVENT_RESIZE', 6),
+         ]),
+     )),
+    ('mjtItem',
+     EnumDecl(
+         name='mjtItem',
+         declname='enum mjtItem_',
+         values=dict([
+             ('mjITEM_END', -2),
+             ('mjITEM_SECTION', -1),
+             ('mjITEM_SEPARATOR', 0),
+             ('mjITEM_STATIC', 1),
+             ('mjITEM_BUTTON', 2),
+             ('mjITEM_CHECKINT', 3),
+             ('mjITEM_CHECKBYTE', 4),
+             ('mjITEM_RADIO', 5),
+             ('mjITEM_RADIOLINE', 6),
+             ('mjITEM_SELECT', 7),
+             ('mjITEM_SLIDERINT', 8),
+             ('mjITEM_SLIDERNUM', 9),
+             ('mjITEM_EDITINT', 10),
+             ('mjITEM_EDITNUM', 11),
+             ('mjITEM_EDITTXT', 12),
+             ('mjNITEM', 13),
+         ]),
+     )),
+    ('mjtCatBit',
+     EnumDecl(
+         name='mjtCatBit',
+         declname='enum mjtCatBit_',
+         values=dict([
+             ('mjCAT_STATIC', 1),
+             ('mjCAT_DYNAMIC', 2),
+             ('mjCAT_DECOR', 4),
+             ('mjCAT_ALL', 7),
+         ]),
+     )),
+    ('mjtMouse',
+     EnumDecl(
+         name='mjtMouse',
+         declname='enum mjtMouse_',
+         values=dict([
+             ('mjMOUSE_NONE', 0),
+             ('mjMOUSE_ROTATE_V', 1),
+             ('mjMOUSE_ROTATE_H', 2),
+             ('mjMOUSE_MOVE_V', 3),
+             ('mjMOUSE_MOVE_H', 4),
+             ('mjMOUSE_ZOOM', 5),
+             ('mjMOUSE_SELECT', 6),
+         ]),
+     )),
+    ('mjtPertBit',
+     EnumDecl(
+         name='mjtPertBit',
+         declname='enum mjtPertBit_',
+         values=dict([
+             ('mjPERT_TRANSLATE', 1),
+             ('mjPERT_ROTATE', 2),
+         ]),
+     )),
+    ('mjtCamera',
+     EnumDecl(
+         name='mjtCamera',
+         declname='enum mjtCamera_',
+         values=dict([
+             ('mjCAMERA_FREE', 0),
+             ('mjCAMERA_TRACKING', 1),
+             ('mjCAMERA_FIXED', 2),
+             ('mjCAMERA_USER', 3),
+         ]),
+     )),
+    ('mjtLabel',
+     EnumDecl(
+         name='mjtLabel',
+         declname='enum mjtLabel_',
+         values=dict([
+             ('mjLABEL_NONE', 0),
+             ('mjLABEL_BODY', 1),
+             ('mjLABEL_JOINT', 2),
+             ('mjLABEL_GEOM', 3),
+             ('mjLABEL_SITE', 4),
+             ('mjLABEL_CAMERA', 5),
+             ('mjLABEL_LIGHT', 6),
+             ('mjLABEL_TENDON', 7),
+             ('mjLABEL_ACTUATOR', 8),
+             ('mjLABEL_CONSTRAINT', 9),
+             ('mjLABEL_SKIN', 10),
+             ('mjLABEL_SELECTION', 11),
+             ('mjLABEL_SELPNT', 12),
+             ('mjLABEL_CONTACTFORCE', 13),
+             ('mjNLABEL', 14),
+         ]),
+     )),
+    ('mjtFrame',
+     EnumDecl(
+         name='mjtFrame',
+         declname='enum mjtFrame_',
+         values=dict([
+             ('mjFRAME_NONE', 0),
+             ('mjFRAME_BODY', 1),
+             ('mjFRAME_GEOM', 2),
+             ('mjFRAME_SITE', 3),
+             ('mjFRAME_CAMERA', 4),
+             ('mjFRAME_LIGHT', 5),
+             ('mjFRAME_WORLD', 6),
+             ('mjNFRAME', 7),
+         ]),
+     )),
+    ('mjtVisFlag',
+     EnumDecl(
+         name='mjtVisFlag',
+         declname='enum mjtVisFlag_',
+         values=dict([
+             ('mjVIS_CONVEXHULL', 0),
+             ('mjVIS_TEXTURE', 1),
+             ('mjVIS_JOINT', 2),
+             ('mjVIS_ACTUATOR', 3),
+             ('mjVIS_CAMERA', 4),
+             ('mjVIS_LIGHT', 5),
+             ('mjVIS_TENDON', 6),
+             ('mjVIS_RANGEFINDER', 7),
+             ('mjVIS_CONSTRAINT', 8),
+             ('mjVIS_INERTIA', 9),
+             ('mjVIS_SCLINERTIA', 10),
+             ('mjVIS_PERTFORCE', 11),
+             ('mjVIS_PERTOBJ', 12),
+             ('mjVIS_CONTACTPOINT', 13),
+             ('mjVIS_CONTACTFORCE', 14),
+             ('mjVIS_CONTACTSPLIT', 15),
+             ('mjVIS_TRANSPARENT', 16),
+             ('mjVIS_AUTOCONNECT', 17),
+             ('mjVIS_COM', 18),
+             ('mjVIS_SELECT', 19),
+             ('mjVIS_STATIC', 20),
+             ('mjVIS_SKIN', 21),
+             ('mjNVISFLAG', 22),
+         ]),
+     )),
+    ('mjtRndFlag',
+     EnumDecl(
+         name='mjtRndFlag',
+         declname='enum mjtRndFlag_',
+         values=dict([
+             ('mjRND_SHADOW', 0),
+             ('mjRND_WIREFRAME', 1),
+             ('mjRND_REFLECTION', 2),
+             ('mjRND_ADDITIVE', 3),
+             ('mjRND_SKYBOX', 4),
+             ('mjRND_FOG', 5),
+             ('mjRND_HAZE', 6),
+             ('mjRND_SEGMENT', 7),
+             ('mjRND_IDCOLOR', 8),
+             ('mjNRNDFLAG', 9),
+         ]),
+     )),
+    ('mjtStereo',
+     EnumDecl(
+         name='mjtStereo',
+         declname='enum mjtStereo_',
+         values=dict([
+             ('mjSTEREO_NONE', 0),
+             ('mjSTEREO_QUADBUFFERED', 1),
+             ('mjSTEREO_SIDEBYSIDE', 2),
+         ]),
+     )),
+])

introspect/enums_test.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+"""Tests for enums.py."""
+
+from absl.testing import absltest
+
+from google3.third_party.mujoco.introspect import enums
+
+
+class EnumsTest(absltest.TestCase):
+
+  # "simple" enum that just starts at zero and increment by one.
+  def test_mjtJoint(self):  # pylint: disable=invalid-name
+    enum_decl = enums.ENUMS['mjtJoint']
+    self.assertEqual(enum_decl.name, 'mjtJoint')
+    self.assertEqual(enum_decl.declname, 'enum mjtJoint_')
+    self.assertEqual(
+        tuple(enum_decl.values.items()), (('mjJNT_FREE', 0),
+                                          ('mjJNT_BALL', 1),
+                                          ('mjJNT_SLIDE', 2),
+                                          ('mjJNT_HINGE', 3)))
+
+  # all values explicitly specified
+  def test_mjtEnableBit(self):  # pylint: disable=invalid-name
+    enum_decl = enums.ENUMS['mjtEnableBit']
+    self.assertEqual(enum_decl.name, 'mjtEnableBit')
+    self.assertEqual(enum_decl.declname, 'enum mjtEnableBit_')
+    self.assertEqual(
+        tuple(enum_decl.values.items()), (('mjENBL_OVERRIDE', 1<<0),
+                                          ('mjENBL_ENERGY', 1<<1),
+                                          ('mjENBL_FWDINV', 1<<2),
+                                          ('mjENBL_SENSORNOISE', 1<<3),
+                                          ('mjNENABLE', 4)))
+
+  # values mostly increment by one with occasional overrides
+  def test_mjtGeom(self):  # pylint: disable=invalid-name
+    enum_decl = enums.ENUMS['mjtGeom']
+    self.assertEqual(enum_decl.name, 'mjtGeom')
+    self.assertEqual(enum_decl.declname, 'enum mjtGeom_')
+
+    self.assertEqual(enum_decl.values['mjGEOM_PLANE'], 0)
+    self.assertEqual(enum_decl.values['mjGEOM_HFIELD'], 1)
+    self.assertEqual(enum_decl.values['mjGEOM_SPHERE'], 2)
+    # Skip a few...
+
+    self.assertEqual(enum_decl.values['mjGEOM_ARROW'], 100)
+    self.assertEqual(enum_decl.values['mjGEOM_ARROW1'], 101)
+    self.assertEqual(enum_decl.values['mjGEOM_ARROW2'], 102)
+    # Skip a few...
+
+    self.assertEqual(enum_decl.values['mjGEOM_NONE'], 1001)
+
+if __name__ == '__main__':
+  absltest.main()

introspect/functions_test.py

@@ -0,0 +1,68 @@
+# 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.
+# ==============================================================================
+"""Tests for functions.py."""
+
+from absl.testing import absltest
+
+from google3.third_party.mujoco.introspect import ast_nodes
+from google3.third_party.mujoco.introspect import functions
+from google3.third_party.mujoco.introspect import type_parsing
+
+
+class FunctionsTest(absltest.TestCase):
+
+  def test_mj_copyData(self):  # pylint: disable=invalid-name
+    func_decl = functions.FUNCTIONS['mj_copyData']
+    self.assertEqual(func_decl.name, 'mj_copyData')
+    self.assertEqual(func_decl.return_type, type_parsing.parse_type('mjData*'))
+    self.assertEqual(
+        func_decl.parameters,
+        (ast_nodes.FunctionParameterDecl(
+            name='dest', type=type_parsing.parse_type('mjData*')),
+         ast_nodes.FunctionParameterDecl(
+             name='m', type=type_parsing.parse_type('const mjModel*')),
+         ast_nodes.FunctionParameterDecl(
+             name='src', type=type_parsing.parse_type('const mjData*'))))
+    self.assertEqual(
+        func_decl.doc, 'Copy mjData. '
+        'm is only required to contain the size fields from MJMODEL_INTS.')
+
+  def test_mju_transformSpatial(self):  # pylint: disable=invalid-name
+    func_decl = functions.FUNCTIONS['mju_transformSpatial']
+    self.assertEqual(func_decl.name, 'mju_transformSpatial')
+    self.assertEqual(func_decl.return_type, type_parsing.parse_type('void'))
+    self.assertEqual(
+        func_decl.parameters,
+        (ast_nodes.FunctionParameterDecl(
+            name='res', type=type_parsing.parse_type('mjtNum[6]')),
+         ast_nodes.FunctionParameterDecl(
+             name='vec', type=type_parsing.parse_type('const mjtNum[6]')),
+         ast_nodes.FunctionParameterDecl(
+             name='flg_force', type=type_parsing.parse_type('int')),
+         ast_nodes.FunctionParameterDecl(
+             name='newpos', type=type_parsing.parse_type('const mjtNum[3]')),
+         ast_nodes.FunctionParameterDecl(
+             name='oldpos', type=type_parsing.parse_type('const mjtNum[3]')),
+         ast_nodes.FunctionParameterDecl(
+             name='rotnew2old',
+             type=type_parsing.parse_type('const mjtNum[9]'))))
+    self.assertEqual(
+        func_decl.doc, '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.')
+
+
+if __name__ == '__main__':
+  absltest.main()

introspect/type_parsing.py

@@ -0,0 +1,148 @@
+# 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.
+# ==============================================================================
+"""Functions for parsing C type declarations."""
+
+import collections
+import re
+from typing import Mapping, MutableSequence, Optional, Sequence, Tuple, Union
+
+from google3.third_party.mujoco.introspect import ast_nodes
+
+ARRAY_EXTENTS_PATTERN = re.compile(r'(\[[^\]]+\]\s*)+\Z')
+ARRAY_N_PATTERN = re.compile(r'\[([^\]]+)\]')
+STARTS_WITH_CONST_PATTERN = re.compile(r'\Aconst(?![A-Za-z0-9_])')
+ENDS_WITH_CONST_PATTERN = re.compile(r'(?<![A-Za-z0-9_])const\Z')
+
+
+def _parse_qualifiers(
+    type_name: str,
+    qualifiers: Sequence[str]) -> Tuple[str, Mapping[str, bool]]:
+  """Separates qualifiers from the rest of the type name."""
+  parts = re.split(r'\s+', type_name)
+  counter = collections.defaultdict(lambda: 0)
+  non_qualifiers = []
+  for part in parts:
+    if part in qualifiers:
+      counter[part] += 1
+      if counter[part] > 1:
+        ValueError('duplicate qualifier: {part!r}')
+    else:
+      non_qualifiers.append(part)
+  is_qualifier = dict()
+  for qualifier in qualifiers:
+    is_qualifier[f'is_{qualifier}'] = bool(counter[qualifier])
+  return ' '.join(non_qualifiers), is_qualifier
+
+
+def _parse_maybe_array(
+    type_name: str, innermost_type: Optional[Union[ast_nodes.ValueType,
+                                                   ast_nodes.PointerType]]
+) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
+  """Internal-only helper that parses a type that may be an array type."""
+  array_match = ARRAY_EXTENTS_PATTERN.search(type_name)
+  if array_match:
+    extents = tuple(
+        int(s.strip()) for s in ARRAY_N_PATTERN.findall(array_match.group(0)))
+    inner_type_str = type_name[:array_match.start()]
+    return ast_nodes.ArrayType(
+        inner_type=_parse_maybe_pointer(inner_type_str.strip(), innermost_type),
+        extents=extents)
+  else:
+    return _parse_maybe_pointer(type_name, innermost_type)
+
+
+def _parse_maybe_pointer(
+    type_name: str, innermost_type: Optional[Union[ast_nodes.ValueType,
+                                                   ast_nodes.PointerType]]
+) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
+  """Internal-only helper that parses a type that may be a pointer type."""
+  p = type_name.rfind('*')
+  if p != -1:
+    leftover, is_qualifier = _parse_qualifiers(
+        type_name[p + 1:].strip(), ('const', 'volatile', 'restrict'))
+    if leftover:
+      raise ValueError('invalid qualifier for pointer: {leftover!r}')
+
+    inner_type_str = type_name[:p].strip()
+    if inner_type_str:
+      inner_type = _parse_maybe_pointer(inner_type_str, innermost_type)
+    else:
+      assert innermost_type is not None
+      inner_type = innermost_type
+    return ast_nodes.PointerType(inner_type=inner_type, **is_qualifier)
+  else:
+    assert innermost_type is None  # value type should be innermost
+    type_name, is_qualifier = _parse_qualifiers(
+        type_name.strip(), ('const', 'volatile'))
+    return ast_nodes.ValueType(name=type_name, **is_qualifier)
+
+
+def _peel_nested_parens(input_str: str) -> MutableSequence[str]:
+  """Extracts substrings from a string with nested parentheses.
+
+  The returned sequence starts from the substring enclosed in the innermost
+  parentheses and moves subsequently outwards. The contents of the inner
+  substrings are removed from the outer ones. For example, given the string
+  'lorem ipsum(dolor sit (consectetur adipiscing) amet)sed do eiusmod',
+  this function produces the sequence
+  ['consectetur adipiscing', 'dolor sit  amet', 'lorem ipsumsed do eiusmod'].
+
+  Args:
+    input_str: An input_str string consisting of zero or more nested
+      parentheses.
+
+  Returns:
+    A sequence of substrings enclosed with in respective parentheses. See the
+    description above for the precise detail of the output.
+  """
+  start = input_str.find('(')
+  end = input_str.rfind(')')
+
+  if start == -1 and end == -1:
+    return [input_str]
+  else:
+    # Assertions to be re-raised into a meaningful error by the caller.
+    assert start != -1  # '(' should be present if there is a ')'
+    assert end != -1  # ')' should be present if there is a '('
+    assert start < end  # '(' should come before ')'
+    out = _peel_nested_parens(input_str[start + 1:end])
+    out.append(input_str[:start] + input_str[end + 1:])
+    return out
+
+
+def parse_type(
+    type_name: str
+) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
+  """Parses a string that represents a C type into an AST node."""
+  try:
+    type_str_stack = _peel_nested_parens(type_name.strip())
+  except AssertionError:
+    raise ValueError(f'{type_name!r} contains incorrectly nested parentheses')
+
+  result = None
+  while type_str_stack:
+    try:
+      result = _parse_maybe_array(type_str_stack.pop(), result)
+    except AssertionError:
+      raise ValueError(f'invalid type name {type_name!r}')
+
+  assert result  # hint for pytype that `result` isn't None
+  return result
+
+
+def parse_function_return_type(
+    type_name: str
+) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
+  return parse_type(type_name[:type_name.rfind('(')])

introspect/type_parsing_test.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""Tests for type_parsing.py."""
+
+from absl.testing import absltest
+
+from google3.third_party.mujoco.introspect import ast_nodes
+from google3.third_party.mujoco.introspect import type_parsing
+
+
+class TypeParsingTest(absltest.TestCase):
+
+  def test_parse_complex_type(self):
+    parsed_type = type_parsing.parse_type(
+        'int unsigned volatile long const long'+
+        '(**const(*const restrict*[9])[7])[3][4]')
+
+    expected_type = ast_nodes.ArrayType(
+        extents=[9],
+        inner_type=ast_nodes.PointerType(
+            ast_nodes.PointerType(
+                is_const=True,
+                is_restrict=True,
+                inner_type=ast_nodes.ArrayType(
+                    extents=[7],
+                    inner_type=ast_nodes.PointerType(
+                        is_const=True,
+                        inner_type=ast_nodes.PointerType(
+                            ast_nodes.ArrayType(
+                                extents=(3, 4),
+                                inner_type=ast_nodes.ValueType(
+                                    'int unsigned long long',
+                                    is_const=True, is_volatile=True)
+                                )
+                            )
+                        )
+                    )
+                )
+            )
+        )
+
+    self.assertEqual(parsed_type, expected_type)
+
+
+if __name__ == '__main__':
+  absltest.main()

python/MANIFEST.in

@@ -0,0 +1,3 @@
+include LICENSE *.md
+recursive-include mujoco *.h *.cc CMakeLists.txt
+recursive-include cmake *.cmake

python/README.md

@@ -0,0 +1,62 @@
+# MuJoCo Python Bindings
+
+This package is the canonical Python bindings for the
+[MuJoCo physics engine](https://github.com/deepmind/mujoco).
+These bindings are developed and maintained by DeepMind, and is kept up-to-date
+with the latest developments in MuJoCo itself.
+
+The `mujoco` package provides direct access to raw MuJoCo C API functions,
+structs, constants, and enumerations. Structs are provided as Python classes,
+with Pythonic initialization and deletion semantics.
+
+It is not the aim of this package to provide fully fledged
+scene/environment/game authoring API, as there are already a number of existing
+packages that do this well. However, this package does provide a number of
+lower-level components outside of MuJoCo itself that are likely to be useful to
+most users who access MuJoCo through Python. For example, the `egl`, `glfw`, and
+`osmesa` subpackages contain utilities for setting up OpenGL rendering contexts.
+
+## Installation
+
+The package can be installed from [PyPI](https://pypi.org/project/mujoco/) via
+
+```sh
+pip install mujoco
+```
+
+A copy of the MuJoCo library is provided as part of the package and does **not**
+need to be downloaded or installed separately.
+
+If you wish to modify and build the bindings from source, you should clone the
+entire `mujoco` repository from GitHub, then run the `make_sdist.sh` script to
+generate a [source distribution](https://packaging.python.org/en/latest/glossary/#term-Source-Distribution-or-sdist)
+tarball, then run `pip wheel name_of_sdist.tar.gz`. The `make_sdist.sh` script
+generates additional C++ header files that are needed to build the bindings,
+and also pulls in required files from elsewhere in the repository outside the
+`python` directory into the sdist.
+
+CMake and a C++17 compiler are needed to build the bindings from source.
+
+## Usage
+
+Once installed, the package can be imported via `import mujoco`. Please consult
+our [documentation](https://mujoco.readthedocs.io/en/latest/python.html) for
+further detail on the package's API.
+
+## Versioning
+
+The `major.minor.micro` portion of the version number matches the version of
+MuJoCo that the bindings provide. Optionally, if we release updates to the
+Python bindings themselves that target the same version of MuJoCo, a `-rN`
+suffix is added, for example `2.1.2-r2` represents the second update to the
+bindings for MuJoCo 2.1.2.
+
+## License and Disclaimer
+
+Copyright 2022 DeepMind Technologies Limited
+
+MuJoCo and its Python bindings are licensed under the Apache License,
+Version 2.0. You may obtain a copy of the License at
+https://www.apache.org/licenses/LICENSE-2.0.
+
+This is not an officially supported Google product.

python/make_sdist.sh

@@ -0,0 +1,61 @@
+#!/bin/bash -xe
+# 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
+#
+#     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.
+
+if [[ -z ${VIRTUAL_ENV} ]]; then
+  echo "This script must be run from within a Python virtual environment"
+  exit 1
+fi
+
+# Figure out the path to this script (https://stackoverflow.com/a/246128).
+package_dir="$(cd -- "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
+
+if [[ "$(uname)" == CYGWIN* ]]; then
+  package_dir="$(cygpath -m ${package_dir})"
+  readonly tmp_dir="$(TMPDIR="${LOCALAPPDATA//\\/$'/'}/Temp" mktemp -d)"
+else
+  readonly tmp_dir="$(mktemp -d)"
+fi
+
+python -m pip install --upgrade pip
+python -m pip install absl-py
+pushd ${tmp_dir}
+cp -r "${package_dir}"/* .
+
+# Generate header files.
+old_pythonpath="${PYTHONPATH}"
+if [[ "$(uname)" == CYGWIN* ]]; then
+  export PYTHONPATH="${old_pythonpath};${package_dir}/.."
+else
+  export PYTHONPATH="${old_pythonpath}:${package_dir}/.."
+fi
+python "${package_dir}"/mujoco/codegen/generate_enum_traits.py > \
+    mujoco/enum_traits.h
+python "${package_dir}"/mujoco/codegen/generate_function_traits.py > \
+    mujoco/function_traits.h
+export PYTHONPATH="${old_pythonpath}"
+
+# Copy over the LICENSE file.
+cp "${package_dir}"/../../LICENSE .
+
+# Copy over CMake scripts.
+mkdir cmake
+cp "${package_dir}"/../../cmake/*.cmake cmake
+
+python setup.py sdist --formats=gztar
+tar -tf dist/mujoco-*.tar.gz
+popd
+
+mkdir -p "${package_dir}"/dist
+mv "${tmp_dir}"/dist/* "${package_dir}"/dist

python/mujoco/CMakeLists.txt

@@ -0,0 +1,381 @@
+# 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
+#
+#     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.
+
+project(mujoco_python)
+cmake_minimum_required(VERSION 3.15)
+
+# Support new IN_LIST if() operator.
+set(CMAKE_POLICY_DEFAULT_CMP0057 NEW)
+# INTERPROCEDURAL_OPTIMIZATION is enforced when enabled.
+set(CMAKE_POLICY_DEFAULT_CMP0069 NEW)
+
+enable_language(C)
+enable_language(CXX)
+
+if(MSVC AND MSVC_VERSION GREATER_EQUAL 1927)
+  set(CMAKE_CXX_STANDARD 20) # For forceinline lambdas.
+else()
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_C_VISIBILITY_PRESET hidden)
+set(CMAKE_CXX_VISIBILITY_PRESET hidden)
+set(CMAKE_VISIBILITY_INLINES_HIDDEN ON)
+
+# TODO(stunya) Figure out why this is required.
+separate_arguments(CMDLINE_LINK_OPTIONS UNIX_COMMAND ${CMAKE_SHARED_LINKER_FLAGS})
+add_link_options(${CMDLINE_LINK_OPTIONS})
+
+include(MujocoLinkOptions)
+get_mujoco_extra_link_options(EXTRA_LINK_OPTIONS)
+add_link_options(${EXTRA_LINK_OPTIONS})
+
+if(APPLE)
+  add_compile_options(-Werror=partial-availability -Werror=unguarded-availability)
+  add_link_options(-Wl,-no_weak_imports)
+endif()
+
+find_package(Python3 COMPONENTS Interpreter Development)
+
+include(FindOrFetch)
+
+# ==================== MUJOCO LIBRARY ==========================================
+if(NOT TARGET mujoco)
+  find_library(MUJOCO_LIBRARY mujoco mujoco.2.1.2 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED)
+  find_path(MUJOCO_INCLUDE mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
+  message("MuJoCo is at ${MUJOCO_LIBRARY}")
+  message("MuJoCo headers are at ${MUJOCO_INCLUDE}")
+  add_library(mujoco SHARED IMPORTED)
+  if(WIN32)
+    set_target_properties(mujoco PROPERTIES IMPORTED_IMPLIB "${MUJOCO_LIBRARY}")
+  else()
+    set_target_properties(mujoco PROPERTIES IMPORTED_LOCATION "${MUJOCO_LIBRARY}")
+  endif()
+  target_include_directories(mujoco INTERFACE "${MUJOCO_INCLUDE}")
+  if(APPLE)
+    execute_process(
+      COMMAND otool -XD ${MUJOCO_LIBRARY}
+      COMMAND head -n 1
+      COMMAND xargs dirname
+      COMMAND xargs echo -n
+      OUTPUT_VARIABLE MUJOCO_INSTALL_NAME_DIR
+    )
+    set_target_properties(mujoco PROPERTIES INSTALL_NAME_DIR "${MUJOCO_INSTALL_NAME_DIR}")
+  elseif(UNIX)
+    execute_process(
+      COMMAND objdump -p ${MUJOCO_LIBRARY}
+      COMMAND grep SONAME
+      COMMAND grep -Po [^\\s]+$
+      COMMAND xargs echo -n
+      OUTPUT_VARIABLE MUJOCO_SONAME
+    )
+    set_target_properties(mujoco PROPERTIES IMPORTED_SONAME "${MUJOCO_SONAME}")
+  endif()
+endif()
+
+# ==================== ABSEIL ==================================================
+if(APPLE)
+  set(ABSL_EXTRA_FETCH_ARGS
+      PATCH_COMMAND
+      "sed"
+      "-i"
+      " "
+      "s/-march=armv8-a+crypto/-mcpu=apple-m1+crypto/g"
+      "${CMAKE_BINARY_DIR}/_deps/abseil-cpp-src/absl/copts/GENERATED_AbseilCopts.cmake"
+  )
+else()
+  set(ABSL_EXTRA_FETCH_ARGS "")
+endif()
+findorfetch(
+  USE_SYSTEM_PACKAGE
+  OFF
+  PACKAGE_NAME
+  absl
+  LIBRARY_NAME
+  abseil-cpp
+  GIT_REPO
+  https://github.com/abseil/abseil-cpp
+  GIT_TAG
+  215105818dfde3174fe799600bb0f3cae233d0bf # 20211102.0
+  TARGETS
+  absl::core_headers
+  absl::flat_hash_map
+  absl::span
+  ${ABSL_EXTRA_FETCH_ARGS}
+  EXCLUDE_FROM_ALL
+)
+
+# ==================== EIGEN ===================================================
+add_compile_definitions(EIGEN_MPL2_ONLY)
+findorfetch(
+  USE_SYSTEM_PACKAGE
+  OFF
+  PACKAGE_NAME
+  Eigen3
+  LIBRARY_NAME
+  eigen
+  GIT_REPO
+  https://gitlab.com/libeigen/eigen
+  GIT_TAG
+  3147391d946bb4b6c68edd901f2add6ac1f31f8c # 3.4.0
+  TARGETS
+  Eigen3::Eigen
+  EXCLUDE_FROM_ALL
+)
+
+# ==================== PYBIND11 ================================================
+if(MUJOCO_PYBIND11_DIR)
+  FetchContent_Declare(pybind11 SOURCE_DIR ${MUJOCO_PYBIND11_DIR} EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(pybind11)
+else()
+  findorfetch(
+    USE_SYSTEM_PACKAGE
+    OFF
+    PACKAGE_NAME
+    pybind11
+    LIBRARY_NAME
+    pybind11
+    GIT_REPO
+    https://github.com/pybind/pybind11
+    GIT_TAG
+    4c4b33f14215f8992e2c67ace2a7a4c114c48bb9
+    TARGETS
+    pybind11::pybind11_headers
+    EXCLUDE_FROM_ALL
+  )
+endif()
+
+# ==================== MUJOCO PYTHON BINDINGS ==================================
+
+add_subdirectory(util)
+
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/enum_traits.h)
+  add_library(enum_traits INTERFACE enum_traits.h)
+else()
+  add_custom_command(
+    OUTPUT enum_traits.h
+    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${mujoco_SOURCE_DIR}/mujoco ${Python3_EXECUTABLE}
+            ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_enum_traits.py > enum_traits.h
+    DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_enum_traits.py
+  )
+  add_library(enum_traits INTERFACE enum_traits.h)
+  target_include_directories(
+    enum_traits INTERFACE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}
+  )
+endif()
+target_link_libraries(enum_traits INTERFACE mujoco absl::core_headers)
+
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/function_traits.h)
+  add_library(function_traits INTERFACE function_traits.h)
+else()
+  add_custom_command(
+    OUTPUT function_traits.h
+    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${mujoco_SOURCE_DIR}/mujoco ${Python3_EXECUTABLE}
+            ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_function_traits.py > function_traits.h
+    DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_function_traits.py
+  )
+  add_library(function_traits INTERFACE function_traits.h)
+  target_include_directories(
+    function_traits INTERFACE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}
+  )
+endif()
+target_link_libraries(function_traits INTERFACE mujoco absl::core_headers)
+
+add_library(errors_header INTERFACE errors.h)
+set_target_properties(errors_header PROPERTIES PUBLIC_HEADER errors.h)
+target_link_libraries(errors_header INTERFACE crossplatform func_wrap mujoco)
+
+add_library(raw INTERFACE raw.h)
+set_target_properties(raw PROPERTIES PUBLIC_HEADER raw.h)
+target_link_libraries(raw INTERFACE mujoco)
+
+add_library(
+  structs_header
+  INTERFACE
+  indexer_xmacro.h
+  indexers.h
+  mjdata_meta.h
+  structs.h
+)
+set_target_properties(structs_header PROPERTIES PUBLIC_HEADER structs.h)
+target_link_libraries(
+  structs_header
+  INTERFACE absl::flat_hash_map
+            absl::span
+            mujoco
+            raw
+)
+
+add_library(functions_header INTERFACE functions.h)
+set_target_properties(functions_header PROPERTIES PUBLIC_HEADER functions.h)
+target_link_libraries(
+  functions_header
+  INTERFACE array_traits
+            crossplatform
+            Eigen3::Eigen
+            errors_header
+            func_wrap
+            structs_header
+            tuple_tools
+)
+
+include(CheckAvxSupport)
+get_avx_compile_options(AVX_COMPILE_OPTIONS)
+
+macro(mujoco_pybind11_module name)
+  pybind11_add_module(${name} ${ARGN})
+  target_compile_options(${name} PRIVATE ${AVX_COMPILE_OPTIONS})
+  set_target_properties(${name} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+  if(APPLE)
+    add_custom_command(
+      TARGET ${name}
+      POST_BUILD
+      COMMAND
+        install_name_tool -change
+        $<TARGET_PROPERTY:mujoco,INSTALL_NAME_DIR>/$<TARGET_FILE_NAME:mujoco>
+        @rpath/$<TARGET_FILE_NAME:mujoco> -add_rpath @loader_path $<TARGET_FILE:${name}>
+    )
+  elseif(NOT WIN32)
+    add_custom_command(
+      TARGET ${name}
+      POST_BUILD
+      COMMAND patchelf --remove-needed $<TARGET_SONAME_FILE_NAME:mujoco>
+              $<TARGET_FILE_NAME:${name}>
+    )
+  endif()
+endmacro()
+
+mujoco_pybind11_module(_callbacks callbacks.cc)
+target_link_libraries(
+  _callbacks
+  PRIVATE errors_header
+          mujoco
+          raw
+          structs_header
+)
+
+mujoco_pybind11_module(_constants constants.cc)
+target_link_libraries(_constants PRIVATE mujoco)
+
+mujoco_pybind11_module(_enums enums.cc)
+target_link_libraries(
+  _enums
+  PRIVATE crossplatform
+          enum_traits
+          mujoco
+          tuple_tools
+)
+
+mujoco_pybind11_module(_errors errors.cc)
+target_link_libraries(_errors PRIVATE errors_header)
+
+mujoco_pybind11_module(_functions functions.cc)
+target_link_libraries(
+  _functions
+  PRIVATE Eigen3::Eigen
+          functions_header
+          function_traits
+          mujoco
+          raw
+)
+if(APPLE)
+  # C++17 aligned allocation is not available until macOS 10.14.
+  target_compile_options(_functions PRIVATE -fno-aligned-allocation)
+endif()
+
+mujoco_pybind11_module(_render render.cc)
+target_link_libraries(
+  _render
+  PRIVATE Eigen3::Eigen
+          errors_header
+          functions_header
+          function_traits
+          mujoco
+          raw
+          structs_header
+)
+
+mujoco_pybind11_module(_rollout rollout.cc)
+target_link_libraries(_rollout PRIVATE functions_header mujoco raw)
+
+mujoco_pybind11_module(
+  _structs
+  indexers.cc
+  serialization.h
+  structs.cc
+)
+target_link_libraries(
+  _structs
+  PRIVATE absl::flat_hash_map
+          mujoco
+          raw
+          errors_header
+          func_wrap
+          function_traits
+          structs_header
+)
+
+set(LIBRARIES_FOR_WHEEL
+    "$<TARGET_FILE:_callbacks>"
+    "$<TARGET_FILE:_constants>"
+    "$<TARGET_FILE:_enums>"
+    "$<TARGET_FILE:_errors>"
+    "$<TARGET_FILE:_functions>"
+    "$<TARGET_FILE:_render>"
+    "$<TARGET_FILE:_rollout>"
+    "$<TARGET_FILE:_structs>"
+    "$<TARGET_FILE:mujoco>"
+)
+if(NOT APPLE)
+  set(LIBRARIES_FOR_WHEEL ${LIBRARIES_FOR_WHEEL} $<TARGET_FILE:mujoco_nogl>)
+endif()
+if(NOT (APPLE OR WIN32))
+  set(LIBRARIES_FOR_WHEEL
+      ${LIBRARIES_FOR_WHEEL}
+      "$<TARGET_FILE:glew>"
+      "$<TARGET_FILE:glewegl>"
+      "$<TARGET_FILE:glewosmesa>"
+  )
+endif()
+
+if(MUJOCO_PYTHON_MAKE_WHEEL)
+  add_custom_target(
+    wheel ALL
+    COMMAND "${CMAKE_COMMAND}" -E rm -rf "${CMAKE_CURRENT_BINARY_DIR}/dist"
+    COMMAND "${CMAKE_COMMAND}" -E copy_directory "${CMAKE_CURRENT_SOURCE_DIR}/.."
+            "${CMAKE_CURRENT_BINARY_DIR}/dist"
+    COMMAND "${CMAKE_COMMAND}" -E copy "${CMAKE_CURRENT_SOURCE_DIR}/../../LICENSE"
+            "${CMAKE_CURRENT_BINARY_DIR}/dist/LICENSE"
+    COMMAND "${CMAKE_COMMAND}" -E copy ${LIBRARIES_FOR_WHEEL}
+            "${CMAKE_CURRENT_BINARY_DIR}/dist/mujoco"
+    COMMAND "${Python3_EXECUTABLE}" -m pip wheel --wheel-dir "${CMAKE_BINARY_DIR}" --no-deps -vvv
+            "${CMAKE_CURRENT_BINARY_DIR}/dist"
+  )
+
+  add_dependencies(
+    wheel
+    _callbacks
+    _constants
+    _enums
+    _errors
+    _functions
+    _render
+    _rollout
+    _structs
+    mujoco
+  )
+endif()

python/mujoco/__init__.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+"""Python bindings for MuJoCo."""
+
+import ctypes
+import os
+import platform
+import subprocess
+
+HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include')
+
+_MUJOCO_GL_ENABLE = ('enable', 'enabled', 'on', 'true', '1' , '')
+_MUJOCO_GL_DISABLE = ('disable', 'disabled', 'off', 'false', '0')
+_MUJOCO_GL = os.environ.get('MUJOCO_GL', '').lower().strip()
+_MUJOCO_GL_IS_VALID = True
+
+_SYSTEM = platform.system()
+print(_MUJOCO_GL)
+if _SYSTEM == 'Linux':
+  libglew_name = None
+  if _MUJOCO_GL in _MUJOCO_GL_ENABLE + ('glfw', 'glx'):
+    libglew_name = 'libglew.so'
+  elif _MUJOCO_GL == 'egl':
+    libglew_name = 'libglewegl.so'
+  elif _MUJOCO_GL == 'osmesa':
+    libglew_name = 'libglewosmesa.so'
+  elif _MUJOCO_GL not in _MUJOCO_GL_DISABLE:
+    _MUJOCO_GL_IS_VALID = False
+  if libglew_name is not None:
+    ctypes.CDLL(os.path.join(os.path.dirname(__file__), libglew_name),
+                ctypes.RTLD_GLOBAL)
+    ctypes.CDLL(
+        os.path.join(os.path.dirname(__file__), 'libmujoco.so.2.1.2'),
+        ctypes.RTLD_GLOBAL)
+  else:
+    ctypes.CDLL(
+        os.path.join(os.path.dirname(__file__), 'libmujoco_nogl.so.2.1.2'),
+        ctypes.RTLD_GLOBAL)
+elif _SYSTEM == 'Windows':
+  if _MUJOCO_GL in _MUJOCO_GL_ENABLE + ('glfw', 'wgl'):
+    ctypes.WinDLL(os.path.join(os.path.dirname(__file__), 'mujoco.dll'))
+  elif _MUJOCO_GL in _MUJOCO_GL_DISABLE:
+    ctypes.WinDLL(os.path.join(os.path.dirname(__file__), 'mujoco_nogl.dll'))
+  else:
+    _MUJOCO_GL_IS_VALID = False
+
+if not _MUJOCO_GL_IS_VALID:
+  raise RuntimeError(
+      f'invalid value for environment variable MUJOCO_GL: {_MUJOCO_GL}')
+
+from mujoco._callbacks import *
+from mujoco._constants import *
+from mujoco._enums import *
+from mujoco._errors import *
+from mujoco._functions import *
+from mujoco._structs import *
+
+# pylint: disable=g-import-not-at-top
+if _MUJOCO_GL not in _MUJOCO_GL_DISABLE:
+  from mujoco._render import *
+  if _SYSTEM != 'Linux':
+    from mujoco.glfw import GLContext
+  else:
+    _dl_handle = ctypes.CDLL(None)
+    if hasattr(_dl_handle, 'OSMesaCreateContextExt'):
+      from mujoco.osmesa import GLContext
+    elif hasattr(_dl_handle, 'eglCreateContext'):
+      from mujoco.egl import GLContext
+    else:
+      from mujoco.glfw import GLContext

python/mujoco/bindings_test.py

@@ -0,0 +1,847 @@
+# 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.
+# ==============================================================================
+"""Tests for MuJoCo Python bindings."""
+
+import contextlib
+import copy
+import pickle
+import sys
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import mujoco
+import numpy as np
+
+TEST_XML = r"""
+<mujoco model="test">
+  <compiler coordinate="local" angle="radian" eulerseq="xyz"/>
+  <option timestep="0.002" gravity="0 0 -9.81"/>
+  <visual>
+    <global fovy="50" />
+    <quality shadowsize="51" />
+  </visual>
+  <worldbody>
+    <geom name="myplane" type="plane" size="10 10 1"/>
+    <body pos="0 0 0.1">
+      <geom name="mybox" type="box" size="0.1 0.1 0.1" mass="0.25"/>
+      <freejoint name="myfree"/>
+    </body>
+    <body>
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <site pos="0 0 -1" name="mysite" type="sphere"/>
+      <joint type="hinge" axis="0 1 0"/>
+    </body>
+    <body>
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <joint name="myball" type="ball"/>
+    </body>
+    <body mocap="true" pos="42 0 42">
+      <geom type="sphere" size="0.1"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+TEST_XML_SENSOR = r"""
+<mujoco model="test">
+  <worldbody>
+    <geom name="myplane" type="plane" size="10 10 1"/>
+  </worldbody>
+  <sensor>
+    <user objtype="geom" objname="myplane"
+          datatype="real" needstage="vel" dim="1"/>
+  </sensor>
+</mujoco>
+"""
+
+
+@contextlib.contextmanager
+def temporary_callback(setter, callback):
+  setter(callback)
+  yield
+  setter(None)
+
+
+class MuJoCoBindingsTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.model: mujoco.MjModel = mujoco.MjModel.from_xml_string(TEST_XML)
+    self.data = mujoco.MjData(self.model)
+
+  def test_load_xml_can_handle_name_clash(self):
+    xml_1 = r"""
+<mujoco>
+  <worldbody>
+    <geom name="plane" type="plane" size="1 1 1"/>
+    <include file="model_.xml"/>
+    <include file="model__.xml"/>
+  </worldbody>
+</mujoco>"""
+    xml_2 = rb"""<mujoco><geom name="box" type="box" size="1 1 1"/></mujoco>"""
+    xml_3 = rb"""<mujoco><geom name="ball" type="sphere" size="1"/></mujoco>"""
+    model = mujoco.MjModel.from_xml_string(
+        xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3})
+    self.assertEqual(
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0)
+    self.assertEqual(
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1)
+    self.assertEqual(
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2)
+
+  def test_can_read_array(self):
+    np.testing.assert_array_equal(
+        self.model.body_pos,
+        [[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]])
+
+  def test_can_set_array(self):
+    self.data.qpos = 0.12345
+    np.testing.assert_array_equal(
+        self.data.qpos, [0.12345]*len(self.data.qpos))
+
+  def test_array_is_a_view(self):
+    qpos_ref = self.data.qpos
+    self.data.qpos = 0.789
+    np.testing.assert_array_equal(
+        qpos_ref, [0.789]*len(self.data.qpos))
+
+  def test_array_keeps_struct_alive(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    qpos0 = model.qpos0
+    qpos_spring = model.qpos_spring
+
+    # This only fails reliably under ASAN, which detects heap-use-after-free.
+    # However, often the assertEqual is enough since the memory block is
+    # already reused between mjModel deallocation and the subsequent read.
+    qpos0[:] = 1
+    del model
+    self.assertEqual(qpos0[0], 1)
+
+    # When running under test coverage tools, the refcount of objects can be
+    # higher than normal. To take this into account, we first measure the
+    # refcount of a dummy object with no other referrer.
+    dummy = []
+    base_refcount = sys.getrefcount(dummy) - 1
+
+    # Here `base` is actually a PyCapsule that holds the raw mjModel* rather
+    # than the actual MjModel wrapper object itself.
+    capsule = qpos0.base
+    self.assertEqual(sys.getrefcount(capsule) - base_refcount, 3)
+    del qpos0
+    self.assertEqual(sys.getrefcount(capsule) - base_refcount, 2)
+    del qpos_spring
+    self.assertEqual(sys.getrefcount(capsule) - base_refcount, 1)
+
+  def test_named_indexing_geom_size(self):
+    box_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, 'mybox')
+    self.assertIs(self.model.geom('mybox'), self.model.geom(box_id))
+    self.assertIs(self.model.geom('mybox').size, self.model.geom(box_id).size)
+    self.assertEqual(self.model.geom('mybox').size.shape, (3,))
+
+    # Test that the indexer is returning a view into the underlying struct.
+    size_from_indexer = self.model.geom('mybox').size
+    self.model.geom_size[box_id] = [7, 11, 13]
+    np.testing.assert_array_equal(size_from_indexer, [7, 11, 13])
+    self.model.geom('mybox').size = [5, 3, 2]
+    np.testing.assert_array_equal(self.model.geom_size[box_id], [5, 3, 2])
+
+  def test_named_indexing_geom_quat(self):
+    box_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, 'mybox')
+    self.assertIs(self.model.geom('mybox'), self.model.geom(box_id))
+    self.assertIs(self.model.geom('mybox').quat, self.model.geom(box_id).quat)
+    self.assertEqual(self.model.geom('mybox').quat.shape, (4,))
+
+    # Test that the indexer is returning a view into the underlying struct.
+    quat_from_indexer = self.model.geom('mybox').quat
+    self.model.geom_quat[box_id] = [5, 10, 15, 20]
+    np.testing.assert_array_equal(quat_from_indexer, [5, 10, 15, 20])
+    self.model.geom('mybox').quat = [12, 9, 6, 3]
+    np.testing.assert_array_equal(self.model.geom_quat[box_id], [12, 9, 6, 3])
+
+  def test_named_indexing_ragged_qpos(self):
+    balljoint_id = mujoco.mj_name2id(
+        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball')
+    self.assertIs(self.data.joint('myball'), self.data.joint(balljoint_id))
+    self.assertIs(self.data.joint('myball').qpos,
+                  self.data.joint(balljoint_id).qpos)
+    self.assertEqual(self.data.joint('myball').qpos.shape, (4,))
+
+    # Test that the indexer is returning a view into the underlying struct.
+    qpos_from_indexer = self.data.joint('myball').qpos
+    qpos_idx = self.model.jnt_qposadr[balljoint_id]
+    self.data.qpos[qpos_idx:qpos_idx+4] = [4, 5, 6, 7]
+    np.testing.assert_array_equal(qpos_from_indexer, [4, 5, 6, 7])
+    self.data.joint('myball').qpos = [9, 8, 7, 6]
+    np.testing.assert_array_equal(self.data.qpos[qpos_idx:qpos_idx+4],
+                                  [9, 8, 7, 6])
+
+  def test_named_indexing_ragged2d_cdof(self):
+    freejoint_id = mujoco.mj_name2id(
+        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree')
+    self.assertIs(self.data.joint('myfree'), self.data.joint(freejoint_id))
+    self.assertIs(self.data.joint('myfree').cdof,
+                  self.data.joint(freejoint_id).cdof)
+    self.assertEqual(self.data.joint('myfree').cdof.shape, (6, 6))
+
+    # Test that the indexer is returning a view into the underlying struct.
+    cdof_from_indexer = self.data.joint('myfree').cdof
+    dof_idx = self.model.jnt_dofadr[freejoint_id]
+    self.data.cdof[dof_idx:dof_idx+6, :] = np.reshape(range(36), (6, 6))
+    np.testing.assert_array_equal(cdof_from_indexer,
+                                  np.reshape(range(36), (6, 6)))
+    self.data.joint('myfree').cdof = 42
+    np.testing.assert_array_equal(self.data.cdof[dof_idx:dof_idx+6], [[42]*6]*6)
+
+  def test_addresses_differ_between_structs(self):
+    model2 = mujoco.MjModel.from_xml_string(TEST_XML)
+    data2 = mujoco.MjData(model2)
+
+    self.assertGreater(self.model._address, 0)
+    self.assertGreater(self.data._address, 0)
+    self.assertGreater(model2._address, 0)
+    self.assertGreater(data2._address, 0)
+    self.assertLen({self.model._address, self.data._address,
+                    model2._address, data2._address}, 4)
+
+  def test_mjmodel_can_read_and_write_opt(self):
+    self.assertEqual(self.model.opt.timestep, 0.002)
+    np.testing.assert_array_equal(self.model.opt.gravity, [0, 0, -9.81])
+
+    opt = self.model.opt
+    self.model.opt.timestep = 0.001
+    self.assertEqual(opt.timestep, 0.001)
+
+    gravity = opt.gravity
+    self.model.opt.gravity[1] = 0.1
+    np.testing.assert_array_equal(gravity, [0, 0.1, -9.81])
+    self.model.opt.gravity = 0.2
+    np.testing.assert_array_equal(gravity, [0.2, 0.2, 0.2])
+
+  def test_mjmodel_can_read_and_write_stat(self):
+    self.assertNotEqual(self.model.stat.meanmass, 0)
+
+    stat = self.model.stat
+    self.model.stat.meanmass = 1.2
+    self.assertEqual(stat.meanmass, 1.2)
+
+  def test_mjmodel_can_read_and_write_vis(self):
+    self.assertEqual(self.model.vis.quality.shadowsize, 51)
+
+    self.model.vis.quality.shadowsize = 100
+    self.assertEqual(self.model.vis.quality.shadowsize, 100)
+
+  def test_mjmodel_can_access_names_directly(self):
+    # mjModel offers direct access to names array, to allow usecases other than
+    # id2name
+    model_name = str(self.model.names[0:self.model.names.find(b'\0')], 'utf-8')
+    self.assertEqual(model_name, 'test')
+
+    start_index = self.model.name_geomadr[0]
+    end_index = self.model.names.find(b'\0', start_index)
+    geom_name = str(self.model.names[start_index:end_index], 'utf-8')
+    self.assertEqual(geom_name, 'myplane')
+
+  def test_mjmodel_names_doesnt_copy(self):
+    names = self.model.names
+    self.assertIs(names, self.model.names)
+
+  def test_vis_global_exposed_as_global_(self):
+    self.assertEqual(self.model.vis.global_.fovy, 50)
+    self.model.vis.global_.fovy = 100
+    self.assertEqual(self.model.vis.global_.fovy, 100)
+
+  def test_mjoption_can_make_default(self):
+    opt = mujoco.MjOption()
+    self.assertEqual(opt.timestep, 0.002)
+    np.testing.assert_array_equal(opt.gravity, [0, 0, -9.81])
+
+  def test_mjoption_can_copy(self):
+    opt1 = mujoco.MjOption()
+    opt1.timestep = 0.001
+    opt1.gravity = 2
+
+    opt2 = copy.copy(opt1)
+    self.assertEqual(opt2.timestep, 0.001)
+    np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
+
+    # Make sure opt2 is actually a copy.
+    opt1.timestep = 0.005
+    opt1.gravity = 5
+    self.assertEqual(opt2.timestep, 0.001)
+    np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
+
+  def test_mjmodel_can_copy(self):
+    model_copy = copy.copy(self.model)
+
+    self.assertEqual(
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0),
+        'myfree')
+
+    self.assertEqual(
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0),
+        'myplane')
+    self.assertEqual(
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1),
+        'mybox')
+
+    # Make sure it's a copy.
+    self.model.geom_size[1] = 0.5
+    np.testing.assert_array_equal(self.model.geom_size[1], [0.5, 0.5, 0.5])
+    np.testing.assert_array_equal(model_copy.geom_size[1], [0.1, 0.1, 0.1])
+
+  def test_assets_array_filename_too_long(self):
+    # Longest allowed filename (excluding null byte)
+    limit = mujoco.mjMAXVFSNAME - 1
+    contents = b'<mujoco/>'
+    valid_filename = 'a' * limit
+    mujoco.MjModel.from_xml_path(valid_filename, {valid_filename: contents})
+    invalid_filename = 'a' * (limit + 1)
+    expected_message = (
+        f'Filename length 1000 exceeds 999 character limit: {invalid_filename}')
+    with self.assertRaisesWithLiteralMatch(ValueError, expected_message):
+      mujoco.MjModel.from_xml_path(invalid_filename,
+                                   {invalid_filename: contents})
+
+  def test_mjdata_can_copy(self):
+    self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001,
+                      np.cos(np.pi/8), np.sin(np.pi/8), 0, 0, 0,
+                      1, 0, 0, 0]
+    mujoco.mj_forward(self.model, self.data)
+
+    data_copy = copy.copy(self.data)
+    self.assertEqual(data_copy.ncon, 2)
+
+    # Make sure it's a copy.
+    mujoco.mj_resetData(self.model, self.data)
+    mujoco.mj_forward(self.model, self.data)
+    mujoco.mj_forward(self.model, data_copy)
+    self.assertEqual(self.data.ncon, 4)
+    self.assertEqual(data_copy.ncon, 2)
+
+    mujoco.mj_resetData(self.model, data_copy)
+    mujoco.mj_forward(self.model, data_copy)
+    self.assertEqual(data_copy.ncon, 4)
+
+  def test_mjdata_can_read_warning_array(self):
+    warnings = self.data.warning
+    self.assertLen(warnings, mujoco.mjtWarning.mjNWARNING)
+    self.data.qpos[0] = float('NaN')
+    mujoco.mj_checkPos(self.model, self.data)
+    self.assertEqual(warnings[mujoco.mjtWarning.mjWARN_BADQPOS].number, 1)
+
+  def test_mjcontact_can_copy(self):
+    mujoco.mj_forward(self.model, self.data)
+
+    contact_copy = []
+    for i in range(4):
+      contact_copy.append(copy.copy(self.data.contact[i]))
+    # Sort contacts in anticlockwise order
+    contact_copy = sorted(
+        contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
+    np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
+    np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
+    np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
+    np.testing.assert_allclose(contact_copy[3].pos[:2], [-0.1, 0.1])
+
+    # Make sure they're actually copies.
+    for i in range(4):
+      self.data.contact[i].pos[:2] = 55
+    np.testing.assert_allclose(self.data.contact[0].pos[:2], [55, 55])
+    np.testing.assert_allclose(self.data.contact[1].pos[:2], [55, 55])
+    np.testing.assert_allclose(self.data.contact[2].pos[:2], [55, 55])
+    np.testing.assert_allclose(self.data.contact[3].pos[:2], [55, 55])
+    np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
+    np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
+    np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
+    np.testing.assert_allclose(contact_copy[3].pos[:2], [-0.1, 0.1])
+
+  def test_mj_step(self):
+    displacement = 0.25
+    self.data.qpos[2] += displacement
+    mujoco.mj_forward(self.model, self.data)
+
+    gravity = -self.model.opt.gravity[2]
+    expected_contact_time = np.sqrt(2 * displacement / gravity)
+
+    # Grab a reference to the contacts upfront so that we know that they're
+    # a view into mjData rather than a copy.
+    contact = self.data.contact[:4]
+
+    self.model.opt.timestep = 2**-9  # 0.001953125; allows exact comparisons
+    self.assertEqual(self.data.time, 0)
+    while self.data.time < expected_contact_time:
+      self.assertEqual(self.data.ncon, 0)
+      prev_time = self.data.time
+      mujoco.mj_step(self.model, self.data)
+      self.assertEqual(self.data.time, prev_time + self.model.opt.timestep)
+
+    mujoco.mj_forward(self.model, self.data)
+    self.assertEqual(self.data.ncon, 4)
+
+    # Sort contacts in anticlockwise order
+    sorted_contact = sorted(
+        contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
+    np.testing.assert_allclose(sorted_contact[0].pos[:2], [-0.1, -0.1])
+    np.testing.assert_allclose(sorted_contact[1].pos[:2], [0.1, -0.1])
+    np.testing.assert_allclose(sorted_contact[2].pos[:2], [0.1, 0.1])
+    np.testing.assert_allclose(sorted_contact[3].pos[:2], [-0.1, 0.1])
+
+  def test_mj_step_multiple(self):
+    self.model.opt.timestep = 2**-9  # 0.001953125; allows exact comparisons
+    self.assertEqual(self.data.time, 0)
+    for _ in range(10):
+      prev_time = self.data.time
+      mujoco.mj_step(self.model, self.data, nstep=7)
+      self.assertEqual(self.data.time, prev_time + 7 * self.model.opt.timestep)
+    self.assertIn('Optionally, repeat nstep times.', mujoco.mj_step.__doc__)
+
+  def test_mj_contact_list(self):
+    self.assertLen(self.data.contact, self.model.nconmax)
+
+    expected_pos = []
+    for contact in self.data.contact:
+      expected_pos.append(np.random.uniform(size=3))
+      contact.pos = expected_pos[-1]
+    np.testing.assert_array_equal(self.data.contact.pos, expected_pos)
+
+    expected_friction = []
+    for contact in self.data.contact:
+      expected_friction.append(np.random.uniform(size=5))
+      contact.friction = expected_friction[-1]
+    np.testing.assert_array_equal(self.data.contact.friction, expected_friction)
+
+    expected_H = []  # pylint: disable=invalid-name
+    for contact in self.data.contact:
+      expected_H.append(np.random.uniform(size=36))
+      contact.H = expected_H[-1]
+    np.testing.assert_array_equal(self.data.contact.H, expected_H)
+
+  def test_mj_struct_list_equality(self):
+    model2 = mujoco.MjModel.from_xml_string(TEST_XML)
+    data2 = mujoco.MjData(model2)
+
+    mujoco.mj_forward(self.model, self.data)
+    self.assertEqual(self.data.ncon, 4)
+    mujoco.mj_forward(model2, data2)
+    self.assertEqual(data2.ncon, 4)
+    self.assertEqual(data2.contact[:4], self.data.contact[:4])
+
+    self.data.qpos[3:7] = [np.cos(np.pi/8), np.sin(np.pi/8), 0, 0]
+    self.data.qpos[2] *= (np.sqrt(2) - 1) * 0.1 - 1e-6
+    mujoco.mj_forward(self.model, self.data)
+    self.assertEqual(self.data.ncon, 2)
+    self.assertNotEqual(data2.contact[:2], self.data.contact[:2])
+
+    # Check that we can compare slices of different lengths
+    self.assertNotEqual(data2.contact[:2], self.data.contact[:4])
+
+    # Check that comparing things of different types do not raise an error
+    self.assertNotEqual(self.data.contact, self.data.warning)
+    self.assertNotEqual(self.data.contact, 5)
+
+  @parameterized.named_parameters([
+      ('MjOption', mujoco.MjOption, 'tolerance'),
+      ('MjWarningStat', mujoco.MjWarningStat, 'number'),
+      ('MjTimerStat', mujoco.MjTimerStat, 'number'),
+      ('MjSolverStat', mujoco.MjSolverStat, 'neval'),
+      ('MjContact', mujoco.MjContact, 'dist'),
+      ('MjStatistic', mujoco.MjStatistic, 'extent'),
+      ('MjLROpt', mujoco.MjLROpt, 'maxforce'),
+      ('MjvPerturb', mujoco.MjvPerturb, 'select'),
+      ('MjvCamera', mujoco.MjvCamera, 'fixedcamid'),
+  ])
+  def test_mj_struct_equality(self, cls, attr):
+    struct = cls()
+    struct2 = cls()
+    setattr(struct, attr, 1)
+    self.assertNotEqual(struct, struct2)
+    setattr(struct2, attr, 1)
+    self.assertEqual(struct, struct2)
+
+    self.assertNotEqual(struct, 3)
+    self.assertNotEqual(struct, None)
+
+    # mutable structs shouldn't declare __hash__
+    with self.assertRaises(TypeError):
+      hash(struct)
+
+  def test_mj_struct_equality_array(self):
+    contact1 = mujoco.MjContact()
+    contact2 = mujoco.MjContact()
+    contact1.H[3] = 1
+    self.assertNotEqual(contact1, contact2)
+    contact2.H[3] = 1
+    self.assertEqual(contact1, contact2)
+
+  @parameterized.named_parameters([
+      ('MjOption', mujoco.MjOption, 'tolerance'),
+      ('MjWarningStat', mujoco.MjWarningStat, 'number'),
+      ('MjTimerStat', mujoco.MjTimerStat, 'number'),
+      ('MjSolverStat', mujoco.MjSolverStat, 'neval'),
+      ('MjContact', mujoco.MjContact, 'dist'),
+      ('MjStatistic', mujoco.MjStatistic, 'extent'),
+      ('MjLROpt', mujoco.MjLROpt, 'maxforce'),
+      ('MjvPerturb', mujoco.MjvPerturb, 'select'),
+      ('MjvCamera', mujoco.MjvCamera, 'fixedcamid'),
+  ])
+  def test_mj_struct_repr(self, cls, attr):
+    struct = cls()
+    setattr(struct, attr, 1)
+    representation = repr(struct)
+    self.assertStartsWith(representation, f'<{cls.__name__}')
+    self.assertIn(f'{attr}: 1', representation)
+    self.assertEqual(str(struct), repr(struct))
+
+  def test_mj_struct_repr_for_subclass(self):
+    class MjWarningStatSubclass(mujoco.MjWarningStat):
+      # ptr attribute could cause an infinite recursion, if the repr
+      # implementation simply looked at all attributes.
+      @property
+      def ptr(self):
+        return self
+
+    # repr should include name of subclass.
+    expected_repr = """<MjWarningStatSubclass
+  lastinfo: 0
+  number: 0
+>"""
+
+    self.assertEqual(repr(MjWarningStatSubclass()), expected_repr)
+  def test_mju_rotVecQuat(self):  # pylint: disable=invalid-name
+    vec = [1, 0, 0]
+    quat = [np.cos(np.pi/8), 0, 0, np.sin(np.pi/8)]
+    expected = np.array([1, 1, 0]) / np.sqrt(2)
+
+    # Check that the output argument works, and that the binding returns None.
+    res = np.empty(3, np.float64)
+    self.assertIsNone(mujoco.mju_rotVecQuat(res, vec, quat))
+    np.testing.assert_allclose(res, expected)
+
+    # Check that the function can be called via keyword arguments.
+    mujoco.mju_rotVecQuat(vec=vec, quat=quat, res=res)
+    np.testing.assert_allclose(res, expected)
+
+    # Check that the res argument must have the right size.
+    with self.assertRaises(TypeError):
+      mujoco.mju_rotVecQuat(np.empty(4, np.float64), vec, quat)
+
+    # Check that the vec argument must have the right size.
+    with self.assertRaises(TypeError):
+      mujoco.mju_rotVecQuat(res, [1, 2, 3, 4], quat)
+
+    # Check that the quat argument must have the right size.
+    with self.assertRaises(TypeError):
+      mujoco.mju_rotVecQuat(res, vec, [1, 2, 3])
+
+    # Check that the output argument must have the correct dtype.
+    with self.assertRaises(TypeError):
+      mujoco.mju_rotVecQuat(vec, quat, res=np.empty(3, int))
+
+  def test_mj_jacSite(self):  # pylint: disable=invalid-name
+    mujoco.mj_forward(self.model, self.data)
+    site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')
+
+    # Call mj_jacSite with only jacp.
+    jacp = np.empty((3, 10), np.float64)
+    mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id)
+
+    expected_jacp = np.array(
+        [[0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
+    np.testing.assert_array_equal(jacp, expected_jacp)
+
+    # Call mj_jacSite with only jacr.
+    jacr = np.empty((3, 10), np.float64)
+    mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id)
+
+    expected_jacr = np.array(
+        [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
+    np.testing.assert_array_equal(jacr, expected_jacr)
+
+    # Call mj_jacSite with both jacp and jacr.
+    jacp[:] = 0
+    jacr[:] = 0
+    mujoco.mj_jacSite(self.model, self.data, jacp, jacr, site_id)
+    np.testing.assert_array_equal(jacp, expected_jacp)
+    np.testing.assert_array_equal(jacr, expected_jacr)
+
+    # Check that the jacp argument must have the right size.
+    with self.assertRaises(TypeError):
+      mujoco.mj_jacSite(
+          self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id)
+
+    # Check that the jacr argument must have the right size.
+    with self.assertRaises(TypeError):
+      mujoco.mj_jacSite(
+          self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id)
+
+    # Check that the jacp argument must have the right dtype.
+    with self.assertRaises(TypeError):
+      mujoco.mj_jacSite(
+          self.model, self.data, np.empty(jacp.shape, int), None, site_id)
+
+    # Check that the jacp argument must have the right dtype.
+    with self.assertRaises(TypeError):
+      mujoco.mj_jacSite(
+          self.model, self.data, None, np.empty(jacr.shape, int), site_id)
+
+  def test_docstrings(self):  # pylint: disable=invalid-name
+    self.assertEqual(
+        mujoco.mj_versionString.__doc__,
+        """mj_versionString() -> str
+
+Return the current version of MuJoCo as a null-terminated string.
+""")
+    self.assertEqual(
+        mujoco.mj_Euler.__doc__,
+        """mj_Euler(m: mujoco._structs.MjModel, d: mujoco._structs.MjData) -> None
+
+Euler integrator, semi-implicit in velocity.
+""")
+
+  def test_int_constant(self):
+    self.assertEqual(mujoco.mjMAXVFSNAME, 1000)
+
+  def test_float_constant(self):
+    self.assertEqual(mujoco.mjMAXVAL, 1e10)
+
+  def test_string_constants(self):
+    self.assertLen(mujoco.mjDISABLESTRING, mujoco.mjtDisableBit.mjNDISABLE)
+    self.assertLen(mujoco.mjENABLESTRING, mujoco.mjtEnableBit.mjNENABLE)
+    self.assertLen(mujoco.mjTIMERSTRING, mujoco.mjtTimer.mjNTIMER)
+    self.assertLen(mujoco.mjLABELSTRING, mujoco.mjtLabel.mjNLABEL)
+    self.assertLen(mujoco.mjFRAMESTRING, mujoco.mjtFrame.mjNFRAME)
+    self.assertLen(mujoco.mjVISSTRING, mujoco.mjtVisFlag.mjNVISFLAG)
+    self.assertLen(mujoco.mjRNDSTRING, mujoco.mjtRndFlag.mjNRNDFLAG)
+    self.assertEqual(mujoco.mjDISABLESTRING[11], 'Refsafe')
+    self.assertEqual(mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
+                     ('&Inertia', '0', 'I'))
+
+  def test_enum_values(self):
+    self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, 0)
+    self.assertEqual(mujoco.mjtJoint.mjJNT_BALL, 1)
+    self.assertEqual(mujoco.mjtJoint.mjJNT_SLIDE, 2)
+    self.assertEqual(mujoco.mjtJoint.mjJNT_HINGE, 3)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1<<0)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1<<1)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_SENSORNOISE, 1<<3)
+    self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 4)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_SPHERE, 2)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW, 100)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW1, 101)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW2, 102)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_NONE, 1001)
+
+  def test_enum_from_int(self):
+    self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, mujoco.mjtJoint(0))
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW, mujoco.mjtGeom(value=100))
+    # mjENABLE_FWDINV and mjNENABLE have the same int value. Default to the
+    # first defined one.
+    self.assertEqual(
+        mujoco.mjtEnableBit.mjENBL_FWDINV,
+        mujoco.mjtEnableBit(mujoco.mjtEnableBit.mjNENABLE.value))
+    with self.assertRaises(ValueError):
+      mujoco.mjtJoint(1000)
+    with self.assertRaises(ValueError):
+      mujoco.mjtJoint(-1)
+
+  def test_can_raise_error(self):
+    self.data.pstack = self.data.nstack
+    with self.assertRaisesWithLiteralMatch(mujoco.FatalError, 'Stack overflow'):
+      mujoco.mj_forward(self.model, self.data)
+
+  def test_mjcb_time(self):
+
+    class CallCounter:
+
+      def __init__(self):
+        self.count = 0
+
+      def __call__(self):
+        self.count += 1
+        return self.count - 1
+
+    call_counter = CallCounter()
+    with temporary_callback(mujoco.set_mjcb_time, call_counter):
+      self.assertIs(mujoco.get_mjcb_time(), call_counter)
+
+      # Check that the callback setter and getter aren't callin g the function.
+      self.assertEqual(call_counter.count, 0)
+
+      mujoco.mj_forward(self.model, self.data)
+      self.assertGreater(call_counter.count, 0)
+
+    self.assertIsNone(mujoco.get_mjcb_time())
+
+  def test_mjcb_time_exception(self):
+
+    class TestError(RuntimeError):
+      pass
+
+    def raises_exception():
+      raise TestError('string', (1, 2, 3), {'a': 1, 'b': 2})
+
+    with temporary_callback(mujoco.set_mjcb_time, raises_exception):
+      with self.assertRaises(TestError) as e:
+        mujoco.mj_forward(self.model, self.data)
+      self.assertEqual(
+          e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2}))
+
+    # Should not raise now that we've cleared the callback.
+    mujoco.mj_forward(self.model, self.data)
+
+  def test_mjcb_time_wrong_return_type(self):
+    with temporary_callback(mujoco.set_mjcb_time, lambda: 'string'):
+      with self.assertRaisesWithLiteralMatch(
+          TypeError, 'mjcb_time callback did not return a number'):
+        mujoco.mj_forward(self.model, self.data)
+
+  def test_mjcb_time_not_callable(self):
+    with self.assertRaisesWithLiteralMatch(
+        TypeError, 'callback is not an Optional[Callable]'):
+      mujoco.set_mjcb_time(1)
+
+  def test_mjcb_sensor(self):
+
+    class SensorCallback:
+
+      def __init__(self, test, expected_model, expected_data):
+        self.test = test
+        self.expected_model = expected_model
+        self.expected_data = expected_data
+        self.count = 0
+
+      def __call__(self, m, d, stage):
+        self.test.assertIs(m, self.expected_model)
+        self.test.assertIs(d, self.expected_data)
+        self.test.assertEqual(stage, mujoco.mjtStage.mjSTAGE_VEL)
+        d.sensordata[0] = 17
+        self.count += 1
+
+    model_with_sensor = mujoco.MjModel.from_xml_string(TEST_XML_SENSOR)
+    data_with_sensor = mujoco.MjData(model_with_sensor)
+    sensor_callback = SensorCallback(self, model_with_sensor, data_with_sensor)
+    self.assertEqual(sensor_callback.count, 0)
+
+    with temporary_callback(mujoco.set_mjcb_sensor, sensor_callback):
+      mujoco.mj_forward(model_with_sensor, data_with_sensor)
+
+    self.assertEqual(sensor_callback.count, 1)
+    self.assertEqual(data_with_sensor.sensordata[0], 17)
+
+  def test_can_initialize_mjv_structs(self):
+    self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
+    self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
+    self.assertIsInstance(mujoco.MjvGLCamera(), mujoco.MjvGLCamera)
+    self.assertIsInstance(mujoco.MjvGeom(), mujoco.MjvGeom)
+    self.assertIsInstance(mujoco.MjvLight(), mujoco.MjvLight)
+    self.assertIsInstance(mujoco.MjvOption(), mujoco.MjvOption)
+    self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
+    self.assertIsInstance(mujoco.MjvScene(self.model, 100), mujoco.MjvScene)
+    self.assertIsInstance(mujoco.MjvFigure(), mujoco.MjvFigure)
+
+  def test_mjv_camera(self):
+    camera = mujoco.MjvCamera()
+    camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING
+    # IDs should be integers
+    camera.fixedcamid = 2**31 - 1
+    self.assertEqual(camera.fixedcamid, 2**31 - 1)
+    with self.assertRaises(TypeError):
+      camera.fixedcamid = 0.5
+
+  def test_mjv_scene(self):
+    scene = mujoco.MjvScene(model=self.model, maxgeom=100)
+    # scene.geoms is a fixed-length tuple of length maxgeom.
+    self.assertEqual(scene.ngeom, 0)
+    self.assertEqual(scene.maxgeom, 100)
+    self.assertLen(scene.geoms, scene.maxgeom)
+    # When the scene is updated, geoms are added to the scene
+    # (ngeom is incremented)
+    mujoco.mj_forward(self.model, self.data)
+    mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(),
+                           None, mujoco.MjvCamera(),
+                           mujoco.mjtCatBit.mjCAT_ALL, scene)
+    self.assertGreater(scene.ngeom, 0)
+
+  def test_mjv_scene_without_model(self):
+    scene = mujoco.MjvScene()
+    self.assertEqual(scene.scale, 1.0)
+    self.assertEqual(scene.maxgeom, 0)
+
+  def test_mj_ray(self):
+    # mj_ray has tricky argument types
+    geomid = np.zeros(1, np.int32)
+    mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0,
+                  geomid)
+    mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1],
+                  [0, 0, 0, 0, 0, 0], 0, 0, geomid)
+    # Check that named arguments work
+    mujoco.mj_ray(
+        m=self.model,
+        d=self.data,
+        pnt=[0, 0, 0],
+        vec=[0, 0, 1],
+        geomgroup=None,
+        flg_static=0,
+        bodyexclude=0,
+        geomid=geomid)
+
+  @parameterized.product(flg_html=(False, True), flg_pad=(False, True))
+  def test_mj_printSchema(self, flg_html, flg_pad):  # pylint: disable=invalid-name
+    # Make sure that mj_printSchema doesn't raise an exception
+    # (e.g. because the internal output buffer is too small)
+    self.assertIn('mujoco', mujoco.mj_printSchema(flg_html, flg_pad))
+
+  def test_pickle_mjdata(self):
+    mujoco.mj_step(self.model, self.data)
+    data2 = pickle.loads(pickle.dumps(self.data))
+    attr_to_compare = (
+        'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
+        'warning', 'energy'
+    )
+    self._assert_attributes_equal(data2, self.data, attr_to_compare)
+    for _ in range(10):
+      mujoco.mj_step(self.model, self.data)
+      mujoco.mj_step(self.model, data2)
+    self._assert_attributes_equal(data2, self.data, attr_to_compare)
+
+  def test_pickle_mjmodel(self):
+    model2 = pickle.loads(pickle.dumps(self.model))
+    attr_to_compare = (
+        'nq', 'nmat', 'body_pos', 'names',
+    )
+    self._assert_attributes_equal(model2, self.model, attr_to_compare)
+
+  def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
+    for name in attr_to_compare:
+      actual_value = getattr(actual_obj, name)
+      expected_value = getattr(expected_obj, name)
+      try:
+        if isinstance(expected_value, np.ndarray):
+          np.testing.assert_array_equal(actual_value, expected_value)
+        else:
+          self.assertEqual(actual_value, expected_value)
+      except AssertionError as e:
+        self.fail("Attribute '{}' differs from expected value: {}".format(
+            name, str(e)))
+
+if __name__ == '__main__':
+  absltest.main()

python/mujoco/callbacks.cc

@@ -0,0 +1,413 @@
+// 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 <cstddef>
+#include <cstdint>
+#include <exception>
+#include <limits>
+#include <sstream>
+#include <type_traits>
+
+#include <mujoco.h>
+#include "errors.h"
+#include "structs.h"
+#include "raw.h"
+#include <pybind11/eval.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+
+namespace mujoco::python {
+namespace {
+namespace py = ::pybind11;
+
+[[noreturn]] static void EscapeWithPythonException() {
+  mju_error("Python exception raised");
+  std::terminate();  // not actually reachable, mju_error doesn't return
+}
+
+template <typename T, typename U>
+using enable_if_not_const_t =
+    std::enable_if_t<std::is_same_v<std::remove_const_t<T>, T>, U>;
+
+// MuJoCo passes raw mjModel* and mjData* as arguments to callbacks, but Python
+// callables expect the corresponding MjWrapper objects. To avoid creating new
+// wrappers each time we enter callbacks, we instead maintain a global lookup
+// table that associates raw MuJoCo struct pointers back to the pointers to
+// their corresponding wrappers.
+template <typename Raw>
+static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
+  using LookupFnType = MjWrapper<Raw>* (Raw*);
+  static LookupFnType* const lookup = []() -> LookupFnType* {
+    py::gil_scoped_acquire gil;
+    auto m = py::module_::import("mujoco._structs");
+    pybind11::handle builtins(PyEval_GetBuiltins());
+    if (!builtins.contains(MjWrapper<Raw>::kFromRawPointer)) {
+      return nullptr;
+    } else {
+      try {
+        return reinterpret_cast<LookupFnType*>(
+            builtins[MjWrapper<Raw>::kFromRawPointer]
+                .template cast<std::uintptr_t>());
+      } catch (const py::cast_error&) {
+        return nullptr;
+      }
+    }
+  }();
+
+  MjWrapper<Raw>* wrapper = nullptr;
+  if (lookup) {
+    wrapper = lookup(ptr);
+  } else {
+    {
+      py::gil_scoped_acquire gil;
+      PyErr_SetString(
+          UnexpectedError::GetPyExc(),
+          "_structs module did not register its raw pointer lookup functions");
+    }
+  }
+
+  if (!wrapper) {
+    {
+      py::gil_scoped_acquire gil;
+      PyErr_SetString(
+          UnexpectedError::GetPyExc(),
+          "cannot find the corresponding wrapper for the raw mjStruct");
+    }
+  }
+
+  // Now we find the existing Python instance of our wrapper.
+  // TODO(stunya): Figure out a way to do this without invoking py::detail.
+  {
+    py::gil_scoped_acquire gil;
+    const auto [src, type] =
+        py::detail::type_caster_base<MjWrapper<Raw>>::src_and_type(wrapper);
+    if (type) {
+      py::handle instance =
+          py::detail::find_registered_python_instance(wrapper, type);
+      if (!instance) {
+        if (!PyErr_Occurred()) {
+          PyErr_SetString(
+              UnexpectedError::GetPyExc(),
+              "cannot find the Python instance of the MjWrapper");
+        }
+      } else {
+        return instance;
+      }
+    } else {
+      if (!PyErr_Occurred()) {
+        PyErr_SetString(
+            UnexpectedError::GetPyExc(),
+            "MjWrapper type isn't registered with pybind11");
+      }
+    }
+  }
+
+  EscapeWithPythonException();
+}
+
+template <typename Raw>
+static const py::handle MjWrapperLookup(const Raw* ptr) {
+  return MjWrapperLookup(const_cast<Raw*>(ptr));
+}
+
+template <typename Return, typename... Args>
+static Return
+CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
+  {
+    py::gil_scoped_acquire gil;
+    if (!py_callback) {
+      std::ostringstream msg;
+      msg << "py_" << name << " is null";
+      PyErr_SetString(UnexpectedError::GetPyExc(), msg.str().c_str());
+    } else {
+      py::handle callback(py_callback);
+      try {
+        if constexpr (std::is_void_v<Return>) {
+          callback(args...);
+          return;
+        } else {
+          return callback(args...).template cast<Return>();
+        }
+      } catch (py::error_already_set& e) {
+        e.restore();
+      } catch (const py::cast_error&) {
+        std::ostringstream msg;
+        msg << name << " callback did not return ";
+        if constexpr (std::is_integral_v<Return>) {
+          msg << "an integer";
+        } else if constexpr (std::is_floating_point_v<Return>) {
+          msg << "a number";
+        } else {
+          msg << "the correct type";
+        }
+        PyErr_SetString(PyExc_TypeError, msg.str().c_str());
+      }
+    }
+  }
+  EscapeWithPythonException();
+}
+
+static PyObject* py_mju_user_warning = nullptr;
+static void PyMjuUserWarning(const char* msg) {
+  CallPyCallback<void>("mju_user_warning", py_mju_user_warning, msg);
+}
+
+// We only support ctypes function pointers for these.
+// The PyObject* are only here so that we can return the ctypes pointers back
+// through the getters.
+static PyObject* py_mju_user_malloc = nullptr;
+static PyObject* py_mju_user_free = nullptr;
+
+static PyObject* py_mjcb_passive = nullptr;
+static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) {
+  CallPyCallback<void>("mjcb_passive", py_mjcb_passive,
+                        MjWrapperLookup(m), MjWrapperLookup(d));
+}
+
+static PyObject* py_mjcb_control = nullptr;
+static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) {
+  CallPyCallback<void>("mjcb_control", py_mjcb_control,
+                         MjWrapperLookup(m), MjWrapperLookup(d));
+}
+
+static PyObject* py_mjcb_contactfilter = nullptr;
+static int PyMjcbContactfilter(
+    const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) {
+  return CallPyCallback<int>("mjcb_contactfilter", py_mjcb_contactfilter,
+                               MjWrapperLookup(m), MjWrapperLookup(d),
+                               geom1, geom2);
+}
+
+static PyObject* py_mjcb_sensor = nullptr;
+static void
+PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) {
+  CallPyCallback<void>("mjcb_sensor", py_mjcb_sensor,
+                         MjWrapperLookup(m), MjWrapperLookup(d), stage);
+}
+
+static PyObject* py_mjcb_time = nullptr;
+static mjtNum PyMjcbTime() {
+  return CallPyCallback<mjtNum>("mjcb_time", py_mjcb_time);
+}
+
+static PyObject* py_mjcb_act_dyn = nullptr;
+static mjtNum
+PyMjcbActDyn(const raw::MjModel* m, const raw::MjData* d, int id) {
+  return CallPyCallback<mjtNum>("mjcb_act_dyn", py_mjcb_act_dyn,
+                                MjWrapperLookup(m), MjWrapperLookup(d), id);
+}
+
+static PyObject* py_mjcb_act_gain = nullptr;
+static mjtNum
+PyMjcbActGain(const raw::MjModel* m, const raw::MjData* d, int id) {
+  return CallPyCallback<mjtNum>("mjcb_act_gain", py_mjcb_act_gain,
+                                MjWrapperLookup(m), MjWrapperLookup(d), id);
+}
+
+static PyObject* py_mjcb_act_bias = nullptr;
+static mjtNum
+PyMjcbActBias(const raw::MjModel* m, const raw::MjData* d, int id) {
+  return CallPyCallback<mjtNum>("mjcb_act_bias", py_mjcb_act_bias,
+                                MjWrapperLookup(m), MjWrapperLookup(d), id);
+}
+
+// If the Python object is a ctypes function pointer, returns the corresponding
+// C function pointer. Otherwise, returns a null pointer.
+template <typename FuncPtr>
+static FuncPtr GetCFuncPtr(py::handle h) {
+  struct CTypes { PyObject* cfuncptr; PyObject* cast; PyObject* c_void_p; };
+  static const CTypes ctypes = []() -> CTypes {
+    try {
+      auto m = py::module_::import("ctypes");
+      PyObject* cfuncptr = m.attr("_CFuncPtr").ptr();
+      PyObject* cast = m.attr("cast").ptr();
+      PyObject* c_void_p = m.attr("c_void_p").ptr();
+      Py_XINCREF(cfuncptr);
+      Py_XINCREF(cast);
+      Py_XINCREF(c_void_p);
+      return {cfuncptr, cast, c_void_p};
+    } catch (const py::error_already_set&) {
+      return {nullptr, nullptr, nullptr};
+    }
+  }();
+
+  if (!ctypes.cfuncptr) {
+    throw UnexpectedError("cannot find `ctypes._CFuncPtr`");
+  }
+  const int is_cfuncptr = PyObject_IsInstance(h.ptr(), ctypes.cfuncptr);
+  if (is_cfuncptr == -1) {
+    throw py::error_already_set();
+  } else if (is_cfuncptr) {
+    if (!ctypes.cast) {
+      throw UnexpectedError("cannot find `ctypes.cast`");
+    }
+    if (!ctypes.c_void_p) {
+      throw UnexpectedError("cannot find `ctypes.c_void_p`");
+    }
+    const uintptr_t func_address =
+        py::handle(ctypes.cast)(h, py::handle(ctypes.c_void_p))
+            .attr("value")
+            .template cast<std::uintptr_t>();
+    return reinterpret_cast<FuncPtr>(func_address);
+  } else {
+    return nullptr;
+  }
+}
+
+static bool IsCallable(py::handle h) {
+  static PyObject* const is_callable = []() -> PyObject* {
+    try{
+      PyObject* o = py::eval("callable").ptr();
+      Py_XINCREF(o);
+      return o;
+    } catch (const py::error_already_set&) {
+      return nullptr;
+    }
+  }();
+
+  if (!is_callable) {
+    throw UnexpectedError("cannot find `callable`");
+  }
+
+  return py::handle(is_callable)(h).cast<bool>();
+}
+
+template <typename CFuncPtr>
+void SetCallback(py::handle h, CFuncPtr py_trampoline,
+                 PyObject** py_callback, CFuncPtr* mj_callback) {
+  CFuncPtr cfuncptr = GetCFuncPtr<CFuncPtr>(h);
+  if (h.is_none()) {
+    Py_XDECREF(*py_callback);
+    *py_callback = nullptr;
+    *mj_callback = nullptr;
+  } else if (cfuncptr) {
+    Py_XDECREF(*py_callback);
+    Py_INCREF(h.ptr());
+    *py_callback = h.ptr();
+    *mj_callback = cfuncptr;
+  } else if (IsCallable(h)) {
+    Py_XDECREF(*py_callback);
+    Py_INCREF(h.ptr());
+    *py_callback = h.ptr();
+    *mj_callback = py_trampoline;
+  } else {
+    throw py::type_error("callback is not an Optional[Callable]");
+  }
+}
+
+py::object GetCallback(PyObject* py_callback) {
+  if (!py_callback) {
+    return py::none();
+  }
+  return py::reinterpret_borrow<py::object>(py_callback);
+}
+
+PYBIND11_MODULE(_callbacks, pymodule) {
+  // Setters
+  pymodule.def("set_mju_user_warning", [](py::handle h) {
+    SetCallback(h, PyMjuUserWarning, &py_mju_user_warning, &::mju_user_warning);
+  });
+  pymodule.def("set_mju_user_malloc", [](py::handle h) {
+    if (h.is_none()) {
+      Py_XDECREF(py_mju_user_malloc);
+      py_mju_user_malloc = nullptr;
+    } else {
+      auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_malloc)>(h);
+      if (!cfuncptr) {
+        throw py::type_error("mju_user_malloc must be a C function pointer");
+      }
+      Py_XDECREF(py_mju_user_malloc);
+      Py_XINCREF(h.ptr());
+      py_mju_user_malloc = h.ptr();
+      ::mju_user_malloc = cfuncptr;
+    }
+  });
+  pymodule.def("set_mju_user_free", [](py::handle h) {
+    if (h.is_none()) {
+      Py_XDECREF(py_mju_user_free);
+      py_mju_user_free = nullptr;
+    } else {
+      auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_free)>(h);
+      if (!cfuncptr) {
+        throw py::type_error("mju_user_free must be a C function pointer");
+      }
+      Py_XDECREF(py_mju_user_free);
+      Py_XINCREF(h.ptr());
+      py_mju_user_free = h.ptr();
+      ::mju_user_free = cfuncptr;
+    }
+  });
+  pymodule.def("set_mjcb_passive", [](py::handle h) {
+    SetCallback(h, PyMjcbPassive, &py_mjcb_passive, &::mjcb_passive);
+  });
+  pymodule.def("set_mjcb_control", [](py::handle h) {
+    SetCallback(h, PyMjcbControl, &py_mjcb_control, &::mjcb_control);
+  });
+  pymodule.def("set_mjcb_contactfilter", [](py::handle h) {
+    SetCallback(h, PyMjcbContactfilter,
+                &py_mjcb_contactfilter, &::mjcb_contactfilter);
+  });
+  pymodule.def("set_mjcb_sensor", [](py::handle h) {
+    SetCallback(h, PyMjcbSensor, &py_mjcb_sensor, &::mjcb_sensor);
+  });
+  pymodule.def("set_mjcb_time", [](py::handle h) {
+    SetCallback(h, PyMjcbTime, &py_mjcb_time, &::mjcb_time);
+  });
+  pymodule.def("set_mjcb_act_dyn", [](py::handle h) {
+    SetCallback(h, PyMjcbActDyn, &py_mjcb_act_dyn, &::mjcb_act_dyn);
+  });
+  pymodule.def("set_mjcb_act_gain", [](py::handle h) {
+    SetCallback(h, PyMjcbActGain, &py_mjcb_act_gain, &::mjcb_act_gain);
+  });
+  pymodule.def("set_mjcb_act_bias", [](py::handle h) {
+    SetCallback(h, PyMjcbActBias, &py_mjcb_act_bias, &::mjcb_act_bias);
+  });
+
+  // Getters
+  pymodule.def("get_mju_user_warning", []() {
+    return GetCallback(py_mju_user_warning);
+  });
+  pymodule.def("get_mju_user_malloc", []() {
+    return GetCallback(py_mju_user_malloc);
+  });
+  pymodule.def("get_mju_user_free", []() {
+    return GetCallback(py_mju_user_free);
+  });
+  pymodule.def("get_mjcb_passive", []() {
+    return GetCallback(py_mjcb_passive);
+  });
+  pymodule.def("get_mjcb_control", []() {
+    return GetCallback(py_mjcb_control);
+  });
+  pymodule.def("get_mjcb_contactfilter", []() {
+    return GetCallback(py_mjcb_contactfilter);
+  });
+  pymodule.def("get_mjcb_sensor", []() {
+    return GetCallback(py_mjcb_sensor);
+  });
+  pymodule.def("get_mjcb_time", []() {
+    return GetCallback(py_mjcb_time);
+  });
+  pymodule.def("get_mjcb_act_dyn", []() {
+    return GetCallback(py_mjcb_act_dyn);
+  });
+  pymodule.def("get_mjcb_act_gain", []() {
+    return GetCallback(py_mjcb_act_gain);
+  });
+  pymodule.def("get_mjcb_act_bias", []() {
+    return GetCallback(py_mjcb_act_bias);
+  });
+}  // PYBIND11_MODULE
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/codegen/generate_enum_traits.py

@@ -0,0 +1,89 @@
+# 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.
+# ==============================================================================
+"""Code generator for function_traits.h."""
+
+from typing import Mapping, Sequence
+
+from absl import app
+
+from introspect import ast_nodes
+from introspect import enums
+
+ENUMS: Mapping[str, ast_nodes.EnumDecl] = enums.ENUMS
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  struct_decls = []
+  for enum in ENUMS.values():
+    value_decls = []
+    for k in enum.values:
+      value_decls.append(f'std::make_pair("{k}", ::{enum.name}::{k})')
+    if len(value_decls) < 2:
+      value_decls = ''.join(value_decls)
+    else:
+      value_decls = '\n    ' + ',\n    '.join(value_decls)
+    struct_decls.append(f"""
+struct {enum.name} {{
+  static constexpr char name[] = "{enum.name}";
+  using type = ::{enum.name};
+  static constexpr auto values = std::array{{{value_decls}}};
+}};
+""".strip())
+
+  all_structs = '\n\n'.join(struct_decls)
+
+  all_enum_inits = '\n    ' + '{},\n    '.join(ENUMS.keys()) + '{}'
+
+  print(f"""
+// 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_PYTHON_CODEGEN_ENUM_TRAITS_H_
+#define MUJOCO_PYTHON_CODEGEN_ENUM_TRAITS_H_
+
+#include <array>
+#include <tuple>
+#include <utility>
+
+#include <mujoco.h>
+
+namespace mujoco::python_traits {{
+
+{all_structs}
+
+static constexpr auto kAllEnums = std::make_tuple({all_enum_inits});
+
+}}  // namespace mujoco::python_traits
+
+#endif  // MUJOCO_PYTHON_CODEGEN_ENUM_TRAITS_H_
+""".lstrip())
+
+
+if __name__ == '__main__':
+  app.run(main)

python/mujoco/codegen/generate_function_traits.py

@@ -0,0 +1,108 @@
+# 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.
+# ==============================================================================
+"""Code generator for function_traits.h."""
+
+from typing import Mapping, Sequence
+
+from absl import app
+
+from introspect import ast_nodes
+from introspect import functions
+
+FUNCTIONS: Mapping[str, ast_nodes.FunctionDecl] = functions.FUNCTIONS
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  struct_decls = []
+  for func in FUNCTIONS.values():
+
+    # Modify some parameter types.
+    parameters = []
+    modified = False
+    for p in func.parameters:
+      # Expose array parameters as pointer-to-arrays so that we can determine
+      # array extents in C++ templates.
+      if isinstance(p.type, ast_nodes.ArrayType):
+        parameters.append(ast_nodes.FunctionParameterDecl(
+            name=p.name, type=ast_nodes.PointerType(
+                ast_nodes.ArrayType(
+                    inner_type=p.type.inner_type, extents=p.type.extents))))
+        modified = True
+      else:
+        parameters.append(p)
+
+    if modified:
+      func = ast_nodes.FunctionDecl(
+          name=func.name, return_type=func.return_type,
+          parameters=parameters, doc=func.doc)
+      getfunc = f'*reinterpret_cast<type*>(&::{func.name})'
+    else:
+      getfunc = f'::{func.name}'
+
+    param_names = ', '.join(f'"{p.name}"' for p in parameters)
+
+    struct_decls.append(f"""
+struct {func.name} {{
+  static constexpr char name[] = "{func.name}";
+  static constexpr char doc[] = "{func.doc}";
+  using type = {func.decltype};
+  static constexpr auto param_names = std::make_tuple({param_names});
+
+  MUJOCO_ALWAYS_INLINE static type& GetFunc() {{
+    return {getfunc};
+  }}
+}};
+""".strip())
+
+  all_structs = '\n\n'.join(struct_decls)
+
+  print(f"""
+// 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_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
+#define MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
+
+#include <tuple>
+
+#include <mujoco.h>
+#include "util/crossplatform.h"
+
+namespace mujoco::python_traits {{
+
+{all_structs}
+
+}}  // namespace mujoco::python_traits
+
+#endif  // MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
+""".lstrip())
+
+
+if __name__ == '__main__':
+  app.run(main)

python/mujoco/constants.cc

@@ -0,0 +1,91 @@
+// 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 <vector>
+#include <mjmodel.h>
+#include <mjvisualize.h>
+#include <mujoco.h>
+#include <pybind11/cast.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace mujoco::python {
+namespace {
+
+namespace py = ::pybind11;
+
+template <auto N>
+py::tuple MakeTuple(
+    const char* (&strings)[N]) {
+  py::list result;
+  for (int i = 0; i < N; i++) {
+    result.append(py::str(strings[i]));
+  }
+  return result;
+}
+
+template <auto N>
+py::tuple MakeTuple(const char* (&strings)[N][3]) {
+  py::list result;
+  for (int i = 0; i < N; i++) {
+    result.append(py::make_tuple(
+        py::str(strings[i][0]),
+        py::str(strings[i][1]),
+        py::str(strings[i][2])));
+  }
+  return result;
+}
+
+PYBIND11_MODULE(_constants, pymodule) {
+  #define X(var) pymodule.attr(#var) = var
+
+  // from mjmodel.h
+  X(mjPI);
+  X(mjMAXVAL);
+  X(mjMINMU);
+  X(mjMINIMP);
+  X(mjMAXIMP);
+  X(mjMAXCONPAIR);
+  X(mjMAXVFS);
+  X(mjMAXVFSNAME);
+  X(mjNEQDATA);
+  X(mjNDYN);
+  X(mjNGAIN);
+  X(mjNBIAS);
+  X(mjNREF);
+  X(mjNIMP);
+  X(mjNSOLVER);
+
+  // from mjvisualize.h
+  X(mjNGROUP);
+  X(mjMAXLIGHT);
+  X(mjMAXOVERLAY);
+  X(mjMAXLINE);
+  X(mjMAXLINEPNT);
+  X(mjMAXPLANEGRID);
+
+  // from mujoco.h
+  X(mjVERSION_HEADER);
+
+  #undef X
+  pymodule.attr("mjDISABLESTRING") = MakeTuple(mjDISABLESTRING);
+  pymodule.attr("mjENABLESTRING") = MakeTuple(mjENABLESTRING);
+  pymodule.attr("mjTIMERSTRING") = MakeTuple(mjTIMERSTRING);
+  pymodule.attr("mjLABELSTRING") = MakeTuple(mjLABELSTRING);
+  pymodule.attr("mjFRAMESTRING") = MakeTuple(mjFRAMESTRING);
+  pymodule.attr("mjVISSTRING") = MakeTuple(mjVISSTRING);
+  pymodule.attr("mjRNDSTRING") = MakeTuple(mjRNDSTRING);
+}
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/egl/__init__.py

@@ -0,0 +1,130 @@
+# Copyright 2018 The dm_control Authors
+#
+# 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.
+# ==============================================================================
+"""An EGL context for headless accelerated OpenGL rendering on GPU devices."""
+
+import atexit
+import ctypes
+import os
+
+PYOPENGL_PLATFORM = os.environ.get('PYOPENGL_PLATFORM')
+
+if not PYOPENGL_PLATFORM:
+  os.environ['PYOPENGL_PLATFORM'] = 'egl'
+elif PYOPENGL_PLATFORM.lower() != 'egl':
+  raise ImportError(
+      'Cannot use EGL rendering platform. '
+      'The PYOPENGL_PLATFORM environment variable is set to {!r} '
+      '(should be either unset or \'egl\').')
+
+from mujoco.egl import egl_ext as EGL
+from OpenGL import error
+
+
+def create_initialized_egl_device_display():
+  """Creates an initialized EGL display directly on a device."""
+  all_devices = EGL.eglQueryDevicesEXT()
+  selected_device = os.environ.get('MUJOCO_EGL_DEVICE_ID', None)
+  if selected_device is None:
+    candidates = all_devices
+  else:
+    device_idx = int(selected_device)
+    if not 0 <= device_idx < len(all_devices):
+      raise RuntimeError(
+          f'The MUJOCO_EGL_DEVICE_ID environment variable must be an integer '
+          f'between 0 and {len(all_devices)-1} (inclusive), got {device_idx}.')
+    candidates = all_devices[device_idx:device_idx + 1]
+  for device in candidates:
+    display = EGL.eglGetPlatformDisplayEXT(
+        EGL.EGL_PLATFORM_DEVICE_EXT, device, None)
+    if display != EGL.EGL_NO_DISPLAY and EGL.eglGetError() == EGL.EGL_SUCCESS:
+      # `eglInitialize` may or may not raise an exception on failure depending
+      # on how PyOpenGL is configured. We therefore catch a `GLError` and also
+      # manually check the output of `eglGetError()` here.
+      try:
+        initialized = EGL.eglInitialize(display, None, None)
+      except error.GLError:
+        pass
+      else:
+        if initialized == EGL.EGL_TRUE and EGL.eglGetError() == EGL.EGL_SUCCESS:
+          return display
+  return EGL.EGL_NO_DISPLAY
+
+
+EGL_DISPLAY = create_initialized_egl_device_display()
+if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
+  raise ImportError(
+      'Cannot initialize a EGL device display. This likely means that your EGL '
+      'driver does not support the PLATFORM_DEVICE extension, which is '
+      'required for creating a headless rendering context.')
+atexit.register(EGL.eglTerminate, EGL_DISPLAY)
+
+
+EGL_ATTRIBUTES = (
+    EGL.EGL_RED_SIZE, 8,
+    EGL.EGL_GREEN_SIZE, 8,
+    EGL.EGL_BLUE_SIZE, 8,
+    EGL.EGL_ALPHA_SIZE, 8,
+    EGL.EGL_DEPTH_SIZE, 24,
+    EGL.EGL_STENCIL_SIZE, 8,
+    EGL.EGL_COLOR_BUFFER_TYPE, EGL.EGL_RGB_BUFFER,
+    EGL.EGL_SURFACE_TYPE, EGL.EGL_PBUFFER_BIT,
+    EGL.EGL_RENDERABLE_TYPE, EGL.EGL_OPENGL_BIT,
+    EGL.EGL_NONE
+)
+
+
+class GLContext:
+  """An EGL context for headless accelerated OpenGL rendering on GPU devices."""
+
+  def __init__(self, max_width, max_height):
+    del max_width, max_height  # unused
+    num_configs = ctypes.c_long()
+    config_size = 1
+    config = EGL.EGLConfig()
+    EGL.eglReleaseThread()
+    EGL.eglChooseConfig(
+        EGL_DISPLAY,
+        EGL_ATTRIBUTES,
+        ctypes.byref(config),
+        config_size,
+        num_configs)
+    if num_configs.value < 1:
+      raise RuntimeError(
+          'EGL failed to find a framebuffer configuration that matches the '
+          'desired attributes: {}'.format(EGL_ATTRIBUTES))
+    EGL.eglBindAPI(EGL.EGL_OPENGL_API)
+    self._context = EGL.eglCreateContext(
+        EGL_DISPLAY, config, EGL.EGL_NO_CONTEXT, None)
+    if not self._context:
+      raise RuntimeError('Cannot create an EGL context.')
+
+  def make_current(self):
+    if not EGL.eglMakeCurrent(
+        EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, self._context):
+      raise RuntimeError('Failed to make the EGL context current.')
+
+  def free(self):
+    """Frees resources associated with this context."""
+    if self._context:
+      current_context = EGL.eglGetCurrentContext()
+      if current_context and self._context.address == current_context.address:
+        EGL.eglMakeCurrent(EGL_DISPLAY, EGL.EGL_NO_SURFACE,
+                           EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
+      EGL.eglDestroyContext(EGL_DISPLAY, self._context)
+      EGL.eglReleaseThread()
+    self._context = None
+
+  def __del__(self):
+    self.free()

python/mujoco/egl/egl_ext.py

@@ -0,0 +1,71 @@
+# Copyright 2018 The dm_control Authors
+#
+# 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.
+# ==============================================================================
+"""Extends OpenGL.EGL with definitions necessary for headless rendering."""
+
+import ctypes
+from OpenGL.platform import ctypesloader  # pylint: disable=g-bad-import-order
+try:
+  # Nvidia driver seems to need libOpenGL.so (as opposed to libGL.so)
+  # for multithreading to work properly. We load this in before everything else.
+  ctypesloader.loadLibrary(ctypes.cdll, 'OpenGL', mode=ctypes.RTLD_GLOBAL)
+except OSError:
+  pass
+
+# pylint: disable=g-import-not-at-top
+from OpenGL import EGL
+from OpenGL import error
+
+# From the EGL_EXT_device_enumeration extension.
+PFNEGLQUERYDEVICESEXTPROC = ctypes.CFUNCTYPE(
+    EGL.EGLBoolean,
+    EGL.EGLint,
+    ctypes.POINTER(EGL.EGLDeviceEXT),
+    ctypes.POINTER(EGL.EGLint),
+)
+try:
+  _eglQueryDevicesEXT = PFNEGLQUERYDEVICESEXTPROC(  # pylint: disable=invalid-name
+      EGL.eglGetProcAddress('eglQueryDevicesEXT'))
+except TypeError as e:
+  raise ImportError('eglQueryDevicesEXT is not available.') from e
+
+
+# From the EGL_EXT_platform_device extension.
+EGL_PLATFORM_DEVICE_EXT = 0x313F
+PFNEGLGETPLATFORMDISPLAYEXTPROC = ctypes.CFUNCTYPE(
+    EGL.EGLDisplay, EGL.EGLenum, ctypes.c_void_p, ctypes.POINTER(EGL.EGLint))
+try:
+  eglGetPlatformDisplayEXT = PFNEGLGETPLATFORMDISPLAYEXTPROC(  # pylint: disable=invalid-name
+      EGL.eglGetProcAddress('eglGetPlatformDisplayEXT'))
+except TypeError as e:
+  raise ImportError('eglGetPlatformDisplayEXT is not available.') from e
+
+
+# Wrap raw _eglQueryDevicesEXT function into something more Pythonic.
+def eglQueryDevicesEXT(max_devices=10):  # pylint: disable=invalid-name
+  devices = (EGL.EGLDeviceEXT * max_devices)()
+  num_devices = EGL.EGLint()
+  success = _eglQueryDevicesEXT(max_devices, devices, num_devices)
+  if success == EGL.EGL_TRUE:
+    return [devices[i] for i in range(num_devices.value)]
+  else:
+    raise error.GLError(err=EGL.eglGetError(),
+                        baseOperation=eglQueryDevicesEXT,
+                        result=success)
+
+
+# Expose everything from upstream so that
+# we can use this as a drop-in replacement for OpenGL.EGL.
+# pylint: disable=wildcard-import,g-bad-import-order
+from OpenGL.EGL import *

python/mujoco/enums.cc

@@ -0,0 +1,62 @@
+// 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 <sstream>
+#include <type_traits>
+
+#include "util/crossplatform.h"
+#include "enum_traits.h"
+#include "util/tuple_tools.h"
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+namespace {
+namespace py = ::pybind11;
+template <typename Trait>
+MUJOCO_ALWAYS_INLINE
+void DefEnum(py::module_& m) {
+  py::enum_<typename Trait::type> e(m, Trait::name);
+  for (const auto& [name, enumerator] : Trait::values) {
+    e.value(name, enumerator);
+  }
+
+  e.def(py::init([](int value) {
+          for (const auto& [name, enumerator] : Trait::values) {
+            if (value == enumerator) {
+              return static_cast<typename Trait::type>(value);
+            }
+          }
+          std::ostringstream err;
+          err << "Invalid int value for " << Trait::name << ": " << value;
+          throw py::value_error(err.str());
+        }),
+        py::arg("value"), py::prepend());
+}
+
+template <typename Tuple>
+MUJOCO_ALWAYS_INLINE
+void DefAllEnums(py::module_& m, Tuple&& tuple) {
+  using TupleNoRef = std::remove_reference_t<Tuple>;
+  if constexpr (std::tuple_size_v<TupleNoRef> != 0) {
+    using Trait = std::remove_reference_t<std::tuple_element_t<0, TupleNoRef>>;
+    DefEnum<Trait>(m);
+    DefAllEnums(m, util::tuple_slice<1, void>(tuple));
+  }
+}
+
+PYBIND11_MODULE(_enums, pymodule) {
+  DefAllEnums(pymodule, python_traits::kAllEnums);
+}
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/errors.cc

@@ -0,0 +1,25 @@
+// 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 "errors.h"
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+namespace {
+PYBIND11_MODULE(_errors, m) {
+  m.attr("FatalError") = FatalError::GetPyExc();
+  m.attr("UnexpectedError") = UnexpectedError::GetPyExc();
+}
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/errors.h

@@ -0,0 +1,193 @@
+// 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_PYTHON_ERRORS_H_
+#define MUJOCO_PYTHON_ERRORS_H_
+
+#include <csetjmp>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+
+#include <mjexport.h>
+#include "util/crossplatform.h"
+#include "util/func_wrap.h"
+#include <pybind11/pybind11.h>
+
+// DO NOT USE THIS FUNCTION ELSEWHERE.
+// It should be regarded as part of MuJoCo's internal implementation detail.
+extern "C" {
+MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
+}
+
+// When building for Linux and statically linking against a "hermetic" libc++abi
+// (i.e. where libc++/libc++abi symbols all have "hidden" visibility), exception
+// types do not propagate correctly across shared library boundaries.
+//
+// This occurs even when PYBIND11_EXPORT_EXCEPTION is used, or when the virtual
+// destructor is defined in errors.cc rather than than defined inline in the
+// class definition (following the FAQ comment about "key functions" on
+// https://libcxxabi.llvm.org/). For some reason, the type_info address is
+// different across different shared objects (DSO), even when the symbols for
+// typeinfo, type name, and vtable are left undefined in all but one DSO.
+// Since libc++abi matches exception types by comparing type_info addresses,
+// this breaks exception handling across DSO boundaries.
+//
+// Instead, we manually create and store exception types in PyEval_GetBuiltins,
+// then effectively register separate pybind11 exception translators in each
+// module that all translate to the same exception type (which is what's
+// happening with pybind11's built-in exceptions under this setup).
+//
+// Effectively, we are doing almost the same thing as what a pybind11 is doing
+// with its "internals" struct, but we store the exception types directly in
+// the builtin context rather than in a PyCapsule.
+namespace mujoco::python {
+namespace _impl {
+template <typename T>
+class ErrorBase : public pybind11::builtin_exception {
+ public:
+  virtual ~ErrorBase() = default;
+
+  static PyObject* GetPyExc() {
+    static PyObject* const e = []() {
+      pybind11::gil_scoped_acquire gil;
+      std::string unique_identifier = "__MUJOCO_ERROR_";
+      unique_identifier += T::kName;
+      pybind11::str py_builtin_identifier(unique_identifier);
+
+      // We can end up here while handling another Python exception.
+      // Temporarily clear the Python error indicator since we need to interact
+      // with the interpreter.
+      struct PyErrCache {
+        PyErrCache() { PyErr_Fetch(&type, &value, &traceback); }
+        ~PyErrCache() { if (type) PyErr_Restore(type, value, traceback); }
+        PyObject* type;
+        PyObject* value;
+        PyObject* traceback;
+      };
+      PyErrCache err_cache;
+
+      pybind11::handle builtins(PyEval_GetBuiltins());
+      if (!builtins.contains(py_builtin_identifier)) {
+        std::string full_name = "mujoco.";
+        full_name += T::kName;
+        auto ret =
+            PyErr_NewException(full_name.c_str(), PyExc_Exception, nullptr);
+        builtins[py_builtin_identifier] = pybind11::handle(ret);
+        return ret;
+      } else {
+        return builtins[py_builtin_identifier].ptr();
+      }
+    }();
+    return e;
+  }
+
+  void set_error() const override { PyErr_SetString(T::GetPyExc(), what()); }
+
+ protected:
+  using builtin_exception::builtin_exception;
+};
+
+// We shouldn't throw a C++ exception from a function that's a callback
+// from C. (Usually it would work, but it would be undefined behavior since
+// there'd be no guarantee that C++ can unwind the stack correctly through the
+// C functions. On Windows, longjmp actually uses the same mechanism as
+// C++ exceptions.)
+// Instead, we call setjmp before entering MuJoCo, and do a longjmp from
+// mju_user_error back to C++ to throw an exception.
+static thread_local std::jmp_buf mju_error_jmp_buf;
+static thread_local std::array<char, 1024> mju_error_msg{0};
+
+static void MjErrorHandler(const char* msg) {
+  std::strncpy(mju_error_msg.data(), msg, mju_error_msg.size());
+  std::longjmp(mju_error_jmp_buf, 1);
+}
+
+template <typename InterceptAsType>
+struct MjErrorIntercepter {
+  template <typename Return, typename... Args, typename Callable>
+  MUJOCO_ALWAYS_INLINE
+  static constexpr auto WrapFunc(Callable&& callable) {
+    return [callable](Args... args) MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
+      _mjPRIVATE__set_tls_error_fn(&MjErrorHandler);
+
+      // DON'T MIX RAII WITH SETJMP!
+      // From https://en.cppreference.com/w/cpp/utility/program/longjmp:
+      //   If replacing of std::longjmp with `throw` and setjmp with `catch`
+      //   would execute a non-trivial destructor for any automatic object, the
+      //   behavior of such std::longjmp is undefined.
+      if (setjmp(mju_error_jmp_buf) == 0) {
+        if constexpr (std::is_void_v<decltype(callable(args...))>) {
+          callable(args...);
+          _mjPRIVATE__set_tls_error_fn(nullptr);
+        } else {
+          auto ret = callable(args...);
+          static_assert(std::is_trivially_destructible_v<decltype(ret)>);
+          _mjPRIVATE__set_tls_error_fn(nullptr);
+          return ret;
+        }
+      } else {
+        // This branch is entered via a longjmp back from our mju_error handler.
+        _mjPRIVATE__set_tls_error_fn(nullptr);
+        {
+          // Check if a Python callback has thrown an exception.
+          // We cannot use a py::gil_scoped_acquire here: on Windows its
+          // destructor isn't triggered before pybind returns control to the
+          // interpreter.
+          auto gil = PyGILState_Ensure();
+          if (PyErr_Occurred()) {
+            // We must hold the GIL when we create the py::error_already_set,
+            // since its constructor calls PyErr_Fetch.
+            pybind11::error_already_set err;
+
+            // But the GIL must be released before we throw, otherwise a
+            // deadlock ensues!
+            PyGILState_Release(gil);
+            throw err;
+          }
+          PyGILState_Release(gil);
+        }
+        throw InterceptAsType(std::string(mju_error_msg.data()));
+      }
+    };
+  }
+};
+}  // namespace _impl
+
+class FatalError : public _impl::ErrorBase<FatalError> {
+ public:
+  static constexpr char kName[] = "FatalError";
+  using ErrorBase<FatalError>::ErrorBase;
+  virtual ~FatalError() = default;
+};
+
+class UnexpectedError : public _impl::ErrorBase<UnexpectedError> {
+ public:
+  static constexpr char kName[] = "UnexpectedError";
+  UnexpectedError(const std::string& msg)
+      : ErrorBase(msg +
+                  " (this error not expected to ever occur,"
+                  " please report it to MuJoCo developers!)") {}
+  virtual ~UnexpectedError() = default;
+};
+
+template <typename Callable>
+MUJOCO_ALWAYS_INLINE
+static constexpr auto InterceptMjErrors(Callable&& callable) {
+  return util::_impl::WrapFunc<_impl::MjErrorIntercepter<FatalError>>(
+      std::forward<Callable>(callable));
+}
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_ERRORS_H_

python/mujoco/functions.h

@@ -0,0 +1,206 @@
+// 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_PYTHON_FUNCTIONS_H_
+#define MUJOCO_PYTHON_FUNCTIONS_H_
+
+#include <array>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <tuple>
+#include <type_traits>
+
+#include <Eigen/Core>
+#include <mujoco.h>
+#include "errors.h"
+#include "structs.h"
+#include "util/array_traits.h"
+#include "util/crossplatform.h"
+#include "util/func_wrap.h"
+#include "util/tuple_tools.h"
+#include <pybind11/eigen.h>
+#include <pybind11/pybind11.h>
+
+// Performs a compile-time check that the omitted argument name list is a
+// subset of the underlying function parameter names, and returns a helper that
+// defines a pybind11 function whose argument names obtained by removing
+// __VA_ARGS__ from the param_names tuple in the given MjTraits.
+// (This has to be implemented as a a macro because we cannot perform a
+// constexpr comparison of a tuple that is passed as a function argument.)
+#define DEF_WITH_OMITTED_PY_ARGS(MJTRAITS, ...)                                \
+  static_assert(                                                               \
+      util::is_subset_strings(std::make_tuple(__VA_ARGS__),                    \
+                              MJTRAITS::param_names),                          \
+      "omitted argument names is not a subset of function parameter names");   \
+  DefWithOmittedPyArgsImpl<MJTRAITS, decltype(std::make_tuple(__VA_ARGS__))> { \
+    std::make_tuple(__VA_ARGS__)                                               \
+  }
+
+namespace mujoco::util {
+template <typename RawMj>
+struct wrapped<RawMj*, python::enable_if_mj_struct_t<RawMj>> {
+  MUJOCO_ALWAYS_INLINE
+  static constexpr RawMj* unwrap(python::MjWrapper<RawMj>& wrapper) {
+    return wrapper.get();
+  }
+};
+
+// We use std::optional on pointer arguments to indicate that Python callers
+// can pass None.
+template <typename RawMj>
+struct wrapped<std::optional<RawMj*>, python::enable_if_mj_struct_t<RawMj>> {
+  MUJOCO_ALWAYS_INLINE
+  static constexpr std::optional<RawMj*> unwrap(
+      std::optional<python::MjWrapper<RawMj>*> wrapper) {
+    if (wrapper.has_value()) {
+      return (*wrapper)->get();
+    }
+    return std::nullopt;
+  }
+};
+
+template <typename Arr>
+using enable_if_arithmetic_array_t = std::enable_if_t<
+    std::is_array_v<Arr> &&
+    std::is_arithmetic_v<std::remove_all_extents_t<Arr>>>;
+
+template <typename Arr>
+struct wrapped<Arr*, enable_if_arithmetic_array_t<Arr>> {
+  MUJOCO_ALWAYS_INLINE
+  static constexpr Arr* unwrap(Eigen::Ref<array_eigen_t<Arr>> wrapper) {
+    return reinterpret_cast<Arr*>(wrapper.data());
+  }
+};
+}  // namespace mujoco::util
+
+namespace mujoco::python {
+namespace _impl {
+template <typename, typename>
+struct py_arg_helper {};
+
+template <typename... PyArg, typename... OmittedArg>
+struct py_arg_helper<std::tuple<PyArg...>, std::tuple<OmittedArg...>> {
+  std::tuple<PyArg...> py_args;
+  std::tuple<OmittedArg...> omitted_args;
+
+  static constexpr int n_py_args = std::tuple_size_v<decltype(py_args)>;
+  static constexpr int n_omitted_args =
+      std::tuple_size_v<decltype(omitted_args)>;
+
+  template <typename... T>
+  MUJOCO_ALWAYS_INLINE
+  constexpr void def(::pybind11::module_& m, T&&... t) {
+    constexpr int NExtras = std::tuple_size_v<std::tuple<T...>>;
+    unpack_tuple_as_py_args<0, NExtras>(m, std::forward<T>(t)...);
+  }
+
+  template <int ArgIdx, int NExtras, typename... T>
+  MUJOCO_ALWAYS_INLINE
+  constexpr void unpack_tuple_as_py_args(::pybind11::module_& m, T&&... t) {
+    if constexpr (ArgIdx == n_py_args) {
+      if constexpr (std::tuple_size_v<std::tuple<T...>> ==
+                    NExtras + n_py_args - n_omitted_args) {
+        m.def(std::forward<T>(t)...);
+      } else {
+        // This should ideally be a static_assert, but we need C++20 consteval
+        // to do that. When using the DEF_WITH_OMITTED_PY_ARGS macro, the
+        // static_assert in that macro would trigger first, rendering this
+        // branch unreachable.
+        throw UnexpectedError(
+            "omitted argument names do not match the underlying function "
+            "parameter names");
+      }
+    } else if (is_omitted<ArgIdx>()) {
+      unpack_tuple_as_py_args<ArgIdx+1, NExtras>(m, std::forward<T>(t)...);
+    } else {
+      unpack_tuple_as_py_args<ArgIdx+1, NExtras>(
+          m, std::forward<T>(t)..., ::pybind11::arg(std::get<ArgIdx>(py_args)));
+    }
+  }
+
+  template <int ArgIdx, int OmittedIdx = 0>
+  MUJOCO_ALWAYS_INLINE
+  constexpr bool is_omitted() {
+    if constexpr (OmittedIdx == n_omitted_args) {
+      return false;
+    // string_view comparison can be constexpr
+    } else if (std::string_view(std::get<ArgIdx>(py_args)) ==
+               std::string_view(std::get<OmittedIdx>(omitted_args))) {
+      return true;
+    } else {
+      return is_omitted<ArgIdx, OmittedIdx + 1>();
+    }
+  }
+};
+}  // namespace _impl
+
+template <typename Tuple>
+MUJOCO_ALWAYS_INLINE
+static constexpr auto WithNamedArgs(Tuple&& py_args) {
+  using ArgTuple = std::remove_cv_t<std::remove_reference_t<Tuple>>;
+  return _impl::py_arg_helper<ArgTuple, std::tuple<>>{
+    std::forward<Tuple>(py_args), std::tuple<>()};
+}
+
+template <typename Tuple1, typename Tuple2>
+MUJOCO_ALWAYS_INLINE
+static constexpr auto WithNamedArgs(Tuple1&& py_args, Tuple2&& omitted_args) {
+  using ArgTuple = std::remove_cv_t<std::remove_reference_t<Tuple1>>;
+  using OmittedTuple = std::remove_cv_t<std::remove_reference_t<Tuple2>>;
+  return _impl::py_arg_helper<ArgTuple, OmittedTuple>{
+    std::forward<Tuple1>(py_args), std::forward<Tuple2>(omitted_args)};
+}
+
+template <typename MjTraits>
+MUJOCO_ALWAYS_INLINE
+static constexpr void Def(::pybind11::module_& m) {
+  WithNamedArgs(MjTraits::param_names).def(
+      m, MjTraits::name,
+      util::UnwrapArgs(InterceptMjErrors(MjTraits::GetFunc())),
+      ::pybind11::doc(MjTraits::doc),
+      ::pybind11::call_guard<::pybind11::gil_scoped_release>());
+}
+
+template <typename MjTraits, typename Func>
+MUJOCO_ALWAYS_INLINE
+static constexpr void Def(::pybind11::module_& m, Func&& func) {
+  WithNamedArgs(MjTraits::param_names).def(
+      m, MjTraits::name, util::UnwrapArgs(std::forward<Func>(func)),
+      ::pybind11::doc(MjTraits::doc),
+      ::pybind11::call_guard<::pybind11::gil_scoped_release>());
+}
+
+template <typename MjTraits, typename OmittedArgs, typename Func>
+MUJOCO_ALWAYS_INLINE
+static constexpr void Def(
+    ::pybind11::module_& m, OmittedArgs&& omitted_args, Func&& func) {
+  WithNamedArgs(MjTraits::param_names, std::forward<OmittedArgs>(omitted_args))
+      .def(m, MjTraits::name, util::UnwrapArgs(std::forward<Func>(func)),
+           ::pybind11::doc(MjTraits::doc),
+           ::pybind11::call_guard<::pybind11::gil_scoped_release>());
+}
+
+// Should only be invoked via the DEF_WITH_OMITTED_PY_ARGS macro.
+template <typename MjTraits, typename OmittedArgsTuple>
+struct DefWithOmittedPyArgsImpl {
+  OmittedArgsTuple omitted_args;
+  template <typename Func>
+  constexpr auto operator()(::pybind11::module_& m, Func&& func) {
+    return Def<MjTraits>(m, omitted_args, std::forward<Func>(func));
+  }
+};
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_FUNCTIONS_H_

python/mujoco/glfw/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2017 The dm_control Authors
+#
+# 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.
+# ==============================================================================
+"""An OpenGL context created via GLFW."""
+
+import glfw
+
+
+class GLContext:
+  """An OpenGL context created via GLFW."""
+
+  def __init__(self, max_width, max_height):
+    glfw.init()
+    glfw.window_hint(glfw.VISIBLE, 0)
+    self._context = glfw.create_window(width=max_width, height=max_height,
+                                       title='Invisible window', monitor=None,
+                                       share=None)
+
+  def make_current(self):
+    glfw.make_context_current(self._context)
+
+  def free(self):
+    if self._context:
+      if glfw.get_current_context() == self._context:
+        glfw.make_context_current(None)
+      glfw.destroy_window(self._context)
+      self._context = None
+
+  def __del__(self):
+    self.free()

python/mujoco/indexer_xmacro.h

@@ -0,0 +1,406 @@
+// 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_PYTHON_INDEXER_XMACRO_H_
+#define MUJOCO_PYTHON_INDEXER_XMACRO_H_
+
+#include <mjxmacro.h>
+
+#define MJMODEL_ACTUATOR                                          \
+  X( int,     actuator_, trntype,      nu, 1                    ) \
+  X( int,     actuator_, dyntype,      nu, 1                    ) \
+  X( int,     actuator_, gaintype,     nu, 1                    ) \
+  X( int,     actuator_, biastype,     nu, 1                    ) \
+  X( int,     actuator_, trnid,        nu, 2                    ) \
+  X( int,     actuator_, group,        nu, 1                    ) \
+  X( mjtByte, actuator_, ctrllimited,  nu, 1                    ) \
+  X( mjtByte, actuator_, forcelimited, nu, 1                    ) \
+  X( mjtNum,  actuator_, dynprm,       nu, mjNDYN               ) \
+  X( mjtNum,  actuator_, gainprm,      nu, mjNGAIN              ) \
+  X( mjtNum,  actuator_, biasprm,      nu, mjNBIAS              ) \
+  X( mjtNum,  actuator_, ctrlrange,    nu, 2                    ) \
+  X( mjtNum,  actuator_, forcerange,   nu, 2                    ) \
+  X( mjtNum,  actuator_, gear,         nu, 6                    ) \
+  X( mjtNum,  actuator_, cranklength,  nu, 1                    ) \
+  X( mjtNum,  actuator_, acc0,         nu, 1                    ) \
+  X( mjtNum,  actuator_, length0,      nu, 1                    ) \
+  X( mjtNum,  actuator_, lengthrange,  nu, 2                    ) \
+  X( mjtNum,  actuator_, user,         nu, MJ_M(nuser_actuator) )
+
+#define MJMODEL_BODY                                        \
+  X( int,     body_, parentid,    nbody, 1                ) \
+  X( int,     body_, rootid,      nbody, 1                ) \
+  X( int,     body_, weldid,      nbody, 1                ) \
+  X( int,     body_, mocapid,     nbody, 1                ) \
+  X( int,     body_, jntnum,      nbody, 1                ) \
+  X( int,     body_, jntadr,      nbody, 1                ) \
+  X( int,     body_, dofnum,      nbody, 1                ) \
+  X( int,     body_, dofadr,      nbody, 1                ) \
+  X( int,     body_, geomnum,     nbody, 1                ) \
+  X( int,     body_, geomadr,     nbody, 1                ) \
+  X( mjtByte, body_, simple,      nbody, 1                ) \
+  X( mjtByte, body_, sameframe,   nbody, 1                ) \
+  X( mjtNum,  body_, pos,         nbody, 3                ) \
+  X( mjtNum,  body_, quat,        nbody, 4                ) \
+  X( mjtNum,  body_, ipos,        nbody, 3                ) \
+  X( mjtNum,  body_, iquat,       nbody, 4                ) \
+  X( mjtNum,  body_, mass,        nbody, 1                ) \
+  X( mjtNum,  body_, subtreemass, nbody, 1                ) \
+  X( mjtNum,  body_, inertia,     nbody, 3                ) \
+  X( mjtNum,  body_, invweight0,  nbody, 2                ) \
+  X( mjtNum,  body_, user,        nbody, MJ_M(nuser_body) )
+
+#define MJMODEL_CAMERA                                   \
+  X( int,    cam_, mode,         ncam, 1               ) \
+  X( int,    cam_, bodyid,       ncam, 1               ) \
+  X( int,    cam_, targetbodyid, ncam, 1               ) \
+  X( mjtNum, cam_, pos,          ncam, 3               ) \
+  X( mjtNum, cam_, quat,         ncam, 4               ) \
+  X( mjtNum, cam_, poscom0,      ncam, 3               ) \
+  X( mjtNum, cam_, pos0,         ncam, 3               ) \
+  X( mjtNum, cam_, mat0,         ncam, 9               ) \
+  X( mjtNum, cam_, fovy,         ncam, 1               ) \
+  X( mjtNum, cam_, ipd,          ncam, 1               ) \
+  X( mjtNum, cam_, user,         ncam, MJ_M(nuser_cam) )
+
+#define MJMODEL_EQUALITY                    \
+  X( int,     eq_, type,   neq, 1         ) \
+  X( int,     eq_, obj1id, neq, 1         ) \
+  X( int,     eq_, obj2id, neq, 1         ) \
+  X( mjtByte, eq_, active, neq, 1         ) \
+  X( mjtNum,  eq_, solref, neq, mjNREF    ) \
+  X( mjtNum,  eq_, solimp, neq, mjNIMP    ) \
+  X( mjtNum,  eq_, data,   neq, mjNEQDATA )
+
+#define MJMODEL_EXCLUDE                      \
+  X( int, exclude_, signature, nexclude, 1 )
+
+#define MJMODEL_GEOM                                        \
+  X( int,     geom_, type,        ngeom, 1                ) \
+  X( int,     geom_, contype,     ngeom, 1                ) \
+  X( int,     geom_, conaffinity, ngeom, 1                ) \
+  X( int,     geom_, condim,      ngeom, 1                ) \
+  X( int,     geom_, bodyid,      ngeom, 1                ) \
+  X( int,     geom_, dataid,      ngeom, 1                ) \
+  X( int,     geom_, matid,       ngeom, 1                ) \
+  X( int,     geom_, group,       ngeom, 1                ) \
+  X( int,     geom_, priority,    ngeom, 1                ) \
+  X( mjtByte, geom_, sameframe,   ngeom, 1                ) \
+  X( mjtNum,  geom_, solmix,      ngeom, 1                ) \
+  X( mjtNum,  geom_, solref,      ngeom, mjNREF           ) \
+  X( mjtNum,  geom_, solimp,      ngeom, mjNIMP           ) \
+  X( mjtNum,  geom_, size,        ngeom, 3                ) \
+  X( mjtNum,  geom_, rbound,      ngeom, 1                ) \
+  X( mjtNum,  geom_, pos,         ngeom, 3                ) \
+  X( mjtNum,  geom_, quat,        ngeom, 4                ) \
+  X( mjtNum,  geom_, friction,    ngeom, 3                ) \
+  X( mjtNum,  geom_, margin,      ngeom, 1                ) \
+  X( mjtNum,  geom_, gap,         ngeom, 1                ) \
+  X( mjtNum,  geom_, user,        ngeom, MJ_M(nuser_geom) ) \
+  X( float,   geom_, rgba,        ngeom, 4                )
+
+#define MJMODEL_HFIELD                       \
+  X( mjtNum, hfield_, size, nhfield,     4 ) \
+  X( int,    hfield_, nrow, nhfield,     1 ) \
+  X( int,    hfield_, ncol, nhfield,     1 ) \
+  X( int,    hfield_, adr,  nhfield,     1 ) \
+  X( float,  hfield_, data, nhfielddata, 1 )
+
+#define MJMODEL_JOINT                                     \
+  X( mjtNum,  ,     qpos0,        nq,   1               ) \
+  X( mjtNum,  ,     qpos_spring,  nq,   1               ) \
+  X( int,     jnt_, type,         njnt, 1               ) \
+  X( int,     jnt_, qposadr,      njnt, 1               ) \
+  X( int,     jnt_, dofadr,       njnt, 1               ) \
+  X( int,     jnt_, group,        njnt, 1               ) \
+  X( mjtByte, jnt_, limited,      njnt, 1               ) \
+  X( mjtNum,  jnt_, pos,          njnt, 3               ) \
+  X( mjtNum,  jnt_, axis,         njnt, 3               ) \
+  X( mjtNum,  jnt_, stiffness,    njnt, 1               ) \
+  X( mjtNum,  jnt_, range,        njnt, 2               ) \
+  X( mjtNum,  jnt_, margin,       njnt, 1               ) \
+  X( mjtNum,  jnt_, user,         njnt, MJ_M(nuser_jnt) ) \
+  X( int,     dof_, bodyid,       nv,   1               ) \
+  X( int,     dof_, jntid,        nv,   1               ) \
+  X( int,     dof_, parentid,     nv,   1               ) \
+  X( int,     dof_, Madr,         nv,   1               ) \
+  X( int,     dof_, simplenum,    nv,   1               ) \
+  X( mjtNum,  dof_, solref,       nv,   mjNREF          ) \
+  X( mjtNum,  dof_, solimp,       nv,   mjNIMP          ) \
+  X( mjtNum,  dof_, frictionloss, nv,   1               ) \
+  X( mjtNum,  dof_, armature,     nv,   1               ) \
+  X( mjtNum,  dof_, damping,      nv,   1               ) \
+  X( mjtNum,  dof_, invweight0,   nv,   1               ) \
+  X( mjtNum,  dof_, M0,           nv,   1               )
+
+#define MJMODEL_LIGHT                             \
+  X( int,     light_, mode,          nlight,  1 ) \
+  X( int,     light_, bodyid,        nlight,  1 ) \
+  X( int,     light_, targetbodyid,  nlight,  1 ) \
+  X( mjtByte, light_, directional,   nlight,  1 ) \
+  X( mjtByte, light_, castshadow,    nlight,  1 ) \
+  X( mjtByte, light_, active,        nlight,  1 ) \
+  X( mjtNum,  light_, pos,           nlight,  3 ) \
+  X( mjtNum,  light_, dir,           nlight,  3 ) \
+  X( mjtNum,  light_, poscom0,       nlight,  3 ) \
+  X( mjtNum,  light_, pos0,          nlight,  3 ) \
+  X( mjtNum,  light_, dir0,          nlight,  3 ) \
+  X( float,   light_, attenuation,   nlight,  3 ) \
+  X( float,   light_, cutoff,        nlight,  1 ) \
+  X( float,   light_, exponent,      nlight,  1 ) \
+  X( float,   light_, ambient,       nlight,  3 ) \
+  X( float,   light_, diffuse,       nlight,  3 ) \
+  X( float,   light_, specular,      nlight,  3 )
+
+#define MJMODEL_MATERIAL                   \
+  X( int,     mat_, texid,       nmat, 1 ) \
+  X( mjtByte, mat_, texuniform,  nmat, 1 ) \
+  X( float,   mat_, texrepeat,   nmat, 2 ) \
+  X( float,   mat_, emission,    nmat, 1 ) \
+  X( float,   mat_, specular,    nmat, 1 ) \
+  X( float,   mat_, shininess,   nmat, 1 ) \
+  X( float,   mat_, reflectance, nmat, 1 ) \
+  X( float,   mat_, rgba,        nmat, 4 )
+
+#define MJMODEL_MESH                   \
+  X( int, mesh_, vertadr,     nmesh, 1 ) \
+  X( int, mesh_, vertnum,     nmesh, 1 ) \
+  X( int, mesh_, texcoordadr, nmesh, 1 ) \
+  X( int, mesh_, faceadr,     nmesh, 1 ) \
+  X( int, mesh_, facenum,     nmesh, 1 ) \
+  X( int, mesh_, graphadr,    nmesh, 1 )
+
+#define MJMODEL_NUMERIC                        \
+  X( int,    numeric_, adr,  nnumeric,     1 ) \
+  X( int,    numeric_, size, nnumeric,     1 ) \
+  X( mjtNum, numeric_, data, nnumericdata, 1 )
+
+#define MJMODEL_PAIR                            \
+  X( int,     pair_, dim,       npair, 1      ) \
+  X( int,     pair_, geom1,     npair, 1      ) \
+  X( int,     pair_, geom2,     npair, 1      ) \
+  X( int,     pair_, signature, npair, 1      ) \
+  X( mjtNum,  pair_, solref,    npair, mjNREF ) \
+  X( mjtNum,  pair_, solimp,    npair, mjNIMP ) \
+  X( mjtNum,  pair_, margin,    npair, 1      ) \
+  X( mjtNum,  pair_, gap,       npair, 1      ) \
+  X( mjtNum,  pair_, friction,  npair, 5      )
+
+#define MJMODEL_SENSOR                                         \
+  X( int,    sensor_, type,      nsensor, 1                  ) \
+  X( int,    sensor_, datatype,  nsensor, 1                  ) \
+  X( int,    sensor_, needstage, nsensor, 1                  ) \
+  X( int,    sensor_, objtype,   nsensor, 1                  ) \
+  X( int,    sensor_, objid,     nsensor, 1                  ) \
+  X( int,    sensor_, reftype,   nsensor, 1                  ) \
+  X( int,    sensor_, refid,     nsensor, 1                  ) \
+  X( int,    sensor_, dim,       nsensor, 1                  ) \
+  X( int,    sensor_, adr,       nsensor, 1                  ) \
+  X( mjtNum, sensor_, cutoff,    nsensor, 1                  ) \
+  X( mjtNum, sensor_, noise,     nsensor, 1                  ) \
+  X( mjtNum, sensor_, user,      nsensor, MJ_M(nuser_sensor) )
+
+#define MJMODEL_SITE                                      \
+  X( int,     site_, type,      nsite, 1                ) \
+  X( int,     site_, bodyid,    nsite, 1                ) \
+  X( int,     site_, matid,     nsite, 1                ) \
+  X( int,     site_, group,     nsite, 1                ) \
+  X( mjtByte, site_, sameframe, nsite, 1                ) \
+  X( mjtNum,  site_, size,      nsite, 3                ) \
+  X( mjtNum,  site_, pos,       nsite, 3                ) \
+  X( mjtNum,  site_, quat,      nsite, 4                ) \
+  X( mjtNum,  site_, user,      nsite, MJ_M(nuser_site) ) \
+  X( float,   site_, rgba,      nsite, 4                )
+
+#define MJMODEL_SKIN                       \
+  X( int,   skin_, matid,       nskin, 1 ) \
+  X( float, skin_, rgba,        nskin, 4 ) \
+  X( float, skin_, inflate,     nskin, 1 ) \
+  X( int,   skin_, vertadr,     nskin, 1 ) \
+  X( int,   skin_, vertnum,     nskin, 1 ) \
+  X( int,   skin_, texcoordadr, nskin, 1 ) \
+  X( int,   skin_, faceadr,     nskin, 1 ) \
+  X( int,   skin_, facenum,     nskin, 1 ) \
+  X( int,   skin_, boneadr,     nskin, 1 ) \
+  X( int,   skin_, bonenum,     nskin, 1 )
+
+#define MJMODEL_TENDON                                             \
+  X( int,     tendon, _adr,          ntendon, 1                  ) \
+  X( int,     tendon, _num,          ntendon, 1                  ) \
+  X( int,     tendon, _matid,        ntendon, 1                  ) \
+  X( int,     tendon, _group,        ntendon, 1                  ) \
+  X( mjtByte, tendon, _limited,      ntendon, 1                  ) \
+  X( mjtNum,  tendon, _width,        ntendon, 1                  ) \
+  X( mjtNum,  tendon, _solref_lim,   ntendon, mjNREF             ) \
+  X( mjtNum,  tendon, _solimp_lim,   ntendon, mjNIMP             ) \
+  X( mjtNum,  tendon, _solref_fri,   ntendon, mjNREF             ) \
+  X( mjtNum,  tendon, _solimp_fri,   ntendon, mjNIMP             ) \
+  X( mjtNum,  tendon, _range,        ntendon, 2                  ) \
+  X( mjtNum,  tendon, _margin,       ntendon, 1                  ) \
+  X( mjtNum,  tendon, _stiffness,    ntendon, 1                  ) \
+  X( mjtNum,  tendon, _damping,      ntendon, 1                  ) \
+  X( mjtNum,  tendon, _frictionloss, ntendon, 1                  ) \
+  X( mjtNum,  tendon, _lengthspring, ntendon, 1                  ) \
+  X( mjtNum,  tendon, _length0,      ntendon, 1                  ) \
+  X( mjtNum,  tendon, _invweight0,   ntendon, 1                  ) \
+  X( mjtNum,  tendon, _user,         ntendon, MJ_M(nuser_tendon) ) \
+  X( float,   tendon, _rgba,         ntendon, 4                  )
+
+#define MJMODEL_TEXTURE                   \
+  X( int,     tex_, type,   ntex,     1 ) \
+  X( int,     tex_, height, ntex,     1 ) \
+  X( int,     tex_, width,  ntex,     1 ) \
+  X( int,     tex_, adr,    ntex,     1 ) \
+  X( mjtByte, tex_, rgb,    ntexdata, 1 )
+
+#define MJMODEL_TUPLE                         \
+  X( int,    tuple_, adr,     ntuple,     1 ) \
+  X( int,    tuple_, size,    ntuple,     1 ) \
+  X( int,    tuple_, objtype, ntupledata, 1 ) \
+  X( int,    tuple_, objid,   ntupledata, 1 ) \
+  X( mjtNum, tuple_, objprm,  ntupledata, 1 )
+
+#define MJMODEL_KEYFRAME                         \
+  X( mjtNum, key_, time,  nkey, 1              ) \
+  X( mjtNum, key_, qpos,  nkey, MJ_M(nq)       ) \
+  X( mjtNum, key_, qvel,  nkey, MJ_M(nv)       ) \
+  X( mjtNum, key_, act,   nkey, MJ_M(na)       ) \
+  X( mjtNum, key_, mpos,  nkey, MJ_M(nmocap)*3 ) \
+  X( mjtNum, key_, mquat, nkey, MJ_M(nmocap)*4 )
+
+#define MJMODEL_VIEW_GROUPS                                            \
+  XGROUP( MjModelActuatorViews, actuator, nu,       MJMODEL_ACTUATOR ) \
+  XGROUP( MjModelBodyViews,     body,     nbody,    MJMODEL_BODY     ) \
+  XGROUP( MjModelCameraViews,   cam,      ncam,     MJMODEL_CAMERA   ) \
+  XGROUP( MjModelEqualityViews, eq,       neq,      MJMODEL_EQUALITY ) \
+  XGROUP( MjModelExcludeViews,  exclude,  nexclude, MJMODEL_EXCLUDE  ) \
+  XGROUP( MjModelGeomViews,     geom,     ngeom,    MJMODEL_GEOM     ) \
+  XGROUP( MjModelHfieldViews,   hfield,   nhfield,  MJMODEL_HFIELD   ) \
+  XGROUP( MjModelJointViews,    jnt,      njnt,     MJMODEL_JOINT    ) \
+  XGROUP( MjModelLightViews,    light,    nlight,   MJMODEL_LIGHT    ) \
+  XGROUP( MjModelMaterialViews, mat,      nmat,     MJMODEL_MATERIAL ) \
+  XGROUP( MjModelMeshViews,     mesh,     nmesh,    MJMODEL_MESH     ) \
+  XGROUP( MjModelNumericViews,  numeric,  nnumeric, MJMODEL_NUMERIC  ) \
+  XGROUP( MjModelPairViews,     pair,     npair,    MJMODEL_PAIR     ) \
+  XGROUP( MjModelSensorViews,   sensor,   nsensor,  MJMODEL_SENSOR   ) \
+  XGROUP( MjModelSiteViews,     site,     nsite,    MJMODEL_SITE     ) \
+  XGROUP( MjModelSkinViews,     skin,     nskin,    MJMODEL_SKIN     ) \
+  XGROUP( MjModelTendonViews,   tendon,   ntendon,  MJMODEL_TENDON   ) \
+  XGROUP( MjModelTextureViews,  tex,      ntex,     MJMODEL_TEXTURE  ) \
+  XGROUP( MjModelTupleViews,    tuple,    ntuple,   MJMODEL_TUPLE    ) \
+  XGROUP( MjModelKeyframeViews, key,      nkey,     MJMODEL_KEYFRAME )
+
+#define MJMODEL_VIEW_GROUPS_ALTNAMES        \
+  XGROUP( cam, camera,   MJMODEL_CAMERA   ) \
+  XGROUP( eq,  equality, MJMODEL_EQUALITY ) \
+  XGROUP( jnt, joint,    MJMODEL_JOINT    ) \
+  XGROUP( mat, material, MJMODEL_MATERIAL ) \
+  XGROUP( tex, texture,  MJMODEL_TEXTURE  ) \
+  XGROUP( key, keyframe, MJMODEL_KEYFRAME )
+
+#define MJDATA_ACTUATOR                          \
+  X( mjtNum, ,          ctrl,     nu, 1        ) \
+  X( mjtNum, actuator_, length,   nu, 1        ) \
+  X( mjtNum, actuator_, moment,   nu, MJ_M(nv) ) \
+  X( mjtNum, actuator_, velocity, nu, 1        ) \
+  X( mjtNum, actuator_, force,    nu, 1        )
+
+#define MJDATA_BODY                        \
+  X( mjtNum, , xfrc_applied,   nbody, 6  ) \
+  X( mjtNum, , xpos,           nbody, 3  ) \
+  X( mjtNum, , xquat,          nbody, 4  ) \
+  X( mjtNum, , xmat,           nbody, 9  ) \
+  X( mjtNum, , xipos,          nbody, 3  ) \
+  X( mjtNum, , ximat,          nbody, 9  ) \
+  X( mjtNum, , subtree_com,    nbody, 3  ) \
+  X( mjtNum, , cinert,         nbody, 10 ) \
+  X( mjtNum, , crb,            nbody, 10 ) \
+  X( mjtNum, , cvel,           nbody, 6  ) \
+  X( mjtNum, , subtree_linvel, nbody, 3  ) \
+  X( mjtNum, , subtree_angmom, nbody, 3  ) \
+  X( mjtNum, , cacc,           nbody, 6  ) \
+  X( mjtNum, , cfrc_int,       nbody, 6  ) \
+  X( mjtNum, , cfrc_ext,       nbody, 6  )
+
+#define MJDATA_CAMERA              \
+  X( mjtNum, cam_, xpos, ncam, 3 ) \
+  X( mjtNum, cam_, xmat, ncam, 9 )
+
+#define MJDATA_GEOM                  \
+  X( mjtNum, geom_, xpos, ngeom, 3 ) \
+  X( mjtNum, geom_, xmat, ngeom, 9 )
+
+#define MJDATA_JOINT                                \
+  X( mjtNum, , qpos,            nq,   1           ) \
+  X( mjtNum, , qvel,            nv,   1           ) \
+  X( mjtNum, , qacc_warmstart,  nv,   1           ) \
+  X( mjtNum, , qfrc_applied,    nv,   1           ) \
+  X( mjtNum, , qacc,            nv,   1           ) \
+  X( mjtNum, , xanchor,         njnt, 3           ) \
+  X( mjtNum, , xaxis,           njnt, 3           ) \
+  X( mjtNum, , cdof,            nv,   6           ) \
+  X( mjtNum, , qLDiagInv,       nv,   1           ) \
+  X( mjtNum, , qLDiagSqrtInv,   nv,   1           ) \
+  X( int,    , efc_JT_rownnz,   nv,   1           ) \
+  X( int,    , efc_JT_rowadr,   nv,   1           ) \
+  X( int,    , efc_JT_rowsuper, nv,   1           ) \
+  X( int,    , efc_JT_colind,   nv,   MJ_M(njmax) ) \
+  X( mjtNum, , efc_JT,          nv,   MJ_M(njmax) ) \
+  X( mjtNum, , cdof_dot,        nv,   6           ) \
+  X( mjtNum, , qfrc_bias,       nv,   1           ) \
+  X( mjtNum, , qfrc_passive,    nv,   1           ) \
+  X( mjtNum, , qfrc_actuator,   nv,   1           ) \
+  X( mjtNum, , qfrc_unc,        nv,   1           ) \
+  X( mjtNum, , qacc_unc,        nv,   1           ) \
+  X( mjtNum, , qfrc_constraint, nv,   1           ) \
+  X( mjtNum, , qfrc_inverse,    nv,   1           )
+
+#define MJDATA_LIGHT                   \
+  X( mjtNum, light_, xpos, nlight, 3 ) \
+  X( mjtNum, light_, xdir, nlight, 3 )
+
+#define MJDATA_SENSOR                       \
+  X( mjtNum, sensor, data, nsensordata, 1 )
+
+#define MJDATA_SITE                  \
+  X( mjtNum, site_, xpos, nsite, 3 ) \
+  X( mjtNum, site_, xmat, nsite, 9 )
+
+#define MJDATA_TENDON                            \
+  X( int, ten_,    wrapadr , ntendon, 1        ) \
+  X( int, ten_,    wrapnum , ntendon, 1        ) \
+  X( int, ten_,    J_rownnz, ntendon, 1        ) \
+  X( int, ten_,    J_rowadr, ntendon, 1        ) \
+  X( int, ten_,    J_colind, ntendon, MJ_M(nv) ) \
+  X( mjtNum, ten_, length  , ntendon, 1        ) \
+  X( mjtNum, ten_, J       , ntendon, MJ_M(nv) ) \
+  X( mjtNum, ten_, velocity, ntendon, 1        )
+
+#define MJDATA_VIEW_GROUPS                                           \
+  XGROUP( MjDataActuatorViews, actuator, nu,       MJDATA_ACTUATOR ) \
+  XGROUP( MjDataBodyViews,     body,     nbody,    MJDATA_BODY     ) \
+  XGROUP( MjDataCameraViews,   cam,      ncam,     MJDATA_CAMERA   ) \
+  XGROUP( MjDataGeomViews,     geom,     ngeom,    MJDATA_GEOM     ) \
+  XGROUP( MjDataJointViews,    jnt,      njnt,     MJDATA_JOINT    ) \
+  XGROUP( MjDataLightViews,    light,    nlight,   MJDATA_LIGHT    ) \
+  XGROUP( MjDataSensorViews,   sensor,   nsensor,  MJDATA_SENSOR   ) \
+  XGROUP( MjDataSiteViews,     site,     nsite,    MJDATA_SITE     ) \
+  XGROUP( MjDataTendonViews,   tendon,   ntendon,  MJDATA_TENDON   )
+
+#define MJDATA_VIEW_GROUPS_ALTNAMES       \
+  XGROUP( cam,    camera, MJDATA_CAMERA ) \
+  XGROUP( jnt,    joint,  MJDATA_JOINT  ) \
+  XGROUP( tendon, ten,    MJDATA_TENDON )
+
+#endif  // MUJOCO_PYTHON_INDEXER_XMACRO_H_

python/mujoco/indexers.cc

@@ -0,0 +1,366 @@
+// 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 <memory>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "errors.h"
+#include "indexers.h"
+#include "mjdata_meta.h"
+#include "raw.h"
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+
+namespace py = ::pybind11;
+
+namespace {
+// Parses raw mjModel to create a mapping from names to array indices.
+//
+// Args:
+//   count: Number of names in the map.
+//   name_offsets: Array consisting of indices that correspond to the start of
+//     each name in the `names` array.
+//   names: Character arrays consisting of the concatenation of all names.
+template <typename IntPtr, typename CharPtr>
+NameToIDMap MakeMap(int count, IntPtr name_offsets, CharPtr names)
+{
+  NameToIDMap name_to_id;
+  for (int index = 0; index < count; ++index) {
+    const char* name = &names[name_offsets[index]];
+    if (name[0] != '\0') {
+      name_to_id.insert({name, index});
+    }
+  }
+  return name_to_id;
+}
+
+// Makes an array view into an mjModel/mjData struct field at a given index.
+//
+// Template args:
+//   MjSize: An (int M::*) specifying the number of entities of the same type
+//     as the one whose array view is to be created. This is a way to indicate
+//     the MuJoCo category of the entity itself, e.g. nbody indicates that the
+//     field belongs to a body.
+//   T: Scalar data type of the field.
+//   M: Either raw::MjModel or MjDataMetadata.
+//
+// Args:
+//   base_ptr: Pointer to the first entry in the entire field.
+//   index: Entity ID at which a view is to be created.
+//   shape: The shape of the array as defined by MuJoCo. If MjSize corresponds
+//     to a ragged or indirect field, the shape may be prepended by additional
+//     dimensions as appropriate, e.g. if MjSize is nv then the view gains an
+//     additional dimension of the size len(qvel) of the particular joint.
+//   m: Used for dereferencing MjSize.
+//   owner: The base object whose lifetime is tied to the returned array.
+template <auto MjSize, typename T, typename M>
+py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
+                         const M& m, py::handle owner) {
+  int offset;
+  if (MjSize == &M::nq) {
+    offset = m.jnt_qposadr[index];
+    shape.insert(
+        shape.begin(),
+        ((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
+  } else if (MjSize == &M::nv) {
+    offset = m.jnt_dofadr[index];
+    shape.insert(
+        shape.begin(),
+        ((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
+  } else if (MjSize == &M::nhfielddata) {
+    offset = m.hfield_adr[index];
+    shape.insert(shape.begin(), m.hfield_ncol[index]);
+    shape.insert(shape.begin(), m.hfield_nrow[index]);
+  } else if (MjSize == &M::ntexdata) {
+    offset = m.tex_adr[index];
+    shape.insert(shape.begin(), m.tex_width[index]);
+    shape.insert(shape.begin(), m.tex_height[index]);
+  } else if (MjSize == &M::nsensordata) {
+    offset = m.sensor_adr[index];
+    shape.insert(shape.begin(), m.sensor_dim[index]);
+  } else if (MjSize == &M::nnumericdata) {
+    offset = m.numeric_adr[index];
+    shape.insert(shape.begin(), m.numeric_size[index]);
+  } else if (MjSize == &M::ntupledata) {
+    offset = m.tuple_adr[index];
+    shape.insert(shape.begin(), m.tuple_size[index]);
+  } else {
+    int size = 1;
+    for (int s : shape) {
+      size *= s;
+    }
+    offset = index * size;
+  }
+  return py::array_t<T>(std::move(shape), base_ptr + offset,
+                        py::reinterpret_borrow<py::object>(owner));
+}
+}  // namespace
+
+// M is either a raw::MjModel or MjDataMetadata.
+template <typename M>
+NameToIDMaps::NameToIDMaps(const M& m)
+    : body(MakeMap(m.nbody, m.name_bodyadr, m.names)),
+      jnt(MakeMap(m.njnt, m.name_jntadr, m.names)),
+      geom(MakeMap(m.ngeom, m.name_geomadr, m.names)),
+      site(MakeMap(m.nsite, m.name_siteadr, m.names)),
+      cam(MakeMap(m.ncam, m.name_camadr, m.names)),
+      light(MakeMap(m.nlight, m.name_lightadr, m.names)),
+      mesh(MakeMap(m.nmesh, m.name_meshadr, m.names)),
+      skin(MakeMap(m.nskin, m.name_skinadr, m.names)),
+      hfield(MakeMap(m.nhfield, m.name_hfieldadr, m.names)),
+      tex(MakeMap(m.ntex, m.name_texadr, m.names)),
+      mat(MakeMap(m.nmat, m.name_matadr, m.names)),
+      pair(MakeMap(m.npair, m.name_pairadr, m.names)),
+      exclude(MakeMap(m.nexclude, m.name_excludeadr, m.names)),
+      eq(MakeMap(m.neq, m.name_eqadr, m.names)),
+      tendon(MakeMap(m.ntendon, m.name_tendonadr, m.names)),
+      actuator(MakeMap(m.nu, m.name_actuatoradr, m.names)),
+      sensor(MakeMap(m.nsensor, m.name_sensoradr, m.names)),
+      numeric(MakeMap(m.nnumeric, m.name_numericadr, m.names)),
+      text(MakeMap(m.ntext, m.name_textadr, m.names)),
+      tuple(MakeMap(m.ntuple, m.name_tupleadr, m.names)),
+      key(MakeMap(m.nkey, m.name_keyadr, m.names)) {}
+
+MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner)
+    : m_(m),
+      owner_(owner),
+      name_to_id_(*m)
+#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
+  , field##_(m->nfield, std::nullopt)
+      MJMODEL_VIEW_GROUPS
+#undef XGROUP
+{}
+
+#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjModelFieldGroupedViews& MjModelIndexer::field(int i) {             \
+    if (i > field##_.size()) {                                         \
+      throw py::index_error("index out of range");                     \
+    }                                                                  \
+    auto& indexer = field##_[i];                                       \
+    if (!indexer.has_value()) {                                        \
+      indexer.emplace(i, m_, owner_);                                  \
+    }                                                                  \
+    return *indexer;                                                   \
+  }
+MJMODEL_VIEW_GROUPS
+#undef XGROUP
+
+#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjModelFieldGroupedViews& MjModelIndexer::field##_by_name(           \
+      std::string_view name) {                                         \
+    try {                                                              \
+      return field(name_to_id_.field.at(name));                        \
+    } catch (...) {                                                    \
+      throw py::key_error(std::string(name));                          \
+    }                                                                  \
+  }
+MJMODEL_VIEW_GROUPS
+#undef XGROUP
+
+MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
+                             py::handle owner)
+    : d_(d),
+      m_(m),
+      owner_(owner),
+      name_to_id_(*m)
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  , field##_(m->nfield, std::nullopt)
+      MJDATA_VIEW_GROUPS
+#undef XGROUP
+{}
+
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjDataGroupedViews& MjDataIndexer::field(int i) {              \
+    if (i > field##_.size()) {                                   \
+      throw py::index_error("index out of range");               \
+    }                                                            \
+    auto& indexer = field##_[i];                                 \
+    if (!indexer.has_value()) {                                  \
+      indexer.emplace(i, d_, m_, owner_);                        \
+    }                                                            \
+    return *indexer;                                             \
+  }
+MJDATA_VIEW_GROUPS
+#undef XGROUP
+
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjDataGroupedViews& MjDataIndexer::field##_by_name(            \
+      std::string_view name) {                                   \
+    try {                                                        \
+      return field(name_to_id_.field.at(name));                  \
+    } catch (...) {                                              \
+      throw py::key_error(std::string(name));                    \
+    }                                                            \
+  }
+MJDATA_VIEW_GROUPS
+#undef XGROUP
+
+#define MAKE_SHAPE(dim)                                              \
+  [n = (dim)]() -> std::vector<int> {                                \
+    if constexpr (std::string_view(#dim) == std::string_view("1")) { \
+      return {};                                                     \
+    } else {                                                         \
+      return {n};                                                    \
+    }                                                                \
+  }()
+
+#undef MJ_M
+#define MJ_M(n) m_->n
+#define X(type, prefix, var, dim0, dim1)                            \
+  py::array_t<type> XGROUP::var() {                                 \
+    if (!var##_.has_value()) {                                      \
+      var##_.emplace(MakeArray<&raw::MjModel::dim0>(                \
+          m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
+    }                                                               \
+    return *var##_;                                                 \
+  }
+
+#define XGROUP MjModelActuatorViews
+MJMODEL_ACTUATOR
+#undef XGROUP
+
+#define XGROUP MjModelBodyViews
+MJMODEL_BODY
+#undef XGROUP
+
+#define XGROUP MjModelCameraViews
+MJMODEL_CAMERA
+#undef XGROUP
+
+#define XGROUP MjModelEqualityViews
+MJMODEL_EQUALITY
+#undef XGROUP
+
+#define XGROUP MjModelExcludeViews
+MJMODEL_EXCLUDE
+#undef XGROUP
+
+#define XGROUP MjModelGeomViews
+MJMODEL_GEOM
+#undef XGROUP
+
+#define XGROUP MjModelHfieldViews
+MJMODEL_HFIELD
+#undef XGROUP
+
+#define XGROUP MjModelJointViews
+MJMODEL_JOINT
+#undef XGROUP
+
+#define XGROUP MjModelLightViews
+MJMODEL_LIGHT
+#undef XGROUP
+
+#define XGROUP MjModelMaterialViews
+MJMODEL_MATERIAL
+#undef XGROUP
+
+#define XGROUP MjModelMeshViews
+MJMODEL_MESH
+#undef XGROUP
+
+#define XGROUP MjModelNumericViews
+MJMODEL_NUMERIC
+#undef XGROUP
+
+#define XGROUP MjModelPairViews
+MJMODEL_PAIR
+#undef XGROUP
+
+#define XGROUP MjModelSensorViews
+MJMODEL_SENSOR
+#undef XGROUP
+
+#define XGROUP MjModelSiteViews
+MJMODEL_SITE
+#undef XGROUP
+
+#define XGROUP MjModelSkinViews
+MJMODEL_SKIN
+#undef XGROUP
+
+#define XGROUP MjModelTendonViews
+MJMODEL_TENDON
+#undef XGROUP
+
+#define XGROUP MjModelTextureViews
+MJMODEL_TEXTURE
+#undef XGROUP
+
+#define XGROUP MjModelTupleViews
+MJMODEL_TUPLE
+#undef XGROUP
+
+#define XGROUP MjModelKeyframeViews
+MJMODEL_KEYFRAME
+#undef XGROUP
+
+#undef X
+#define MJ_M(n) m_->n
+#define X(type, prefix, var, dim0, dim1)                            \
+  py::array_t<type> XGROUP::var() {                               \
+    if (!var##_.has_value()) {                                      \
+      var##_.emplace(MakeArray<&MjDataMetadata::dim0>(              \
+          d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
+    }                                                               \
+    return *var##_;                                                 \
+  }
+
+#define XGROUP MjDataActuatorViews
+MJDATA_ACTUATOR
+#undef XGROUP
+
+#define XGROUP MjDataBodyViews
+MJDATA_BODY
+#undef XGROUP
+
+#define XGROUP MjDataCameraViews
+MJDATA_CAMERA
+#undef XGROUP
+
+#define XGROUP MjDataGeomViews
+MJDATA_GEOM
+#undef XGROUP
+
+#define XGROUP MjDataJointViews
+MJDATA_JOINT
+#undef XGROUP
+
+#define XGROUP MjDataLightViews
+MJDATA_LIGHT
+#undef XGROUP
+
+#define XGROUP MjDataSensorViews
+MJDATA_SENSOR
+#undef XGROUP
+
+#define XGROUP MjDataSiteViews
+MJDATA_SITE
+#undef XGROUP
+
+#define XGROUP MjDataTendonViews
+MJDATA_TENDON
+#undef XGROUP
+
+#undef X
+#undef MJ_M
+#define MJ_M(n) n
+
+}  // namespace mujoco::python

python/mujoco/indexers.h

@@ -0,0 +1,190 @@
+// 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_PYTHON_INDEXERS_H_
+#define MUJOCO_PYTHON_INDEXERS_H_
+
+#include <optional>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include <absl/container/flat_hash_map.h>
+#include <mjxmacro.h>
+#include "indexer_xmacro.h"
+#include "mjdata_meta.h"
+#include "raw.h"
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+using NameToIDMap = absl::flat_hash_map<std::string, int>;
+
+struct NameToIDMaps {
+  // M is either a raw::MjModel or MjDataMetadata.
+  template <typename M>
+  explicit NameToIDMaps(const M& m);
+
+  NameToIDMap body;
+  NameToIDMap jnt;
+  NameToIDMap geom;
+  NameToIDMap site;
+  NameToIDMap cam;
+  NameToIDMap light;
+  NameToIDMap mesh;
+  NameToIDMap skin;
+  NameToIDMap hfield;
+  NameToIDMap tex;
+  NameToIDMap mat;
+  NameToIDMap pair;
+  NameToIDMap exclude;
+  NameToIDMap eq;
+  NameToIDMap tendon;
+  NameToIDMap actuator;
+  NameToIDMap sensor;
+  NameToIDMap numeric;
+  NameToIDMap text;
+  NameToIDMap tuple;
+  NameToIDMap key;
+};
+
+
+// Base class for a collection of NumPy views into mjModel fields associated
+// with the same "entity" (for example, fields corresponding to a particular
+// geom, or a particular joint).
+class MjModelGroupedViewsBase {
+ public:
+  MjModelGroupedViewsBase(int index, raw::MjModel* m,
+                          pybind11::handle owner)
+      : index_(index), m_(m), owner_(owner) {}
+
+ protected:
+  int index_;
+  raw::MjModel* m_;
+  pybind11::handle owner_;
+};
+
+#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
+  class MjModelGroupedViews : public MjModelGroupedViewsBase {    \
+   public:                                                        \
+    using MjModelGroupedViewsBase::MjModelGroupedViewsBase;       \
+    FIELD_XMACROS                                                 \
+  };
+
+// Lazily instantiate a NumPy array view object when accessed from Python, but
+// cache it once made so that we can return the same one if requested again.
+#define X(type, prefix, var, dim0, dim1)         \
+  pybind11::array_t<type> var();                 \
+  std::optional<pybind11::array_t<type>> var##_;
+
+MJMODEL_VIEW_GROUPS
+
+#undef XGROUP
+#undef X
+
+// A class for accessing fields corresponding to a single mjModel "entity"
+// (e.g. a particular geom or joint) either by name or by ID.
+class MjModelIndexer {
+ public:
+  MjModelIndexer(raw::MjModel* m, pybind11::handle owner);
+
+#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjModelGroupedViews& field(int i);                              \
+  MjModelGroupedViews& field##_by_name(std::string_view name);
+
+  MJMODEL_VIEW_GROUPS
+#undef XGROUP
+
+ private:
+  raw::MjModel* m_;
+  pybind11::handle owner_;
+  NameToIDMaps name_to_id_;
+
+  // Lazily instantiate a grouped views object when accessed from Python, but
+  // cache it once made so that we can return the same one if requested again.
+#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
+  std::vector<std::optional<MjModelGroupedViews>> field##_;
+
+  MJMODEL_VIEW_GROUPS
+#undef XGROUP
+};
+
+// Base class for a collection of NumPy views into mjData fields associated
+// with the same "entity" (for example, fields corresponding to a particular
+// geom, or a particular joint).
+class MjDataGroupedViewsBase {
+ public:
+  MjDataGroupedViewsBase(int index, raw::MjData* d,
+                         const MjDataMetadata* m,
+                         pybind11::handle owner)
+      : index_(index), d_(d), m_(m), owner_(owner) {}
+
+ protected:
+  int index_;
+  raw::MjData* d_;
+  const MjDataMetadata* m_;
+  pybind11::handle owner_;
+};
+
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  class MjDataGroupedViews : public MjDataGroupedViewsBase {     \
+   public:                                                       \
+    using MjDataGroupedViewsBase::MjDataGroupedViewsBase;        \
+    FIELD_XMACROS                                                \
+  };
+
+// Lazily instantiate a NumPy array view object when accessed from Python, but
+// cache it once made so that we can return the same one if requested again.
+#define X(type, prefix, var, dim0, dim1)         \
+  pybind11::array_t<type> var();                 \
+  std::optional<pybind11::array_t<type>> var##_;
+
+MJDATA_VIEW_GROUPS
+
+#undef XGROUP
+#undef X
+
+// A class for accessing fields corresponding to a single mjModel "entity"
+// (e.g. a particular geom or joint) either by name or by ID.
+class MjDataIndexer {
+ public:
+  MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
+                pybind11::handle owner);
+
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  MjDataGroupedViews& field(int i);                              \
+  MjDataGroupedViews& field##_by_name(std::string_view name);
+
+  MJDATA_VIEW_GROUPS
+#undef XGROUP
+
+ private:
+  raw::MjData* d_;
+  const MjDataMetadata* m_;
+  pybind11::handle owner_;
+  NameToIDMaps name_to_id_;
+
+  // Lazily instantiate a grouped views object when accessed from Python, but
+  // cache it once made so that we can return the same one if requested again.
+#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
+  std::vector<std::optional<MjDataGroupedViews>> field##_;
+
+  MJDATA_VIEW_GROUPS
+#undef XGROUP
+};
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_INDEXERS_H_

python/mujoco/mjdata_meta.h

@@ -0,0 +1,103 @@
+// 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_PYTHON_MJDATA_META_H_
+#define MUJOCO_PYTHON_MJDATA_META_H_
+
+#include <mjxmacro.h>
+#include "raw.h"
+
+namespace mujoco::python {
+namespace _impl {
+template <typename T> class MjWrapper;
+}  // namespace _impl
+
+#define MJDATA_METADATA                 \
+  X( int,  jnt_qposadr,      njnt     ) \
+  X( int,  jnt_dofadr,       njnt     ) \
+  X( int,  hfield_nrow,      nhfield  ) \
+  X( int,  hfield_ncol,      nhfield  ) \
+  X( int,  hfield_adr,       nhfield  ) \
+  X( int,  tex_height,       ntex     ) \
+  X( int,  tex_width,        ntex     ) \
+  X( int,  tex_adr,          ntex     ) \
+  X( int,  sensor_dim,       nsensor  ) \
+  X( int,  sensor_adr,       nsensor  ) \
+  X( int,  numeric_adr,      nnumeric ) \
+  X( int,  numeric_size,     nnumeric ) \
+  X( int,  tuple_adr,        ntuple   ) \
+  X( int,  tuple_size,       ntuple   ) \
+  X( int,  name_bodyadr,     nbody    ) \
+  X( int,  name_jntadr,      njnt     ) \
+  X( int,  name_geomadr,     ngeom    ) \
+  X( int,  name_siteadr,     nsite    ) \
+  X( int,  name_camadr,      ncam     ) \
+  X( int,  name_lightadr,    nlight   ) \
+  X( int,  name_meshadr,     nmesh    ) \
+  X( int,  name_skinadr,     nskin    ) \
+  X( int,  name_hfieldadr,   nhfield  ) \
+  X( int,  name_texadr,      ntex     ) \
+  X( int,  name_matadr,      nmat     ) \
+  X( int,  name_pairadr,     npair    ) \
+  X( int,  name_excludeadr,  nexclude ) \
+  X( int,  name_eqadr,       neq      ) \
+  X( int,  name_tendonadr,   ntendon  ) \
+  X( int,  name_actuatoradr, nu       ) \
+  X( int,  name_sensoradr,   nsensor  ) \
+  X( int,  name_numericadr,  nnumeric ) \
+  X( int,  name_textadr,     ntext    ) \
+  X( int,  name_tupleadr,    ntuple   ) \
+  X( int,  name_keyadr,      nkey     ) \
+  X( char, names,            nnames   )
+
+// A subset of mjModel fields that are required to reconstruct an MjDataWrapper.
+struct MjDataMetadata {
+ public:
+  friend class _impl::MjWrapper<raw::MjData>;
+
+#define X(var) int var;
+  MJMODEL_INTS
+#undef X
+#define X(type, var, n) std::shared_ptr<type[]> var;
+  MJDATA_METADATA
+#undef X
+
+ private:
+  MjDataMetadata() = default;
+  MjDataMetadata(const MjDataMetadata& other) = default;
+  MjDataMetadata(MjDataMetadata&& other) = default;
+  explicit MjDataMetadata(const raw::MjModel* m)
+      :
+#define X(var) var(m->var),
+        MJMODEL_INTS
+#undef X
+
+#define X(dtype, var, n)                            \
+  var(                                              \
+      [](dtype* src, int len) {                     \
+        dtype* dst = new dtype[len];                \
+        std::memcpy(dst, src, len * sizeof(dtype)); \
+        return dst;                                 \
+      }(m->var, m->n)),
+
+        MJDATA_METADATA
+#undef X
+        dummy_() {}
+
+  bool dummy_;  // Dummy variable to terminate X macro sequences.
+};
+
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_MJDATA_META_H_

python/mujoco/osmesa/__init__.py

@@ -0,0 +1,81 @@
+# Copyright 2018 The dm_control Authors
+#
+# 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.
+# ==============================================================================
+"""An OSMesa context for software-based OpenGL rendering."""
+
+import os
+
+PYOPENGL_PLATFORM = os.environ.get('PYOPENGL_PLATFORM')
+
+if not PYOPENGL_PLATFORM:
+  os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
+elif PYOPENGL_PLATFORM.lower() != 'osmesa':
+  raise ImportError(
+      'Cannot use OSMesa rendering platform. '
+      'The PYOPENGL_PLATFORM environment variable is set to {!r} '
+      '(should be either unset or \'osmesa\').'
+      .format(PYOPENGL_PLATFORM))
+
+# pylint: disable=g-import-not-at-top
+from OpenGL import GL
+from OpenGL import osmesa
+from OpenGL.GL import arrays
+
+_DEPTH_BITS = 24
+_STENCIL_BITS = 8
+_ACCUM_BITS = 0
+
+
+class GLContext:
+  """An OSMesa context for software-based OpenGL rendering."""
+
+  def __init__(self, max_width, max_height):
+    """Initializes this OSMesa context."""
+    self._context = osmesa.OSMesaCreateContextExt(
+        osmesa.OSMESA_RGBA,
+        _DEPTH_BITS,
+        _STENCIL_BITS,
+        _ACCUM_BITS,
+        None,  # sharelist
+    )
+    if not self._context:
+      raise RuntimeError('Failed to create OSMesa GL context.')
+
+    self._height = max_height
+    self._width = max_width
+
+    # Allocate a buffer to render into.
+    self._buffer = arrays.GLfloatArray.zeros((max_height, max_width, 4))
+
+  def make_current(self):
+    if self._context:
+      success = osmesa.OSMesaMakeCurrent(
+          self._context,
+          self._buffer,
+          GL.GL_FLOAT,
+          self._width,
+          self._height)
+      if not success:
+        raise RuntimeError('Failed to make OSMesa context current.')
+
+  def free(self):
+    """Frees resources associated with this context."""
+    if self._context and self._context == osmesa.OSMesaGetCurrentContext():
+      osmesa.OSMesaMakeCurrent(None, None, GL.GL_FLOAT, 0, 0)
+    osmesa.OSMesaDestroyContext(self._context)
+    self._buffer = None
+    self._context = None
+
+  def __del__(self):
+    self.free()

python/mujoco/raw.h

@@ -0,0 +1,60 @@
+// 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_PYTHON_RAW_H_
+#define MUJOCO_PYTHON_RAW_H_
+
+#include <mjdata.h>
+#include <mjmodel.h>
+#include <mjrender.h>
+#include <mjvisualize.h>
+
+// Type aliases for MuJoCo C structs to allow us refer to consistently refer
+// to them under the "raw" namespace.
+namespace mujoco::raw {
+
+using MjContact = ::mjContact;
+using MjData = ::mjData;
+using MjLROpt = ::mjLROpt;
+using MjModel = ::mjModel;
+using MjOption = ::mjOption;
+using MjSolverStat = ::mjSolverStat;
+using MjStatistic = ::mjStatistic;
+using MjTimerStat = ::mjTimerStat;
+using MjVisual = ::mjVisual;
+using MjVisualGlobal = decltype(::mjVisual::global);
+using MjVisualQuality = decltype(::mjVisual::quality);
+using MjVisualHeadlight = decltype(::mjVisual::headlight);
+using MjVisualMap = decltype(::mjVisual::map);
+using MjVisualScale = decltype(::mjVisual::scale);
+using MjVisualRgba = decltype(::mjVisual::rgba);
+using MjWarningStat = ::mjWarningStat;
+
+// From mjrender.h
+using MjrRect = ::mjrRect;
+using MjrContext = ::mjrContext;
+
+// From mjvisualize.h
+using MjvPerturb = ::mjvPerturb;
+using MjvCamera = ::mjvCamera;
+using MjvGLCamera = ::mjvGLCamera;
+using MjvGeom = ::mjvGeom;
+using MjvLight = ::mjvLight;
+using MjvOption = ::mjvOption;
+using MjvScene = ::mjvScene;
+using MjvFigure = ::mjvFigure;
+
+}  // namespace mujoco::raw
+
+#endif  // MUJOCO_PYTHON_RAW_H_

python/mujoco/render.cc

@@ -0,0 +1,276 @@
+// 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 <array>
+
+#include <Eigen/Core>
+#include <mjrender.h>
+#include <mujoco.h>
+#include "errors.h"
+#include "function_traits.h"
+#include "functions.h"
+#include "raw.h"
+#include "structs.h"
+#include <pybind11/eigen.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace mujoco::python {
+namespace _impl {
+template <>
+class MjWrapper<raw::MjrContext> : public WrapperBase<raw::MjrContext> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjModelWrapper& model, int fontscale);
+  MjWrapper(const MjWrapper&) = delete;
+  MjWrapper(MjWrapper&&) = default;
+  ~MjWrapper() = default;
+
+#define X(var) py_array_or_tuple_t<mjtNum> var
+  X(fogRGBA);
+  X(auxWidth);
+  X(auxHeight);
+  X(auxSamples);
+  X(auxFBO);
+  X(auxFBO_r);
+  X(auxColor);
+  X(auxColor_r);
+  X(textureType);
+  X(texture);
+  X(skinvertVBO);
+  X(skinnormalVBO);
+  X(skintexcoordVBO);
+  X(skinfaceVBO);
+  X(charWidth);
+  X(charWidthBig);
+#undef X
+};
+
+using MjrContextWrapper = MjWrapper<raw::MjrContext>;
+
+template <>
+struct enable_if_mj_struct<raw::MjrContext> {
+  using type = void;
+};
+
+static void MjrContextCapsuleDestructor(PyObject* pyobj) {
+  auto* ptr =
+      static_cast<raw::MjrContext*>(PyCapsule_GetPointer(pyobj, nullptr));
+  mjr_freeContext(ptr);
+  delete ptr;
+}
+
+#define X(var) var(InitPyArray(ptr_->var, owner_))
+#define X_SKIN(var) var(InitPyArray(std::array{ptr_->nskin}, ptr_->var, owner_))
+MjrContextWrapper::MjWrapper()
+    : WrapperBase([]() {
+        raw::MjrContext *const ctx = new raw::MjrContext;
+        mjr_defaultContext(ctx);
+        return ctx;
+      }()),
+      X(fogRGBA),
+      X(auxWidth),
+      X(auxHeight),
+      X(auxSamples),
+      X(auxFBO),
+      X(auxFBO_r),
+      X(auxColor),
+      X(auxColor_r),
+      X(textureType),
+      X(texture),
+      X_SKIN(skinvertVBO),
+      X_SKIN(skinnormalVBO),
+      X_SKIN(skintexcoordVBO),
+      X_SKIN(skinfaceVBO),
+      X(charWidth),
+      X(charWidthBig) {}
+
+MjrContextWrapper::MjWrapper(const MjModelWrapper& model, int fontscale)
+    : WrapperBase([fontscale](const raw::MjModel* m) {
+        raw::MjrContext *const ctx = new raw::MjrContext;
+        mjr_defaultContext(ctx);
+        InterceptMjErrors(mjr_makeContext)(m, ctx, fontscale);
+        return ctx;
+      }(model.get()), &MjrContextCapsuleDestructor),
+      X(fogRGBA),
+      X(auxWidth),
+      X(auxHeight),
+      X(auxSamples),
+      X(auxFBO),
+      X(auxFBO_r),
+      X(auxColor),
+      X(auxColor_r),
+      X(textureType),
+      X(texture),
+      X_SKIN(skinvertVBO),
+      X_SKIN(skinnormalVBO),
+      X_SKIN(skintexcoordVBO),
+      X_SKIN(skinfaceVBO),
+      X(charWidth),
+      X(charWidthBig) {}
+#undef X_SKIN
+#undef X
+}  // namespace _impl
+
+namespace {
+PYBIND11_MODULE(_render, pymodule) {
+  namespace py = ::pybind11;
+  namespace traits = python_traits;
+
+  using _impl::MjModelWrapper;
+  using _impl::MjrContextWrapper;
+
+  // Import the _structs module so that pybind11 knows about Python bindings
+  // for MjWrapper types and therefore generates prettier docstrings.
+  py::module::import("mujoco._structs");
+
+  py::class_<raw::MjrRect> mjrRect(pymodule, "MjrRect");
+  mjrRect.def(py::init([](int left, int bottom, int width, int height) {
+                return raw::MjrRect{left, bottom, width, height};
+              }),
+              py::arg("left"), py::arg("bottom"), py::arg("width"),
+              py::arg("height"));
+  mjrRect.def("__copy__",
+              [](const raw::MjrRect& other) { return raw::MjrRect(other); });
+  mjrRect.def("__deepcopy__", [](const raw::MjrRect& other, py::dict) {
+    return raw::MjrRect(other);
+  });
+#define X(var) mjrRect.def_readwrite(#var, &raw::MjrRect::var)
+  X(left);
+  X(bottom);
+  X(width);
+  X(height);
+#undef X
+
+  py::class_<MjrContextWrapper> mjrContext(pymodule, "MjrContext");
+  mjrContext.def(py::init<>());
+  mjrContext.def(py::init<const MjModelWrapper&, int>());
+#define X(var)                                                       \
+  mjrContext.def_property(                                           \
+      #var, [](const MjrContextWrapper& c) { return c.get()->var; }, \
+      [](MjrContextWrapper& c, mjtNum rhs) { c.get()->var = rhs; })
+  X(lineWidth);
+  X(shadowClip);
+  X(shadowScale);
+  X(fogStart);
+  X(fogEnd);
+  X(shadowSize);
+  X(offWidth);
+  X(offHeight);
+  X(offSamples);
+  X(fontScale);
+  X(offFBO);
+  X(offFBO_r);
+  X(offColor);
+  X(offColor_r);
+  X(offDepthStencil);
+  X(offDepthStencil_r);
+  X(shadowFBO);
+  X(shadowTex);
+  X(ntexture);
+  X(basePlane);
+  X(baseMesh);
+  X(baseHField);
+  X(baseBuiltin);
+  X(baseFontNormal);
+  X(baseFontShadow);
+  X(baseFontBig);
+  X(rangePlane);
+  X(rangeMesh);
+  X(rangeHField);
+  X(rangeBuiltin);
+  X(rangeFont);
+  X(nskin);
+  X(charHeight);
+  X(charHeightBig);
+  X(glewInitialized);
+  X(windowAvailable);
+  X(windowSamples);
+  X(windowStereo);
+  X(windowDoublebuffer);
+  X(currentBuffer);
+#undef X
+
+#define X(var)                      \
+  mjrContext.def_property_readonly( \
+      #var, [](const MjrContextWrapper& c) { return c.get()->var; })
+  X(nskin);
+#undef X
+
+#define X(var) DefinePyArray(mjrContext, #var, &MjrContextWrapper::var)
+  X(fogRGBA);
+  X(auxWidth);
+  X(auxHeight);
+  X(auxSamples);
+  X(auxFBO);
+  X(auxFBO_r);
+  X(auxColor);
+  X(auxColor_r);
+  X(textureType);
+  X(texture);
+  X(skinvertVBO);
+  X(skinnormalVBO);
+  X(skintexcoordVBO);
+  X(skinfaceVBO);
+  X(charWidth);
+  X(charWidthBig);
+#undef X
+
+  using EigenUnsignedCharVectorX = Eigen::Vector<unsigned char, Eigen::Dynamic>;
+  using EigenFloatVectorX = Eigen::Vector<float, Eigen::Dynamic>;
+
+  // Skipped: mjr_defaultContext (have MjrContext.__init__)
+  // Skipped: mjr_makeContext (have MjrContext.__init__)
+  Def<traits::mjr_changeFont>(pymodule);
+  Def<traits::mjr_addAux>(pymodule);
+  // Skipped: mjr_freeContext (have MjrContext.__del__)
+  Def<traits::mjr_uploadTexture>(pymodule);
+  Def<traits::mjr_uploadMesh>(pymodule);
+  Def<traits::mjr_uploadHField>(pymodule);
+  Def<traits::mjr_restoreBuffer>(pymodule);
+  Def<traits::mjr_setBuffer>(pymodule);
+  Def<traits::mjr_readPixels>(
+      pymodule, [](std::optional<py::array_t<uint8_t>> rgb,
+                   std::optional<py::array_t<float>> depth,
+                   const raw::MjrRect* viewport, const raw::MjrContext* con) {
+        return InterceptMjErrors(::mjr_readPixels)(
+            rgb.has_value() ? rgb->mutable_data() : nullptr,
+            depth.has_value() ? depth->mutable_data() : nullptr, *viewport,
+            con);
+      });
+  Def<traits::mjr_drawPixels>(
+      pymodule,
+      [](std::optional<Eigen::Ref<const EigenUnsignedCharVectorX>> rgb,
+         std::optional<Eigen::Ref<const EigenFloatVectorX>> depth,
+         const raw::MjrRect* viewport, const raw::MjrContext* con) {
+        return InterceptMjErrors(::mjr_drawPixels)(
+            rgb.has_value() ? rgb->data() : nullptr,
+            depth.has_value() ? depth->data() : nullptr, *viewport, con);
+      });
+  Def<traits::mjr_blitBuffer>(pymodule);
+  Def<traits::mjr_setAux>(pymodule);
+  Def<traits::mjr_blitAux>(pymodule);
+  Def<traits::mjr_text>(pymodule);
+  Def<traits::mjr_overlay>(pymodule);
+  Def<traits::mjr_maxViewport>(pymodule);
+  Def<traits::mjr_rectangle>(pymodule);
+  Def<traits::mjr_label>(pymodule);
+  Def<traits::mjr_figure>(pymodule);
+  Def<traits::mjr_render>(pymodule);
+  Def<traits::mjr_finish>(pymodule);
+  Def<traits::mjr_getError>(pymodule);
+  Def<traits::mjr_findRect>(pymodule);
+}  // PYBIND11_MODULE
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/render_test.py

@@ -0,0 +1,81 @@
+# 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.
+# ==============================================================================
+"""Tests for MuJoCo Python rendering."""
+
+from absl.testing import absltest
+import mujoco
+import numpy as np
+
+
+@absltest.skipUnless(hasattr(mujoco, 'mjr_render'),
+                     'MuJoCo rendering is disabled')
+class MuJoCoRenderTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.gl = mujoco.GLContext(640, 480)
+    self.gl.make_current()
+
+  def tearDown(self):
+    super().tearDown()
+    del self.gl
+
+  def test_can_render(self):
+    """Test that the bindings can successfully render a simple image.
+
+    This test sets up a basic MuJoCo rendering context similar to the example in
+    https://mujoco.readthedocs.io/en/latest/programming.html#visualization
+    It calls `mjr_rectangle` rather than `mjr_render` so that we can assert an
+    exact rendered image without needing golden data. The purpose of this test
+    is to ensure that the bindings can correctly return pixels in Python, rather
+    than to test MuJoCo's rendering pipeline itself.
+    """
+
+    self.model = mujoco.MjModel.from_xml_string('<mujoco><worldbody/></mujoco>')
+    self.data = mujoco.MjData(self.model)
+
+    scene = mujoco.MjvScene(self.model, maxgeom=0)
+    mujoco.mjv_updateScene(
+        self.model, self.data, mujoco.MjvOption(), mujoco.MjvPerturb(),
+        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL.value, scene)
+
+    context = mujoco.MjrContext(
+        self.model,
+        mujoco.mjtFontScale.mjFONTSCALE_150.value)
+    mujoco.mjr_setBuffer(
+        mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, context)
+
+    # MuJoCo's default render buffer size is 640x480.
+    full_rect = mujoco.MjrRect(0, 0, 640, 480)
+    mujoco.mjr_rectangle(full_rect, 0, 0, 0, 1)
+
+    blue_rect = mujoco.MjrRect(56, 67, 234, 123)
+    mujoco.mjr_rectangle(blue_rect, 0, 0, 1, 1)
+
+    expected_upside_down_image = np.zeros((480, 640, 3), dtype=np.uint8)
+    expected_upside_down_image[67:67+123, 56:56+234, 2] = 255
+
+    upside_down_image = np.empty((480, 640, 3), dtype=np.uint8)
+    mujoco.mjr_readPixels(upside_down_image, None, full_rect, context)
+    np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
+
+    # Check that mjr_readPixels can accept a flattened array.
+    upside_down_image[:] = 0
+    mujoco.mjr_readPixels(
+        np.reshape(upside_down_image, -1), None, full_rect, context)
+    np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
+
+if __name__ == '__main__':
+  absltest.main()

python/mujoco/rollout.cc

@@ -0,0 +1,263 @@
+// 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 <array>
+#include <cstdio>
+#include <iostream>
+#include <optional>
+#include <sstream>
+#include <string>
+
+#include "functions.h"
+#include "raw.h"
+#include <pybind11/buffer_info.h>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace mujoco::python {
+
+namespace {
+
+namespace py = ::pybind11;
+
+const auto rollout_doc = R"(
+Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
+
+  input arguments (required):
+    model              an instance of MjModel
+    data               an associated instance of MjData
+    nstate             an integer, number of initial states from which to roll out trajectories
+    nstep              an integer, number of steps to be taken for each trajectory
+  input arguments (optional):
+    initial_state      (nstate x nqva)                 nstate initial state vectors, nqva=nq+nv+na
+    initial_time       (nstate x 1)                    nstate initial times
+    initial_warmstart  (nstate x nv)                   nstate qacc_warmstart vectors
+    ctrl               (nstate x nstep x nu)           nstate length-nstep controls
+    qfrc_applied       (nstate x nstep x nv)           nstate length-nstep generalized forces
+    xfrc_applied       (nstate x nstep x nbody*6)      nstate length-nstep Cartesian wrenches
+    mocap              (nstate x nstep x nmocap*7)     nstate length-nstep mocap body poses
+  output arguments (optional):
+    state              (nstate x nstep x nqva)         nstate length-nstep states
+    sensordata         (nstate x nstep x nsendordata)  nstate length-nstep sensordatas
+)";
+
+// C-style rollout function, assumes all arguments are valid
+// all input fields of d are initialised, contents at call time do not matter
+// after returning, d will contain the last step of the last rollout
+void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
+                     const mjtNum* state0, const mjtNum* ctrl,
+                     const mjtNum* qfrc, const mjtNum* xfrc,
+                     const mjtNum* mocap, const mjtNum* time0,
+                     const mjtNum* warmstart0,
+                     mjtNum* state, mjtNum* sensordata) {
+  // model sizes
+  int nq = m->nq;
+  int nv = m->nv;
+  int na = m->na;
+  int nqva = nq + nv + na;
+  int nu = m->nu;
+  int nbody = m->nbody;
+  int nmocap = m->nmocap;
+  int nsensordata = m->nsensordata;
+
+  // loop over initial states
+  for (int s=0; s < nstate; s++) {
+
+    // set initial state
+    if (state0) {
+      mju_copy(d->qpos, state0 + s*nqva, nq);
+      mju_copy(d->qvel, state0 + s*nqva + nq, nv);
+      mju_copy(d->act,  state0 + s*nqva + nq + nv, na);
+    } else {
+      mju_copy(d->qpos, m->qpos0, nq);
+      mju_zero(d->qvel, nv);
+      mju_zero(d->act, na);
+    }
+
+    // set initial time
+    d->time = time0 ? time0[s] : 0;
+
+    // set warmstart accelerations
+    if (warmstart0) {
+      mju_copy(d->qacc_warmstart, warmstart0 + s*nv, nv);
+    } else {
+      mju_zero(d->qacc_warmstart, nv);
+    }
+
+    // clear control inputs if unspecified
+    if (s == 0) {
+      if (!ctrl) {
+        mju_zero(d->ctrl, nu);
+      }
+      if (!qfrc) {
+        mju_zero(d->qfrc_applied, nv);
+      }
+      if (!xfrc) {
+        mju_zero(d->xfrc_applied, 6*nbody);
+      }
+      if (!mocap) {
+        for (int j=0; j<nbody; j++) {
+          int id = m->body_mocapid[j];
+          if (id>=0) {
+            mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j);
+            mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j);
+          }
+        }
+      }
+    }
+
+    // roll out trajectories
+    for (int t = 0; t < nstep; t++) {
+      // controls
+      if (ctrl) {
+        mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu);
+      }
+      // generalized forces
+      if (qfrc) {
+        mju_copy(d->qfrc_applied, qfrc + s*nstep*nv + t*nv, nv);
+      }
+      // Cartesian wrenches
+      if (xfrc) {
+        mju_copy(d->xfrc_applied, xfrc + s*nstep*6*nbody + t*6*nbody, 6*nbody);
+      }
+      // mocap bodies
+      if (mocap) {
+        mju_copy(d->mocap_pos,
+                 mocap + s*nstep*7*nmocap + t*7*nmocap, 3*nmocap);
+        mju_copy(d->mocap_quat,
+                 mocap + s*nstep*7*nmocap + t*7*nmocap + 3*nmocap, 4*nmocap);
+      }
+
+      // step
+      mj_step(m, d);
+
+      // copy out new state
+      if (state) {
+        mju_copy(state + s*nstep*nqva + t*nqva,           d->qpos, nq);
+        mju_copy(state + s*nstep*nqva + t*nqva + nq,      d->qvel, nv);
+        mju_copy(state + s*nstep*nqva + t*nqva + nq + nv, d->act,  na);
+      }
+      // copy out sensor values
+      if (sensordata) {
+        mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata,
+                 d->sensordata, nsensordata);
+      }
+    }
+  }
+}
+
+
+// check size of optional argument to rollout(), return raw pointer
+mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
+                      const char* name, int nstate, int nstep, int dim) {
+  // if empty return nullptr
+  if (!arg.has_value()) {
+    return nullptr;
+  }
+
+  // get info
+  py::buffer_info info = arg->request();
+
+  // check size
+  int expected_size = nstate * nstep * dim;
+  if (info.size != expected_size) {
+    std::ostringstream msg;
+    msg << name << ".size should be " << expected_size <<  ", got " << info.size;
+    throw py::value_error(msg.str());
+  }
+  return static_cast<mjtNum*>(info.ptr);
+}
+
+
+PYBIND11_MODULE(_rollout, pymodule) {
+  namespace py = ::pybind11;
+  using PyCArray = py::array_t<mjtNum, py::array::c_style>;
+
+  // roll out open loop trajectories from multiple initial states
+  // get subsequent states and corresponding sensor values
+  pymodule.def(
+      "rollout",
+      [](const MjModelWrapper& m, MjDataWrapper& d, int nstate, int nstep,
+         std::optional<const PyCArray> init_state,
+         std::optional<const PyCArray> init_time,
+         std::optional<const PyCArray> init_warmstart,
+         std::optional<const PyCArray> ctrl,
+         std::optional<const PyCArray> qfrc,
+         std::optional<const PyCArray> xfrc,
+         std::optional<const PyCArray> mocap,
+         std::optional<const PyCArray> state,
+         std::optional<const PyCArray> sensordata
+         ) {
+
+        const raw::MjModel* model = m.get();
+        raw::MjData* data = d.get();
+
+        // check that some steps need to be taken, return if not
+        if (nstate < 1 || nstep < 1) {
+          return;
+        }
+
+        // get raw pointers
+        int nqva = model->nq + model->nv + model->na;
+        mjtNum* init_state_ptr =
+            get_array_ptr(init_state, "initial_state", nstate, 1, nqva);
+        mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu);
+        mjtNum* qfrc_ptr =
+            get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv);
+        mjtNum* xfrc_ptr =
+            get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody);
+        mjtNum* mocap_ptr =
+            get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap);
+        mjtNum* init_time_ptr =
+            get_array_ptr(init_time, "init_time", nstate, 1, 1);
+        mjtNum* init_warmstart_ptr =
+            get_array_ptr(init_warmstart, "init_warmstart", nstate, 1, model->nv);
+        mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva);
+        mjtNum* sensordata_ptr =
+            get_array_ptr(sensordata, "sensordata", nstate, nstep, model->nsensordata);
+
+        // perform rollouts
+        {
+          // release the GIL
+          py::gil_scoped_release no_gil;
+
+          // call unsafe rollout function
+          InterceptMjErrors(_unsafe_rollout)(
+              model, data, nstate, nstep, init_state_ptr, ctrl_ptr, qfrc_ptr,
+              xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr,
+              sensordata_ptr);
+        }
+      },
+      py::arg("model"),
+      py::arg("data"),
+      py::arg("nstate"),
+      py::arg("nstep"),
+      py::arg("initial_state")     = py::none(),
+      py::arg("initial_time")      = py::none(),
+      py::arg("initial_warmstart") = py::none(),
+      py::arg("ctrl")              = py::none(),
+      py::arg("qfrc_applied")      = py::none(),
+      py::arg("xfrc_applied")      = py::none(),
+      py::arg("mocap")             = py::none(),
+      py::arg("state")             = py::none(),
+      py::arg("sensordata")        = py::none(),
+      py::doc(rollout_doc)
+  );
+
+}  // namespace
+
+}
+
+}

python/mujoco/rollout.py

@@ -0,0 +1,205 @@
+# 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.
+# ==============================================================================
+"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
+
+from mujoco import _rollout
+import numpy as np
+
+
+def rollout(model, data, initial_state=None, ctrl=None,
+            *,  # require following arguments to be named
+            skip_checks=False,
+            nstate=None,
+            nstep=None,
+            initial_time=None,
+            initial_warmstart=None,
+            qfrc_applied=None,
+            xfrc_applied=None,
+            mocap=None,
+            state=None,
+            sensordata=None):
+  """Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
+
+    This function serves as a Python wrapper for the C++ functionality in
+    `rollout.cc`, please see documentation therein. This python funtion will
+    infer `nstate` and `nstep`, tile input arguments with singleton dimensions,
+    and allocate output arguments if none are given.
+  """
+  # don't infer nstate/nstep, don't support singleton expansion, don't allocate
+  # output arrays, just call rollout
+  if skip_checks:
+    _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
+                     initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
+                     state, sensordata)
+    return state, sensordata
+
+  # check types
+  if nstate and not isinstance(nstate, int):
+    raise ValueError('nstate must be an integer')
+  if nstep and not isinstance(nstep, int):
+    raise ValueError('nstep must be an integer')
+  _check_must_be_numeric(
+      initial_state=initial_state,
+      initial_time=initial_time,
+      initial_warmstart=initial_warmstart,
+      ctrl=ctrl,
+      qfrc_applied=qfrc_applied,
+      xfrc_applied=xfrc_applied,
+      mocap=mocap,
+      state=state,
+      sensordata=sensordata)
+
+  # check number of dimensions
+  _check_number_of_dimensions(2,
+                              initial_state=initial_state,
+                              initial_time=initial_time,
+                              initial_warmstart=initial_warmstart)
+  _check_number_of_dimensions(3,
+                              ctrl=ctrl,
+                              qfrc_applied=qfrc_applied,
+                              xfrc_applied=xfrc_applied,
+                              mocap=mocap,
+                              state=state,
+                              sensordata=sensordata)
+
+  # ensure 2D, make contiguous, row-major (C ordering)
+  initial_state = _ensure_2d(initial_state)
+  initial_time = _ensure_2d(initial_time)
+  initial_warmstart = _ensure_2d(initial_warmstart)
+
+  # ensure 3D, make contiguous, row-major (C ordering)
+  ctrl = _ensure_3d(ctrl)
+  qfrc_applied = _ensure_3d(qfrc_applied)
+  xfrc_applied = _ensure_3d(xfrc_applied)
+  mocap = _ensure_3d(mocap)
+  state = _ensure_3d(state)
+  sensordata = _ensure_3d(sensordata)
+
+  # check trailing dimensions
+  _check_trailing_dimension(model.nq + model.nv + model.na,
+                            initial_state=initial_state, state=state)
+  _check_trailing_dimension(1, initial_time=initial_time)
+  _check_trailing_dimension(model.nu, ctrl=ctrl)
+  _check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied)
+  _check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied)
+  _check_trailing_dimension(model.nmocap*7, mocap=mocap)
+  _check_trailing_dimension(model.nsensordata, sensordata=sensordata)
+
+  # infer nstate, check for incompatibilities
+  nstate = _infer_dimension(0, nstate or 1,
+                            initial_state=initial_state,
+                            initial_time=initial_time,
+                            initial_warmstart=initial_warmstart,
+                            ctrl=ctrl,
+                            qfrc_applied=qfrc_applied,
+                            xfrc_applied=xfrc_applied,
+                            mocap=mocap,
+                            state=state,
+                            sensordata=sensordata)
+
+  # infer nstep, check for incompatibilities
+  nstep = _infer_dimension(1, nstep or 1,
+                           ctrl=ctrl,
+                           qfrc_applied=qfrc_applied,
+                           xfrc_applied=xfrc_applied,
+                           mocap=mocap,
+                           state=state,
+                           sensordata=sensordata)
+
+  # tile input arrays if required (singleton expansion)
+  initial_state = _tile_if_required(initial_state, nstate)
+  initial_time = _tile_if_required(initial_time, nstate)
+  initial_warmstart = _tile_if_required(initial_warmstart, nstate)
+  ctrl = _tile_if_required(ctrl, nstate, nstep)
+  qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep)
+  xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep)
+  mocap = _tile_if_required(mocap, nstate, nstep)
+
+  # allocate output if not provided
+  if state is None:
+    state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+  if sensordata is None:
+    sensordata = np.empty((nstate, nstep, model.nsensordata))
+
+  # call rollout
+  _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
+                   initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
+                   state, sensordata)
+
+  # return squeezed outputs
+  return state.squeeze(), sensordata.squeeze()
+
+def _check_must_be_numeric(**kwargs):
+  for key, value in kwargs.items():
+    if value is None:
+      continue
+    if not isinstance(value, np.ndarray) and not isinstance(value, float):
+      raise ValueError(f'{key} must be a numpy array or float')
+
+def _check_number_of_dimensions(ndim, **kwargs):
+  for key, value in kwargs.items():
+    if value is None:
+      continue
+    if value.ndim > ndim:
+      raise ValueError(f'{key} can have at most {ndim} dimensions')
+
+def _check_trailing_dimension(dim, **kwargs):
+  for key, value in kwargs.items():
+    if value is None:
+      continue
+    if value.shape[-1] != dim:
+      raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}')
+
+def _ensure_2d(arg):
+  if arg is None:
+    return None
+  else:
+    return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
+
+def _ensure_3d(arg):
+  if arg is None:
+    return None
+  else:
+    # np.atleast_3d adds both leading and trailing dims, we want only leading
+    if arg.ndim == 0:
+      arg = arg[np.newaxis, np.newaxis, np.newaxis, ...]
+    elif arg.ndim == 1:
+      arg = arg[np.newaxis, np.newaxis, ...]
+    elif arg.ndim == 2:
+      arg = arg[np.newaxis, ...]
+    return np.ascontiguousarray(arg, dtype=np.float64)
+
+def _infer_dimension(dim, value, **kwargs):
+  for name, array in kwargs.items():
+    if array is None:
+      continue
+    if array.shape[dim] != value:
+      if value == 1:
+        value = array.shape[dim]
+      elif array.shape[dim] != 1:
+        raise ValueError(
+            f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}'
+        )
+  return value
+
+def _tile_if_required(array, dim0, dim1=None):
+  if array is None:
+    return
+  reps = np.ones(array.ndim, dtype=int)
+  if array.shape[0] == 1:
+    reps[0] = dim0
+  if dim1 is not None and array.shape[1] == 1:
+    reps[1] = dim1
+  return np.tile(array, reps)

python/mujoco/rollout_test.py

@@ -0,0 +1,488 @@
+# 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.
+# ==============================================================================
+"""tests for rollout function."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import mujoco
+import numpy as np
+import concurrent.futures
+import threading
+from mujoco import rollout
+
+#--------------------------- models used for testing ---------------------------
+
+TEST_XML = r"""
+<mujoco>
+  <worldbody>
+    <light pos="0 0 2"/>
+    <geom type="plane" size="5 5 .1"/>
+    <body pos="0 0 .1">
+      <joint name="yaw" axis="0 0 1"/>
+      <joint name="pitch" axis="0 1 0"/>
+      <geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
+      <geom type="box" pos="1 0 0" size=".1 .1 .1"/>
+      <site name="site" pos="1 0 0"/>
+    </body>
+  </worldbody>
+  <actuator>
+    <general joint="pitch" gainprm="100"/>
+    <general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
+  </actuator>
+  <sensor>
+    <accelerometer site="site"/>
+  </sensor>
+</mujoco>
+"""
+
+TEST_XML_NO_SENSORS = r"""
+<mujoco>
+  <worldbody>
+    <light pos="0 0 2"/>
+    <geom type="plane" size="5 5 .1"/>
+    <body pos="0 0 .1">
+      <joint name="yaw" axis="0 0 1"/>
+      <joint name="pitch" axis="0 1 0"/>
+      <geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
+      <geom type="box" pos="1 0 0" size=".1 .1 .1"/>
+      <site name="site" pos="1 0 0"/>
+    </body>
+  </worldbody>
+  <actuator>
+    <general joint="pitch" gainprm="100"/>
+    <general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
+  </actuator>
+</mujoco>
+"""
+
+TEST_XML_NO_ACTUATORS = r"""
+<mujoco>
+  <worldbody>
+    <light pos="0 0 2"/>
+    <geom type="plane" size="5 5 .1"/>
+    <body pos="0 0 .1">
+      <joint name="yaw" axis="0 0 1"/>
+      <joint name="pitch" axis="0 1 0"/>
+      <geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
+      <geom type="box" pos="1 0 0" size=".1 .1 .1"/>
+      <site name="site" pos="1 0 0"/>
+    </body>
+  </worldbody>
+  <sensor>
+    <accelerometer site="site"/>
+  </sensor>
+</mujoco>
+"""
+
+TEST_XML_MOCAP = r"""
+<mujoco>
+  <worldbody>
+    <body name="1" mocap="true">
+    </body>
+    <body name="2" mocap="true">
+    </body>
+  </worldbody>
+  <sensor>
+    <framepos objtype="xbody" objname="1"/>
+    <framequat objtype="xbody" objname="2"/>
+  </sensor>
+</mujoco>
+"""
+
+TEST_XML_EMPTY = r"""
+<mujoco>
+</mujoco>
+"""
+
+ALL_MODELS = {'TEST_XML': TEST_XML,
+              'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
+              'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
+              'TEST_XML_EMPTY': TEST_XML_EMPTY}
+
+#------------------------------- tests -----------------------------------------
+
+class MuJoCoRolloutTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(42)
+
+  #----------------------------- test basic operation
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_single_step(self, model_name):
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    ctrl = np.random.randn(model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    mujoco.mj_resetData(model, data)
+    py_state, py_sensordata = step(model, data, initial_state, ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_single_rollout(self, model_name):
+    nstep = 3  # number of timesteps
+
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    ctrl = np.random.randn(nstep, model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    py_state, py_sensordata = single_rollout(model, data, initial_state,
+                                             ctrl=ctrl)
+    np.testing.assert_array_equal(state, np.asarray(py_state))
+    np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata))
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_multi_step(self, model_name):
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    nstate = 5  # number of initial states
+
+    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    ctrl = np.random.randn(nstate, 1, model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    mujoco.mj_resetData(model, data)
+    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+                                            ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_single_rollout_fixed_ctrl(self, model_name):
+    nstep = 3
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    ctrl = np.random.randn(model.nu)
+    state = np.empty((nstep, model.nq + model.nv + model.na))
+    sensordata = np.empty((nstep, model.nsensordata))
+    rollout.rollout(model, data, initial_state, ctrl,
+                    state=state, sensordata=sensordata)
+
+    ctrl = np.tile(ctrl, (nstep, 1))  # repeat??
+    py_state, py_sensordata = single_rollout(model, data, initial_state,
+                                             ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_multi_rollout(self, model_name):
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    nstate = 2  # number of initial states
+    nstep = 3  # number of timesteps
+
+    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    ctrl = np.random.randn(nstate, nstep, model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+                                            ctrl=ctrl)
+    np.testing.assert_array_equal(py_state, py_state)
+    np.testing.assert_array_equal(py_sensordata, py_sensordata)
+
+  @parameterized.parameters(ALL_MODELS.keys())
+  def test_multi_rollout_fixed_ctrl_infer_from_output(self, model_name):
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    nstate = 2  # number of initial states
+    nstep = 3  # number of timesteps
+
+    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    ctrl = np.random.randn(nstate, 1, model.nu)  # 1 control in the time dimension
+    state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+                                        state=state)
+
+    ctrl = np.repeat(ctrl, nstep, axis=1)
+    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+                                            ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+  @parameterized.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]],
+                         model_name=list(ALL_MODELS.keys()))
+  def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name):
+    model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    data = mujoco.MjData(model)
+
+    nstate = 4  # number of initial states
+
+    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+
+    # arg_nstep is the horizon for {ctrl, qfrc_applied, xfrc_applied}, respectively
+    ctrl = np.random.randn(nstate, arg_nstep[0], model.nu)
+    qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv)
+    xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6)
+
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+                                        qfrc_applied=qfrc_applied,
+                                        xfrc_applied=xfrc_applied)
+
+    # tile singleton arguments
+    nstep = max(arg_nstep)
+    if arg_nstep[0] == 1:
+      ctrl = np.repeat(ctrl, nstep, axis=1)
+    if arg_nstep[1] == 1:
+      qfrc_applied = np.repeat(qfrc_applied, nstep, axis=1)
+    if arg_nstep[2] == 1:
+      xfrc_applied = np.repeat(xfrc_applied, nstep, axis=1)
+
+    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+                                            ctrl=ctrl,
+                                            qfrc_applied=qfrc_applied,
+                                            xfrc_applied=xfrc_applied)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+  #----------------------------- test threaded operation
+
+  def test_threading(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    num_workers = 32
+    nstate = 10000
+    nstep = 5
+    initial_state = np.random.randn(nstate, model.nq+model.nv+model.na)
+    state = np.zeros((nstate, nstep, model.nq+model.nv+model.na))
+    sensordata = np.zeros((nstate, nstep, model.nsensordata))
+    ctrl = np.random.randn(nstate, nstep, model.nu)
+
+    thread_local = threading.local()
+
+    def thread_initializer():
+      thread_local.data = mujoco.MjData(model)
+
+    def call_rollout(initial_state, ctrl, state):
+      rollout.rollout(model, thread_local.data, skip_checks=True,
+                      nstate=initial_state.shape[0], nstep=nstep,
+                      initial_state=initial_state, ctrl=ctrl, state=state)
+
+    n = initial_state.shape[0] // num_workers  # integer division
+    chunks = []  # a list of tuples, one per worker
+    for i in range(num_workers-1):
+      chunks.append(
+          (initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n]))
+    # last chunk, absorbing the remainder:
+    chunks.append(
+        (initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:],
+         state[(num_workers-1)*n:]))
+
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=num_workers, initializer=thread_initializer) as executor:
+      futures = []
+      for chunk in chunks:
+        futures.append(executor.submit(call_rollout, *chunk))
+      for future in concurrent.futures.as_completed(futures):
+        future.result()
+
+    data = mujoco.MjData(model)
+    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+                                            ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
+
+  #----------------------------- test advanced operation
+
+  def test_time(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    data = mujoco.MjData(model)
+
+    nstate = 1
+    nstep = 3
+
+    initial_time = np.array([[2.]])
+    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    ctrl = np.random.randn(nstate, nstep, model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+                                        initial_time=initial_time)
+
+    self.assertAlmostEqual(data.time, 2 + nstep*model.opt.timestep)
+
+  def test_warmstart(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    data = mujoco.MjData(model)
+
+    state0 = np.zeros(model.nq + model.nv + model.na)
+    ctrl = np.zeros(model.nu)
+    state1, _ = step(model, data, state0, ctrl=ctrl)
+    initial_warmstart = data.qacc_warmstart.copy()
+
+    state2, _ = step(model, data, state1, ctrl=ctrl)
+
+    state, _ = rollout.rollout(model, data, state1, ctrl)
+    assert np.linalg.norm(state-state2) > 0
+
+    state, _ = rollout.rollout(model, data, state1, ctrl,
+                               initial_warmstart=initial_warmstart)
+    np.testing.assert_array_equal(state, state2)
+
+  def test_mocap(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_MOCAP)
+    data = mujoco.MjData(model)
+
+    initial_state = np.zeros(model.nq + model.nv + model.na)
+    pos1 = np.array((1., 2., 3.))
+    quat1 = np.array((1., 2., 3., 4.))
+    quat1 /= np.linalg.norm(quat1)
+    pos2 = np.array((2., 3., 4.))
+    quat2 = np.array((2., 3., 4., 5.))
+    quat2 /= np.linalg.norm(quat2)
+    mocap = np.hstack((pos1, quat1, pos2, quat2))
+
+    state, sensordata = rollout.rollout(model, data, initial_state, mocap=mocap)
+
+    np.testing.assert_array_almost_equal(sensordata[:3], pos1)
+    np.testing.assert_array_almost_equal(sensordata[3:], quat2)
+
+  #----------------------------- test correctness
+
+  def test_intercept_mj_errors(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    data = mujoco.MjData(model)
+
+    initial_state = np.zeros(model.nq + model.nv + model.na)
+    ctrl = np.zeros((3, model.nu))
+
+    model.opt.solver = 10  # invalid solver type
+    with self.assertRaisesWithLiteralMatch(mujoco.FatalError,
+                                           'Unknown solver type 10'):
+      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+  def test_invalid(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    data = mujoco.MjData(model)
+
+    initial_state = np.zeros(model.nq + model.nv + model.na)
+
+    ctrl = 'string'
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'ctrl must be a numpy array or float'):
+      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    qfrc_applied = np.zeros((2, 3, 4, 5))
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'qfrc_applied can have at most 3 dimensions'):
+      state, sensordata = rollout.rollout(model, data, initial_state,
+                                          qfrc_applied=qfrc_applied)
+
+  def test_bad_sizes(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    data = mujoco.MjData(model)
+
+    initial_state = np.random.randn(model.nq + model.nv + model.na+1)
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'trailing dimension of initial_state must be 5, got 6'):
+      state, sensordata = rollout.rollout(model, data, initial_state)
+
+    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    ctrl = np.random.randn(model.nu+1)
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'trailing dimension of ctrl must be 2, got 3'):
+      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    ctrl = np.random.randn(2, model.nu)
+    qfrc_applied = np.random.randn(3, model.nv)  # incompatible horizon
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'):
+      state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+                                          qfrc_applied=qfrc_applied)
+
+  def test_stateless(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART.value
+    data = mujoco.MjData(model)
+
+    # call step with a clean mjData
+    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    ctrl = np.random.randn(model.nu)
+    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+
+    # fill mjData with some debug value, see that we still get the same outputs
+    mujoco.mj_resetDataDebug(model, data, 255)
+    debug_state, debug_sensordata = rollout.rollout(model, data, initial_state,
+                                                    ctrl)
+
+    np.testing.assert_array_equal(state, debug_state)
+    np.testing.assert_array_equal(sensordata, debug_sensordata)
+
+
+#--------------- Python implementation of rollout functionality ----------------
+
+def get_state(data):
+  return np.hstack((data.qpos, data.qvel, data.act))
+
+def set_state(model, data, state):
+  data.qpos = state[:model.nq]
+  data.qvel = state[model.nq:model.nq+model.nv]
+  data.act = state[model.nq+model.nv:model.nq+model.nv+model.na]
+
+def step(model, data, state, **kwargs):
+  if state is not None:
+    set_state(model, data, state)
+  for key, value in kwargs.items():
+    if value is not None:
+      setattr(data, key, np.reshape(value, getattr(data, key).shape))
+  mujoco.mj_step(model, data)
+  return (get_state(data), data.sensordata)
+
+def single_rollout(model, data, initial_state, **kwargs):
+  arg_nstep = set([a.shape[0] for a in kwargs.values()])
+  assert len(arg_nstep) == 1  # nstep dimensions must match
+  nstep = arg_nstep.pop()
+
+  state = np.empty((nstep, model.nq + model.nv + model.na))
+  sensordata = np.empty((nstep, model.nsensordata))
+
+  mujoco.mj_resetData(model, data)
+  for t in range(nstep):
+    kwargs_t = {}
+    for key, value in kwargs.items():
+      kwargs_t[key] = value[0 if value.ndim == 1 else t]
+    state[t], sensordata[t] = step(model, data,
+                                   initial_state if t==0 else None,
+                                   **kwargs_t)
+  return state, sensordata
+
+def multi_rollout(model, data, initial_state, **kwargs):
+  nstate = initial_state.shape[0]
+  arg_nstep = set([a.shape[1] for a in kwargs.values()])
+  assert len(arg_nstep) == 1  # nstep dimensions must match
+  nstep = arg_nstep.pop()
+
+  state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+  sensordata = np.empty((nstate, nstep, model.nsensordata))
+  for s in range(nstate):
+    kwargs_s = {key : value[s] for key, value in kwargs.items()}
+    state_s, sensordata_s = single_rollout(model, data, initial_state[s],
+                                           **kwargs_s)
+    state[s] = state_s
+    sensordata[s] = sensordata_s
+  return state.squeeze(), sensordata.squeeze()
+
+if __name__ == '__main__':
+  absltest.main()

python/mujoco/serialization.h

@@ -0,0 +1,72 @@
+// 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_PYTHON_SERIALIZATION_H_
+#define MUJOCO_PYTHON_SERIALIZATION_H_
+
+#include <iostream>
+
+#include <mjtnum.h>
+
+namespace mujoco::python::_impl {
+
+// For now, the serialization format is architecture- and version-dependent:
+// It assumes that the writer and the reader have the same endianess, same
+// numeric value sizes and same struct definitions.
+// It is not safe to load serialized data from another version of MuJoCo or from
+// a machine with a different architecture.
+
+static_assert(sizeof(char) == 1);
+static_assert(sizeof(int) == 4);
+static_assert(sizeof(mjtNum) == 8);
+
+inline void WriteChar(std::ostream& output, char c) {
+  output.write(&c, 1);
+}
+
+inline char ReadChar(std::istream& input) {
+  char c = '\0';
+  input.read(&c, 1);
+  return c;
+}
+
+inline void WriteInt(std::ostream& output, int i) {
+  output.write(reinterpret_cast<char*>(&i), sizeof(int));
+}
+
+inline int ReadInt(std::istream& input) {
+  int i = 0;
+  input.read(reinterpret_cast<char*>(&i), sizeof(int));
+  return i;
+}
+
+inline void WriteBytes(std::ostream& output, const void* src, size_t nbytes) {
+  // Start by writing nbytes itself, so it can be validated at the time of
+  // reading.
+  WriteInt(output, nbytes);
+  output.write(reinterpret_cast<const char*>(src), nbytes);
+}
+
+inline void ReadBytes(std::istream& input, void* dest, size_t nbytes) {
+  size_t actual_nbytes = ReadInt(input);
+  if (actual_nbytes != nbytes) {
+    input.setstate(input.rdstate() | std::ios_base::failbit);
+    return;
+  }
+  input.read(reinterpret_cast<char*>(dest), nbytes);
+}
+
+}  // namespace mujoco::python::_impl
+
+#endif  // MUJOCO_PYTHON_SERIALIZATION_H_

python/mujoco/structs.h

@@ -0,0 +1,978 @@
+// 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_PYTHON_STRUCTS_H_
+#define MUJOCO_PYTHON_STRUCTS_H_
+
+#include <array>
+#include <istream>
+#include <memory>
+#include <optional>
+#include <ostream>
+#include <unordered_map>
+#include <string>
+#include <vector>
+
+#include <absl/types/span.h>
+#include <mujoco.h>
+#include <mjxmacro.h>
+#include "indexers.h"
+#include "mjdata_meta.h"
+#include "raw.h"
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+namespace _impl {
+template <typename T>
+class WrapperBase {
+ public:
+  WrapperBase(WrapperBase&& other) = default;
+
+  T* get() { return ptr_; }
+  const T* get() const { return ptr_; }
+
+  const pybind11::handle owner() const { return owner_; }
+
+ protected:
+  static void DefaultCapsuleDestructor(PyObject* pyobj) {
+    T* ptr = static_cast<T*>(PyCapsule_GetPointer(pyobj, nullptr));
+    if (ptr) {
+      delete ptr;
+    }
+  }
+
+  // Takes ownership of ptr.
+  explicit WrapperBase(T* ptr,
+                       void (*destructor)(PyObject*) = DefaultCapsuleDestructor)
+      : ptr_(ptr),
+        owner_(pybind11::capsule(ptr_, /* name = */ nullptr, destructor)) {}
+
+  // `ptr` is owned by `owner`.
+  WrapperBase(T* ptr, pybind11::handle owner)
+      : ptr_(ptr),
+        owner_(pybind11::reinterpret_borrow<pybind11::object>(owner)) {}
+
+  T* ptr_;
+  pybind11::object owner_;
+};
+
+template <typename T, typename = void>
+struct py_array_or_tuple {
+  using type = pybind11::tuple;
+};
+
+template <typename T>
+struct py_array_or_tuple<T, std::enable_if_t<std::is_arithmetic_v<T>>> {
+  using type = pybind11::array_t<T>;
+};
+
+// A type that resolves to a NumPy array if the dtype is numeric, and
+// a Python tuple otherwise.
+template <typename T> using py_array_or_tuple_t =
+    typename py_array_or_tuple<T>::type;
+
+template <typename T>
+struct enable_if_mj_struct {};
+
+template <typename T>
+class MjWrapper {};
+
+template <typename T>
+class StructListBase {
+ public:
+  StructListBase(T* ptr, int num, pybind11::handle owner) : ptr_(ptr) {
+    for (int i = 0; i < num; ++i) {
+      wrappers_.push_back(std::make_shared<MjWrapper<T>>(&ptr[i], owner));
+    }
+  }
+
+  StructListBase(const StructListBase& other) = delete;
+  StructListBase(StructListBase&& other) = default;
+
+  MjWrapper<T>& operator[](int i) {
+    if (i < 0 || i >= wrappers_.size()) {
+      throw pybind11::index_error();
+    }
+    return *wrappers_[i];
+  }
+
+  int size() const {
+    return wrappers_.size();
+  }
+
+ protected:
+  // Slicing
+  StructListBase(StructListBase& other, pybind11::slice slice) {
+    pybind11::size_t start, stop, step, slicelength;
+    slice.compute(other.size(), &start, &stop, &step, &slicelength);
+    ptr_ = &other.ptr_[start];
+    for (int i = start; i < stop; i += step) {
+      wrappers_.push_back(other.wrappers_[i]);
+    }
+  }
+
+  T* ptr_;
+
+  // Using shared_ptr here so that we get identical Python objects when slicing.
+  std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
+};
+
+template <typename T>
+struct is_mj_struct_list { static constexpr bool value = false; };
+
+template <typename T>
+class MjStructList {};
+
+// ==================== MJOPTION ===============================================
+template <>
+class MjWrapper<raw::MjOption> : public WrapperBase<raw::MjOption> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjOption* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var, dim) py_array_or_tuple_t<mjtNum> var;
+  MJOPTION_VECTORS
+  #undef X
+};
+
+using MjOptionWrapper = MjWrapper<raw::MjOption>;
+
+template <>
+struct enable_if_mj_struct<raw::MjOption> { using type = void; };
+
+// ==================== MJVISUAL ===============================================
+template <>
+class MjWrapper<raw::MjVisualHeadlight>
+    : public WrapperBase<raw::MjVisualHeadlight> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper& other);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjVisualHeadlight* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<float> var
+  X(ambient);
+  X(diffuse);
+  X(specular);
+  #undef X
+};
+
+using MjVisualHeadlightWrapper = MjWrapper<raw::MjVisualHeadlight>;
+
+template <>
+struct enable_if_mj_struct<raw::MjVisualHeadlight> { using type = void; };
+
+template <>
+class MjWrapper<raw::MjVisualRgba> : public WrapperBase<raw::MjVisualRgba> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper& other);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjVisualRgba* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<float> var
+  X(fog);
+  X(haze);
+  X(force);
+  X(inertia);
+  X(joint);
+  X(actuator);
+  X(actuatornegative);
+  X(actuatorpositive);
+  X(com);
+  X(camera);
+  X(light);
+  X(selectpoint);
+  X(connect);
+  X(contactpoint);
+  X(contactforce);
+  X(contactfriction);
+  X(contacttorque);
+  X(contactgap);
+  X(rangefinder);
+  X(constraint);
+  X(slidercrank);
+  X(crankbroken);
+  #undef X
+};
+
+using MjVisualRgbaWrapper = MjWrapper<raw::MjVisualRgba>;
+
+template <>
+struct enable_if_mj_struct<raw::MjVisualRgba> { using type = void; };
+
+template <>
+class MjWrapper<raw::MjVisual> : public WrapperBase<raw::MjVisual> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper& other);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjVisual* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  MjVisualHeadlightWrapper headlight;
+  MjVisualRgbaWrapper rgba;
+};
+
+using MjVisualWrapper = MjWrapper<raw::MjVisual>;
+
+template <>
+struct enable_if_mj_struct<raw::MjVisual> { using type = void; };
+
+// ==================== MJSTATISTIC ============================================
+template <>
+class MjWrapper<raw::MjStatistic> : public WrapperBase<raw::MjStatistic> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjStatistic* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<mjtNum> var
+  X(center);
+  #undef X
+};
+
+using MjStatisticWrapper = MjWrapper<raw::MjStatistic>;
+
+template <>
+struct enable_if_mj_struct<raw::MjStatistic> { using type = void; };
+
+// ==================== MJWARNINGSTAT ==========================================
+template <>
+class MjWrapper<raw::MjWarningStat> : public WrapperBase<raw::MjWarningStat> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjWarningStat* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+};
+
+using MjWarningStatWrapper = MjWrapper<raw::MjWarningStat>;
+
+template <>
+struct enable_if_mj_struct<raw::MjWarningStat> { using type = void; };
+
+template <>
+class MjStructList<raw::MjWarningStat>
+    : public StructListBase<raw::MjWarningStat> {
+ public:
+  MjStructList(raw::MjWarningStat* ptr, int num, pybind11::handle owner);
+
+  using StructListBase::operator[];
+  using StructListBase::size;
+  MjStructList Slice(pybind11::slice slice) {
+    return MjStructList(*this, slice);
+  }
+
+#define X(type, var) pybind11::array_t<type> var
+  X(int, lastinfo);
+  X(int, number);
+#undef X
+
+ protected:
+  MjStructList(MjStructList& other, pybind11::slice slice);
+};
+
+using MjWarningStatList = MjStructList<raw::MjWarningStat>;
+
+template <>
+struct py_array_or_tuple<raw::MjWarningStat> {
+  using type = MjWarningStatList;
+};
+
+template <>
+struct is_mj_struct_list<raw::MjWarningStat> {
+  static constexpr bool value = true;
+};
+
+// ==================== MJTIMERSTAT ============================================
+template <>
+class MjWrapper<raw::MjTimerStat> : public WrapperBase<raw::MjTimerStat> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjTimerStat* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+};
+
+using MjTimerStatWrapper = MjWrapper<raw::MjTimerStat>;
+
+template <>
+struct enable_if_mj_struct<raw::MjTimerStat> { using type = void; };
+
+template <>
+class MjStructList<raw::MjTimerStat> : public StructListBase<raw::MjTimerStat> {
+ public:
+  MjStructList(raw::MjTimerStat* ptr, int num, pybind11::handle owner);
+
+  using StructListBase::operator[];
+  using StructListBase::size;
+  MjStructList Slice(pybind11::slice slice) {
+    return MjStructList(*this, slice);
+  }
+
+#define X(type, var) pybind11::array_t<type> var
+  X(mjtNum, duration);
+  X(int, number);
+#undef X
+
+ protected:
+  MjStructList(MjStructList& other, pybind11::slice slice);
+};
+
+using MjTimerStatList = MjStructList<raw::MjTimerStat>;
+
+template <>
+struct py_array_or_tuple<raw::MjTimerStat> {
+  using type = MjTimerStatList;
+};
+
+template <>
+struct is_mj_struct_list<raw::MjTimerStat> {
+  static constexpr bool value = true;
+};
+
+// ==================== MJSOLVERSTAT ===========================================
+template <>
+class MjWrapper<raw::MjSolverStat> : public WrapperBase<raw::MjSolverStat> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjSolverStat* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+};
+
+using MjSolverStatWrapper = MjWrapper<raw::MjSolverStat>;
+
+template <>
+struct enable_if_mj_struct<raw::MjSolverStat> { using type = void; };
+
+template <>
+class MjStructList<raw::MjSolverStat>
+    : public StructListBase<raw::MjSolverStat> {
+ public:
+  MjStructList(raw::MjSolverStat* ptr, int num, pybind11::handle owner);
+
+  using StructListBase::operator[];
+  using StructListBase::size;
+  MjStructList Slice(pybind11::slice slice) {
+    return MjStructList(*this, slice);
+  }
+
+#define X(type, var) pybind11::array_t<type> var
+  X(mjtNum, improvement);
+  X(mjtNum, gradient);
+  X(mjtNum, lineslope);
+  X(int, nactive);
+  X(int, nchange);
+  X(int, neval);
+  X(int, nupdate);
+#undef X
+
+ protected:
+  MjStructList(MjStructList& other, pybind11::slice slice);
+};
+
+using MjSolverStatList = MjStructList<raw::MjSolverStat>;
+
+template <>
+struct py_array_or_tuple<raw::MjSolverStat> {
+  using type = MjSolverStatList;
+};
+
+template <>
+struct is_mj_struct_list<raw::MjSolverStat> {
+  static constexpr bool value = true;
+};
+
+// ==================== MJMODEL ================================================
+template <>
+class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
+ public:
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&);
+  ~MjWrapper();
+
+  MjModelIndexer& indexer() { return indexer_; }
+
+  void Serialize(std::ostream& output) const;
+  static MjWrapper<raw::MjModel> Deserialize(std::istream& input);
+
+  static MjWrapper LoadXMLFile(
+      const std::string& filename,
+      const std::optional<
+          std::unordered_map<std::string, pybind11::bytes>>& assets);
+
+  static MjWrapper LoadBinaryFile(
+      const std::string& filename,
+      const std::optional<
+          std::unordered_map<std::string, pybind11::bytes>>& assets);
+
+  static MjWrapper LoadXML(
+      const std::string& xml,
+      const std::optional<
+          std::unordered_map<std::string, pybind11::bytes>>& assets);
+
+  static constexpr char kFromRawPointer[] =
+      "__MUJOCO_STRUCTS_MJMODELWRAPPER_LOOKUP";
+  static MjWrapper* FromRawPointer(raw::MjModel* m) noexcept;
+
+  MjOptionWrapper opt;
+  MjVisualWrapper vis;
+  MjStatisticWrapper stat;
+
+#define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
+  MJMODEL_POINTERS
+#undef X
+  // TODO(nimrod): Exclude text_data and names from the MJMODEL_POINTERS macro.
+  pybind11::bytes text_data_bytes;
+  pybind11::bytes names_bytes;
+
+ private:
+  explicit MjWrapper(raw::MjModel* ptr);
+
+  MjModelIndexer indexer_;
+};
+
+using MjModelWrapper = MjWrapper<raw::MjModel>;
+
+template <>
+struct enable_if_mj_struct<raw::MjModel> { using type = void; };
+
+// ==================== MJCONTACT ==============================================
+template <>
+class MjWrapper<raw::MjContact> : public WrapperBase<raw::MjContact> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjContact* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<mjtNum> var
+  X(pos);
+  X(frame);
+  X(friction);
+  X(solref);
+  X(solimp);
+  X(H);
+  #undef X
+};
+
+using MjContactWrapper = MjWrapper<raw::MjContact>;
+
+template <>
+struct enable_if_mj_struct<raw::MjContact> { using type = void; };
+
+template <>
+class MjStructList<raw::MjContact> : public StructListBase<raw::MjContact> {
+ public:
+  MjStructList(raw::MjContact* ptr, int num, pybind11::handle owner);
+
+  using StructListBase::operator[];
+  using StructListBase::size;
+  MjStructList Slice(pybind11::slice slice) {
+    return MjStructList(*this, slice);
+  }
+
+#define X(type, var) pybind11::array_t<type> var
+  X(mjtNum, dist);
+  X(mjtNum, pos);
+  X(mjtNum, frame);
+  X(mjtNum, includemargin);
+  X(mjtNum, friction);
+  X(mjtNum, solref);
+  X(mjtNum, solimp);
+  X(mjtNum, mu);
+  X(mjtNum, H);
+  X(int, dim);
+  X(int, geom1);
+  X(int, geom2);
+  X(int, exclude);
+  X(int, efc_address);
+#undef X
+
+ protected:
+  MjStructList(MjStructList& other, pybind11::slice slice);
+};
+
+using MjContactList = MjStructList<raw::MjContact>;
+
+template <>
+struct py_array_or_tuple<raw::MjContact> {
+  using type = MjContactList;
+};
+
+template <>
+struct is_mj_struct_list<raw::MjContact> {
+  static constexpr bool value = true;
+};
+
+// ==================== MJDATA =================================================
+template <>
+class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
+ public:
+  explicit MjWrapper(const MjModelWrapper& model);
+  MjWrapper(const MjWrapper& other);
+  MjWrapper(MjWrapper&&);
+  ~MjWrapper();
+
+  MjDataIndexer& indexer() { return indexer_; }
+
+  void Serialize(std::ostream& output) const;
+  static MjWrapper<raw::MjData> Deserialize(std::istream& input);
+
+  static constexpr char kFromRawPointer[] =
+      "__MUJOCO_STRUCTS_MJDATAWRAPPER_LOOKUP";
+  static MjWrapper* FromRawPointer(raw::MjData* m) noexcept;
+
+#define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
+  MJDATA_POINTERS
+#undef X
+
+  py_array_or_tuple_t<raw::MjWarningStat> warning;
+  py_array_or_tuple_t<raw::MjTimerStat> timer;
+  py_array_or_tuple_t<raw::MjSolverStat> solver;
+  py_array_or_tuple_t<mjtNum> solver_fwdinv;
+  py_array_or_tuple_t<mjtNum> energy;
+
+ private:
+  // Internal constructor which takes ownership of given mjData pointer.
+  // Used for deserialization.
+  explicit MjWrapper(MjDataMetadata&& metadata, raw::MjData* d);
+  raw::MjData* Copy() const;
+
+  MjDataMetadata metadata_;
+  MjDataIndexer indexer_;
+};
+
+using MjDataWrapper = MjWrapper<raw::MjData>;
+
+template <>
+struct enable_if_mj_struct<raw::MjData> { using type = void; };
+
+// ==================== MJVPERTURB =============================================
+template <>
+class MjWrapper<raw::MjvPerturb> : public WrapperBase<raw::MjvPerturb> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<mjtNum> var
+  X(refpos);
+  X(refquat);
+  X(localpos);
+  #undef X
+};
+
+using MjvPerturbWrapper = MjWrapper<raw::MjvPerturb>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvPerturb> { using type = void; };
+
+// ==================== MJVCAMERA ==============================================
+template <>
+class MjWrapper<raw::MjvCamera> : public WrapperBase<raw::MjvCamera> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<mjtNum> var
+  X(lookat);
+  #undef X
+};
+
+using MjvCameraWrapper = MjWrapper<raw::MjvCamera>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvCamera> { using type = void; };
+
+// ==================== MJVGLCAMERA ============================================
+template <>
+class MjWrapper<raw::MjvGLCamera> : public WrapperBase<raw::MjvGLCamera> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjvGLCamera* ptr, pybind11::handle owner);
+  explicit MjWrapper(raw::MjvGLCamera&& other);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<float> var
+  X(pos);
+  X(forward);
+  X(up);
+  #undef X
+};
+
+using MjvGLCameraWrapper = MjWrapper<raw::MjvGLCamera>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvGLCamera> { using type = void; };
+
+// ==================== MJVGEOM ================================================
+template <>
+class MjWrapper<raw::MjvGeom> : public WrapperBase<raw::MjvGeom> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjvGeom* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<float> var
+  X(texrepeat);
+  X(size);
+  X(pos);
+  X(mat);
+  X(rgba);
+  #undef X
+};
+
+using MjvGeomWrapper = MjWrapper<raw::MjvGeom>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvGeom> { using type = void; };
+
+// ==================== MJVLIGHT ===============================================
+template <>
+class MjWrapper<raw::MjvLight> : public WrapperBase<raw::MjvLight> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjvLight* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<float> var
+  X(pos);
+  X(dir);
+  X(attenuation);
+  X(ambient);
+  X(diffuse);
+  X(specular);
+  #undef X
+};
+
+using MjvLightWrapper = MjWrapper<raw::MjvLight>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvLight> { using type = void; };
+
+// ==================== MJVOPTION ==============================================
+template <>
+class MjWrapper<raw::MjvOption> : public WrapperBase<raw::MjvOption> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  ~MjWrapper() = default;
+
+  #define X(var) py_array_or_tuple_t<mjtByte> var
+  X(geomgroup);
+  X(sitegroup);
+  X(jointgroup);
+  X(tendongroup);
+  X(actuatorgroup);
+  X(flags);
+  #undef X
+};
+
+using MjvOptionWrapper = MjWrapper<raw::MjvOption>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvOption> { using type = void; };
+
+// ==================== MJVSCENE ===============================================
+template <>
+class MjWrapper<raw::MjvScene> : public WrapperBase<raw::MjvScene> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(const MjModelWrapper& model, int maxgeom);
+  ~MjWrapper() = default;
+
+  int nskinvert;
+
+  #define X(dtype, var) py_array_or_tuple_t<dtype> var
+  X(mjvGeom, geoms);
+  X(int, geomorder);
+  X(int, skinfacenum);
+  X(int, skinvertadr);
+  X(int, skinvertnum);
+  X(float, skinvert);
+  X(float, skinnormal);
+  #undef X
+
+  #define X(dtype, var) py_array_or_tuple_t<dtype> var
+  X(raw::MjvLight, lights);
+  X(raw::MjvGLCamera, camera);
+  X(float, translate);
+  X(float, rotate);
+  X(mjtByte, flags);
+  X(float, framergb);
+  #undef X
+};
+
+using MjvSceneWrapper = MjWrapper<raw::MjvScene>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvScene> { using type = void; };
+
+// ==================== MJVFIGURE ==============================================
+template <>
+class MjWrapper<raw::MjvFigure> : public WrapperBase<raw::MjvFigure> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  ~MjWrapper() = default;
+
+  #define X(dtype, var) py_array_or_tuple_t<dtype> var
+  X(int, flg_ticklabel);
+  X(int, gridsize);
+  X(float, gridrgb);
+  X(float, figurergba);
+  X(float, panergba);
+  X(float, legendrgba);
+  X(float, textrgb);
+  X(float, linergb);
+  X(float, range);
+  X(int, highlight);
+  X(int, linepnt);
+  X(float, linedata);
+  X(int, xaxispixel);
+  X(int, yaxispixel);
+  X(float, xaxisdata);
+  X(float, yaxisdata);
+  #undef X
+
+  pybind11::array linename;
+};
+
+using MjvFigureWrapper = MjWrapper<raw::MjvFigure>;
+
+template <>
+struct enable_if_mj_struct<raw::MjvFigure> { using type = void; };
+}  // namespace _impl
+
+template <typename T>
+using MjWrapper = typename std::conditional_t<
+    std::is_const_v<T>, const _impl::MjWrapper<std::remove_const_t<T>>,
+    _impl::MjWrapper<T>>;
+
+template <typename T>
+using enable_if_mj_struct_t =
+    typename _impl::enable_if_mj_struct<std::remove_const_t<T>>::type;
+
+using _impl::MjOptionWrapper;
+using _impl::MjVisualHeadlightWrapper;
+using _impl::MjVisualRgbaWrapper;
+using _impl::MjVisualWrapper;
+using _impl::MjStatisticWrapper;
+using _impl::MjWarningStatWrapper;
+using _impl::MjTimerStatWrapper;
+using _impl::MjSolverStatWrapper;
+using _impl::MjModelWrapper;
+using _impl::MjDataWrapper;
+using _impl::MjContactWrapper;
+using _impl::MjvPerturbWrapper;
+using _impl::MjvCameraWrapper;
+using _impl::MjvGLCameraWrapper;
+using _impl::MjvGeomWrapper;
+using _impl::MjvLightWrapper;
+using _impl::MjvOptionWrapper;
+using _impl::MjvSceneWrapper;
+using _impl::MjvFigureWrapper;
+
+template <typename T>
+using MjStructList = typename std::conditional_t<
+    std::is_const_v<T>, const _impl::MjStructList<std::remove_const_t<T>>,
+    _impl::MjStructList<T>>;
+
+template <typename T>
+static constexpr bool is_mj_struct_list_v =
+    _impl::is_mj_struct_list<std::remove_const_t<T>>::value;
+
+using _impl::MjContactList;
+using _impl::MjWarningStatList;
+using _impl::MjTimerStatList;
+using _impl::MjSolverStatList;
+
+// ==================== HELPER FUNCTIONS FOR BINDING ARRAYS ====================
+
+// Python array initialization.
+// If T is an arithmetic (i.e. numeric) type, returns a NumPy array that wraps
+// around a preexisting data buffer.
+template <typename T, typename Shape>
+std::enable_if_t<std::is_arithmetic_v<T>, pybind11::array_t<T>>
+static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
+  int size = 1;
+  for (const auto& i : shape) {
+    size *= i;
+  }
+  if (shape.empty() || size == 0) {
+    return pybind11::array_t<T>(shape);
+  } else {
+    return pybind11::array_t<T>(shape, buf, owner);
+  }
+}
+
+// Same as above, but where we can determine array dimensions through the
+// C array type directly.
+template <typename T, typename Int, Int N>
+std::enable_if_t<(N > 0) && std::is_arithmetic_v<T>, pybind11::array_t<T>>
+static InitPyArray(T (&buf)[N], pybind11::handle owner) {
+  return pybind11::array_t<T>({N}, &buf[0], owner);
+}
+
+template <typename T, typename Int, Int N1, Int N2>
+std::enable_if_t<(N1*N2 > 0) && std::is_arithmetic_v<T>, pybind11::array_t<T>>
+static InitPyArray(T (&buf)[N1][N2], pybind11::handle owner) {
+  return pybind11::array_t<T>(
+      {N1, N2}, &buf[0][0], owner);
+}
+
+template <typename T, typename Shape>
+std::enable_if_t<is_mj_struct_list_v<T>, MjStructList<T>>
+static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
+  return MjStructList<T>(buf, shape[0], owner);
+}
+
+// For arrays of non-arithmetic type, we create tuple of tuples of MjWrapper<T>.
+template <typename T, typename Shape>
+std::enable_if_t<!std::is_arithmetic_v<T> && !is_mj_struct_list_v<T>,
+                 pybind11::tuple>
+static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
+  int size = 1;
+  for (const auto& i : shape) {
+    size *= i;
+  }
+  if (shape.empty() || !size) {
+    return pybind11::tuple();
+  }
+
+  pybind11::list out;
+  const auto n = shape[0];
+  if (shape.size() == 1) {
+    for (int i = 0; i < n; ++i) {
+      out.append(MjWrapper<T>(&buf[i], owner));
+    }
+  } else {
+    auto block_shape = absl::MakeConstSpan(shape).subspan(1);
+    auto block_size = std::accumulate(
+        block_shape.begin(), block_shape.end(),
+        1, std::multiplies<decltype(n)>());
+    for (int i = 0; i < n; ++i) {
+      out.append(InitPyArray(block_shape, &buf[i * block_size], owner));
+    }
+  }
+  return out;
+}
+
+// Same as above, but where we can determine array dimensions through the
+// C array type directly.
+template <typename T, typename Int, Int N>
+std::enable_if_t<(N > 0) && !std::is_arithmetic_v<T> &&
+                 !std::is_array_v<T> && !is_mj_struct_list_v<T>,
+                 pybind11::tuple>
+static InitPyArray(T (&buf)[N], pybind11::handle owner) {
+  return InitPyArray(std::array{N}, buf, owner);
+}
+
+template <typename T, typename Int, Int N1, Int N2>
+std::enable_if_t<(N1*N2 > 0) && !std::is_arithmetic_v<T> &&
+                 !std::is_array_v<T> && !is_mj_struct_list_v<T>,
+                 pybind11::tuple>
+static InitPyArray(T (&buf)[N1][N2], pybind11::handle owner) {
+  return InitPyArray(std::array{N1, N2}, buf, owner);
+}
+
+template <typename T, typename Int, Int N>
+std::enable_if_t<is_mj_struct_list_v<T>, MjStructList<T>>
+static InitPyArray(T (&buf)[N], pybind11::handle owner) {
+  return MjStructList<T>(buf, N, owner);
+}
+
+// Helpers for defining array/tuple properties in pybind11 classes.
+//
+// Defines a NumPy array property of a Python class that supports assignments.
+// Specifically, we implement the setter such that `obj.arr = val` is the same
+// as `obj.arr[:] = val`.
+//
+// Use `DefinePyArray(c, "somearray", &MjStructHolder<C>::somearray)`
+// as a drop-in replacement for
+// `c.def_readonly("somearray", &MjStructHolder<C>::somearray)`.
+template <typename T, typename C, typename... O>
+static void DefinePyArray(pybind11::class_<C, O...> c, const char* name,
+                          pybind11::array_t<T> C::* arr) {
+  namespace py = pybind11;
+  c.def_property(
+      name,
+      [arr](const C& wrapper) { return wrapper.*arr; },
+      [arr](const C& wrapper, py::handle rhs) -> void {
+        (wrapper.*arr)[py::slice(py::none(), py::none(), py::none())] = rhs;
+      }
+  );
+}
+
+// For array of non-arithmetic type, we bind to tuples rather than NumPy array.
+// These can't be assigned to directly so we just use def_readonly.
+template <typename T, typename C, typename... O>
+static void DefinePyArray(pybind11::class_<C, O...> c, const char* name,
+                          T C::* arr) {
+  c.def_property_readonly(
+      name,
+      [arr](const C& wrapper) -> auto& { return wrapper.*arr; });
+}
+
+template <typename Raw, int N, typename C, typename... O>
+static void DefinePyStr(pybind11::class_<C, O...> c, const char* name,
+                        char (Raw::* arr)[N]) {
+  c.def_property(
+      name,
+      [arr](const C& c) { return pybind11::str(c.get()->*arr); },
+      [name = std::string(name), arr](C& c, std::string_view rhs) {
+        constexpr int kMaxLen = sizeof(c.get()->*arr);
+        const int actual_len = rhs.size();
+        if (actual_len >= kMaxLen) {
+          std::ostringstream msg;
+          msg << "len(" << name << ") cannot exceed " << kMaxLen - 1
+              << ": got length " << actual_len;
+          throw pybind11::value_error(msg.str());
+        }
+        rhs.copy(c.get()->*arr, actual_len);
+        (c.get()->*arr)[actual_len] = '\0';
+      });
+}
+
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_STRUCTS_H_

python/mujoco/util/CMakeLists.txt

@@ -0,0 +1,96 @@
+# 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
+#
+#     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.
+
+if(MSVC AND MSVC_VERSION GREATER_EQUAL 1927)
+  set(CMAKE_CXX_STANDARD 20) # For forceinline lambdas.
+else()
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+# INTERPROCEDURAL_OPTIMIZATION is enforced when enabled.
+set(CMAKE_POLICY_DEFAULT_CMP0069 NEW)
+
+if(APPLE)
+  add_compile_options(-Werror=partial-availability -Werror=unguarded-availability)
+  add_link_options(-Wl,-no_weak_imports)
+endif()
+
+add_library(crossplatform INTERFACE crossplatform.h)
+set_target_properties(crossplatform PROPERTIES PUBLIC_HEADER crossplatform.h)
+target_include_directories(crossplatform INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
+
+add_library(array_traits INTERFACE array_traits.h)
+set_target_properties(array_traits PROPERTIES PUBLIC_HEADER array_traits.h)
+target_include_directories(array_traits INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
+target_link_libraries(array_traits INTERFACE crossplatform Eigen3::Eigen)
+
+add_library(func_traits INTERFACE func_traits.h)
+set_target_properties(func_traits PROPERTIES PUBLIC_HEADER func_traits.h)
+target_include_directories(func_traits INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
+
+add_library(tuple_tools INTERFACE tuple_tools.h)
+set_target_properties(tuple_tools PROPERTIES PUBLIC_HEADER tuple_tools.h)
+target_include_directories(tuple_tools INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
+target_link_libraries(tuple_tools INTERFACE crossplatform)
+
+add_library(func_wrap INTERFACE func_wrap.h)
+set_target_properties(func_wrap PROPERTIES PUBLIC_HEADER func_wrap.h)
+target_include_directories(func_wrap INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
+target_link_libraries(
+  func_wrap
+  INTERFACE crossplatform
+            Eigen3::Eigen
+            array_traits
+            func_traits
+)
+
+if(MUJOCO_TEST_PYTHON_UTIL)
+  add_executable(array_traits_test array_traits_test.cc)
+  target_link_libraries(
+    array_traits_test
+    array_traits
+    gmock
+    gtest_main
+  )
+  gtest_add_tests(TARGET array_traits_test SOURCES array_traits_test.cc)
+
+  add_executable(func_traits_test func_traits_test.cc)
+  target_link_libraries(
+    func_traits_test
+    func_traits
+    gmock
+    gtest_main
+  )
+  gtest_add_tests(TARGET func_traits_test SOURCES func_traits_test.cc)
+
+  add_executable(func_wrap_test func_wrap_test.cc)
+  target_link_libraries(
+    func_wrap_test
+    func_wrap
+    gmock
+    gtest_main
+  )
+  gtest_add_tests(TARGET func_wrap_test SOURCES func_wrap_test.cc)
+
+  add_executable(tuple_tools_test tuple_tools_test.cc)
+  target_link_libraries(
+    tuple_tools_test
+    func_wrap
+    gmock
+    gtest_main
+  )
+  gtest_add_tests(TARGET tuple_tools_test SOURCES tuple_tools_test.cc)
+endif()

python/mujoco/util/array_traits.h

@@ -0,0 +1,140 @@
+// 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_PYTHON_UTIL_ARRAY_TRAITS_H_
+#define MUJOCO_PYTHON_UTIL_ARRAY_TRAITS_H_
+
+#include <type_traits>
+
+#include "crossplatform.h"
+#include <Eigen/Eigen>
+#include <unsupported/Eigen/CXX11/Tensor>
+
+namespace mujoco::util {
+
+// Forward declaration so that the public interface appears at the top of file.
+namespace _impl {
+template <typename T, int... N> struct c_array;
+template <typename T> struct c_array_traits;
+}  // namespace _impl
+
+// Array type from scalar type and extents. This is intended to be used to
+// deduce array extents as template integer parameters.
+// For example c_array_t<double, 9, 4, 7> is the same as double[9][4][7].
+template <typename T, int... N>
+using c_array_t = typename _impl::c_array<T, N...>::type;
+
+// Scalar type from an array, reference-to-array, or pointer-to-array type.
+template <typename T>
+using array_scalar_t = typename _impl::c_array_traits<T>::scalar_type;
+
+// The number of dimensions of an array type. If the array type is regarded as
+// a tensor then this corresponds to the tensor rank.
+template <typename T>
+static constexpr int array_ndim_v = _impl::c_array_traits<T>::ndim;
+
+// Makes an Eigen::Matrix or Eigen::Tensor object with the same data type and
+// shape as the given array type. The Eigen object returned is always row-major
+// (i.e. C ordering) and dense. If the array type is one- or two-dimensional
+// then an Eigen::Matrix with compile-time constant shape is returned.
+// Otherwise, an Eigen::Tensor is returned.
+template <typename ArrType>
+constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
+  return _impl::c_array_traits<ArrType>::template MakeEigen<>();
+}
+
+template <typename ArrType>
+using array_eigen_t = std::conditional_t<
+    std::is_const_v<ArrType>,
+    const decltype(MakeEigen<std::remove_const_t<ArrType>>()),
+    // Still need remove_const here since the type substitution in the false
+    // branch always occurs regardless of the condition, and
+    // MakeEigen<const T>() is invalid.
+    decltype(MakeEigen<std::remove_const_t<ArrType>>())>;
+
+// =====================================================================
+// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
+// =====================================================================
+namespace _impl {
+template <typename T, int... N>
+struct c_array {};
+
+template <typename T>
+struct c_array<T> {
+  using type = T;
+  static constexpr int ndim = 0;
+};
+
+template <typename T, int M, int... N>
+struct c_array<T, M, N...> {
+  using type = typename c_array<T, N...>::type[M];
+  static constexpr int ndim = c_array<T, N...>::ndim + 1;
+};
+
+template <typename T>
+struct c_array_traits {
+  static constexpr int ndim = 0;
+  using scalar_type = std::remove_reference_t<T>;
+
+  template <int... N>
+  static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
+    if constexpr (c_array<T, N...>::ndim <= 2) {
+      return Eigen::Matrix<T, N..., Eigen::RowMajor>();
+    } else {
+      return Eigen::Tensor<
+          T, c_array<T, N...>::ndim, Eigen::RowMajor, Eigen::DenseIndex>(N...);
+    }
+  }
+};
+
+template <typename T, int N>
+struct c_array_traits<T[N]> {
+  // Recursively peel off the innermost extent.
+  static constexpr int ndim = c_array_traits<T>::ndim + 1;
+  using scalar_type = typename c_array_traits<T>::scalar_type;
+
+  template <int... M>
+  static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
+    return c_array_traits<T>::template MakeEigen<M..., N>();
+  }
+};
+
+template <typename T, int N>
+struct c_array_traits<T(&)[N]> {
+  // Delegate everything to the T[N] case.
+  static constexpr int ndim = c_array_traits<T[N]>::ndim;
+  using scalar_type = typename c_array_traits<T[N]>::scalar_type;
+
+  template <int... M>
+  static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
+    return c_array_traits<T[N]>::template MakeEigen<M...>();
+  }
+};
+
+template <typename T, int N>
+struct c_array_traits<T(*)[N]> {
+  // Delegate everything to the T[N] case.
+  static constexpr int ndim = c_array_traits<T[N]>::ndim;
+  using scalar_type = typename c_array_traits<T[N]>::scalar_type;
+
+  template <int... M>
+  static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
+    return c_array_traits<T[N]>::template MakeEigen<M...>();
+  }
+};
+
+}  // namespace _impl
+}  // namespace mujoco::util
+
+#endif  // MUJOCO_PYTHON_UTIL_ARRAY_TRAITS_H_

python/mujoco/util/array_traits_test.cc

@@ -0,0 +1,95 @@
+// 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 "array_traits.h"
+
+#include <type_traits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mujoco::util {
+namespace {
+
+using ::testing::ElementsAre;
+
+TEST(ArrayTraitsTest, CArrayType) {
+  struct Foo {};
+  static_assert(std::is_same_v<c_array_t<int, 3, 4, 5>,
+                               int[3][4][5]>);
+  static_assert(std::is_same_v<c_array_t<double>,
+                               double>);
+  static_assert(std::is_same_v<c_array_t<Foo, 7>,
+                               Foo[7]>);
+  static_assert(std::is_same_v<c_array_t<Foo*, 3, 4, 5, 6>,
+                               Foo*[3][4][5][6]>);
+}
+
+TEST(ArrayTraitsTest, ArrayNdim) {
+  struct Foo {};
+  EXPECT_EQ(array_ndim_v<int[3][4][5]>, 3);
+  EXPECT_EQ(array_ndim_v<double(*)[3][4]>, 2);
+  EXPECT_EQ(array_ndim_v<Foo*[3][4]>, 2);
+  EXPECT_EQ(array_ndim_v<Foo(&)[3][4][5][6]>, 4);
+}
+
+TEST(ArrayTraitsTest, ArrayScalarType) {
+  struct Foo {};
+  static_assert(std::is_same_v<
+      array_scalar_t<int[3][4][5]>,
+      int
+  >);
+  static_assert(std::is_same_v<
+      array_scalar_t<double(*)[3][4]>,
+      double
+  >);
+  static_assert(std::is_same_v<
+      array_scalar_t<Foo*[3][4]>,
+      Foo*
+  >);
+  static_assert(std::is_same_v<
+      array_scalar_t<Foo(&)[3][4][5][6]>,
+      Foo
+  >);
+}
+
+TEST(ArrayTraitsTest, MakeEigen) {
+  {
+    auto eigen = MakeEigen<float[3]>();
+    static_assert(std::is_same_v<
+        decltype(eigen),
+        Eigen::Vector3f
+    >);
+  }
+  {
+    auto eigen = MakeEigen<int[2][3]>();
+    static_assert(std::is_same_v<
+        decltype(eigen),
+        Eigen::Matrix<int, 2, 3, Eigen::RowMajor>
+    >);
+  }
+  {
+    auto eigen = MakeEigen<double[2][3][4]>();
+    static_assert(std::is_same_v<
+        decltype(eigen),
+        Eigen::Tensor<double, 3, Eigen::RowMajor, Eigen::DenseIndex>
+    >);
+    EXPECT_THAT(eigen.dimensions(), ElementsAre(2, 3, 4));
+  }
+}
+
+
+
+}  // namespace
+}  // namespace mujoco::util

python/mujoco/util/crossplatform.h

@@ -0,0 +1,47 @@
+// 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_PYTHON_UTIL_CROSSPLATFORM_H_
+#define MUJOCO_PYTHON_UTIL_CROSSPLATFORM_H_
+
+#ifdef __has_attribute
+#define MUJOCO_HAS_ATTRIBUTE(x) __has_attribute(x)
+#else
+#define MUJOCO_HAS_ATTRIBUTE(x) 0
+#endif
+
+#if MUJOCO_HAS_ATTRIBUTE(always_inline) || \
+    (defined(__GNUC__) && !defined(__clang__))
+#define MUJOCO_ALWAYS_INLINE __attribute__((always_inline))
+#define MUJOCO_ALWAYS_INLINE_LAMBDA MUJOCO_ALWAYS_INLINE
+#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE MUJOCO_ALWAYS_INLINE_LAMBDA mutable
+#elif defined(_MSC_VER)
+#define MUJOCO_ALWAYS_INLINE __forceinline
+#if _MSC_VER >= 1927 && _MSVC_LANG >= 202002L
+#define MUJOCO_ALWAYS_INLINE_LAMBDA [[msvc::forceinline]]
+#endif
+#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE mutable MUJOCO_ALWAYS_INLINE_LAMBDA
+#else
+#define MUJOCO_ALWAYS_INLINE
+#endif
+
+#ifndef MUJOCO_ALWAYS_INLINE_LAMBDA
+#define MUJOCO_ALWAYS_INLINE_LAMBDA
+#endif
+
+#ifndef MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE
+#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE
+#endif
+
+#endif  // MUJOCO_PYTHON_UTIL_CROSSPLATFORM_H_

python/mujoco/util/func_traits.h

@@ -0,0 +1,114 @@
+// 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_PYTHON_FUNC_TRAITS_H_
+#define MUJOCO_PYTHON_FUNC_TRAITS_H_
+
+#include <tuple>
+#include <type_traits>
+
+namespace mujoco::util {
+
+// Forward declaration so that the public interface appears at the top of file.
+namespace _impl {
+template <typename, typename = void> struct is_callable;
+template <typename, int, bool = false> struct func_arg;
+}  // namespace _impl
+
+// True if T is callable, i.e. if T is either a function pointer/reference or
+// is an instance of a type with operator().
+template <typename T>
+static constexpr bool is_callable_v =
+    _impl::is_callable<std::remove_reference_t<T>>::value;
+
+// Type of the Nth argument of a function or functor, where N=0 refers to the
+// first argument. If N exceeds the number arguments for F then
+// func_arg_t<F, N> is void.
+template <typename F, int N = 0>
+using func_arg_t = typename _impl::func_arg<F, N, (N > 0)>::type;
+
+template <typename F>
+static constexpr int func_arg_count_v = _impl::func_arg<F, 0>::count;
+
+// =====================================================================
+// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
+// =====================================================================
+namespace _impl {
+template <typename T, typename>
+struct is_callable {
+  static constexpr bool value = false;
+};
+
+template <typename T>
+struct is_callable<T, std::void_t<decltype(&T::operator())>> {
+  static constexpr bool value = true;
+};
+
+template <typename Return, typename... Args>
+struct is_callable<Return(Args...)> {
+  static constexpr bool value = true;
+};
+
+template <typename Return, typename... Args>
+struct is_callable<Return (*)(Args...)> {
+  static constexpr bool value = true;
+};
+
+// Support functors by looking at its member function Func::operator().
+template <typename Func, int N, bool Recursing>
+struct func_arg {
+  using call = decltype(
+      &std::remove_const_t<std::remove_reference_t<Func>>::operator());
+  using type = typename func_arg<call, N>::type;
+  static constexpr int count = func_arg<call, N>::count;
+};
+
+// Base case (N == 0) for function: resolve to Arg0.
+template <typename Ret, typename Arg0, typename... Args>
+struct func_arg<Ret(Arg0, Args...), 0> {
+  using type = Arg0;
+  static constexpr int count = 1 + std::tuple_size_v<std::tuple<Args...>>;
+};
+
+// Recursive case (N > 0) for function: discard Arg0 it and resolve to N-1.
+template <typename Ret, int N, typename Arg0, typename... Args>
+struct func_arg<Ret(Arg0, Args...), N, true> {
+  using type = typename func_arg<Ret(Args...), N - 1, (N > 1)>::type;
+  static constexpr int count = 1 + std::tuple_size_v<std::tuple<Args...>>;
+};
+
+// Specialization for non-const member functions.
+template <typename C, typename Ret, int N, typename... Args>
+struct func_arg<Ret (C::*)(Args...), N> {
+  using type = typename func_arg<Ret(Args...), N, (N > 0)>::type;
+  static constexpr int count = std::tuple_size_v<std::tuple<Args...>>;
+};
+
+// Specialization for const member functions (matches lambda::operator()).
+template <typename C, typename Ret, int N, typename... Args>
+struct func_arg<Ret (C::*)(Args...) const, N> {
+  using type = typename func_arg<Ret(Args...), N, (N > 0)>::type;
+  static constexpr int count = std::tuple_size_v<std::tuple<Args...>>;
+};
+
+// Functions with no argument: always resolve to void.
+template <typename Ret, int N, bool Recursing>
+struct func_arg<Ret(), N, Recursing> {
+  using type = void;
+  static constexpr int count = 0;
+};
+}  // namespace _impl
+}  // namespace mujoco::util
+
+#endif  // MUJOCO_PYTHON_FUNC_TRAITS_H_

python/mujoco/util/func_traits_test.cc

@@ -0,0 +1,156 @@
+// 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 "func_traits.h"
+
+#include <functional>
+#include <type_traits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mujoco::util {
+namespace {
+
+TEST(FuncTraitsTest, IsCallable) {
+  EXPECT_FALSE(is_callable_v<int>);
+  EXPECT_TRUE(is_callable_v<double(double)>);
+  EXPECT_TRUE(is_callable_v<void(*)(void)>);
+  EXPECT_TRUE((is_callable_v<int(&)(int, float)>));
+  EXPECT_TRUE(is_callable_v<std::function<void(void)>>);
+  {
+    auto lambda = [](){};
+    EXPECT_TRUE(is_callable_v<decltype(lambda)>);
+  }
+  {
+    auto mutable_lambda = []() mutable {};
+    EXPECT_TRUE(is_callable_v<decltype(mutable_lambda)>);
+  }
+  {
+    struct Functor { void operator()() {} };
+    EXPECT_TRUE(is_callable_v<Functor>);
+  }
+  {
+    struct ConstFunctor { void operator()() const {} };
+    EXPECT_TRUE(is_callable_v<ConstFunctor>);
+  }
+  {
+    struct NotCallable { void Foo() {} };
+    EXPECT_FALSE(is_callable_v<NotCallable>);
+    EXPECT_FALSE(is_callable_v<decltype(&NotCallable::Foo)>);
+  }
+}
+
+TEST(FuncTraitsTest, FuncArgType) {
+  static_assert(std::is_same_v<
+      func_arg_t<void()>,
+      void
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int)>,
+      int
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&)>,
+      int&&
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&), 0>,
+      int&&
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&), 1>,
+      char&
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&), 2>,
+      const float&
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&), 3>,
+      void
+  >);
+  static_assert(std::is_same_v<
+      func_arg_t<bool(int&&, char&, const float&), 7>,
+      void
+  >);
+
+  {
+    auto lambda = [](bool, double&, float&&){};
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda)>,
+        bool
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda), 0>,
+        bool
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda), 1>,
+        double&
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda), 2>,
+        float&&
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda), 3>,
+        void
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(lambda), 10>,
+        void
+    >);
+  }
+  {
+    struct Functor { void operator()(char, void*) {} };
+    static_assert(std::is_same_v<
+        func_arg_t<Functor>,
+        char
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<Functor, 0>,
+        char
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<Functor, 1>,
+        void*
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<Functor, 2>,
+        void
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<Functor, 5>,
+        void
+    >);
+  }
+}
+
+TEST(FuncTraitsTest, FuncArgCount) {
+  EXPECT_EQ(func_arg_count_v<void()>, 0);
+  EXPECT_EQ(func_arg_count_v<bool(int)>, 1);
+  EXPECT_EQ(func_arg_count_v<bool(int&&, char&, const float&)>, 3);
+  {
+    auto lambda = [](bool, double&, float&&){};
+    EXPECT_EQ(func_arg_count_v<decltype(lambda)>, 3);
+  }
+  {
+    struct Functor { void operator()(char, void*) {} };
+    EXPECT_EQ(func_arg_count_v<Functor>, 2);
+  }
+}
+
+}  // namespace
+}  // namespace mujoco::util

python/mujoco/util/func_wrap.h

@@ -0,0 +1,193 @@
+// 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_PYTHON_UTIL_FUNC_WRAP_H_
+#define MUJOCO_PYTHON_UTIL_FUNC_WRAP_H_
+
+#include <type_traits>
+#include <utility>
+
+#include <Eigen/Eigen>
+#include "array_traits.h"
+#include "crossplatform.h"
+#include "func_traits.h"
+
+namespace mujoco::util {
+
+// Represents an argument type T of a C++ function that is callable from Python
+// via pybind11. Template specializations of this struct defines how to unwrap
+// arguments from pybind11 before passing them the underlying C++ function.
+//
+// This is used to help bind functions whose argument types are not related
+// to the types that are registered with pybind11, and where it is not
+// desirable/appropriate for the function's argument types to be registered.
+//
+// Usage:
+//   In the compilation unit that binds a function, specialize the template for
+//   each argument type that needs to be unwrapped, e.g. if a function expects
+//   an argument of type SomeArgType, but the type that is known to pybind11
+//   is SomeWrapperType, then the specialization looks like:
+//
+//      template <> wrapped py_arg<SomeArgType*> {
+//        static constexpr SomeArgType* unwrap(SomeWrapperType* wrapped_arg) {
+//          return wrapped_arg->get_the_underlying_thing();
+//        }
+//      };
+template <typename T, typename = void>
+struct wrapped {
+  MUJOCO_ALWAYS_INLINE
+  static constexpr T unwrap(T arg)  {
+    return arg;
+  }
+};
+
+// The wrapper type for T that can be unwrapped via wrapped<T>::unwrap.
+template <typename T> using wrapper_t =
+    typename util::func_arg_t<decltype(wrapped<T>::unwrap)>;
+
+namespace _impl {
+
+template <typename T, typename = void>
+struct arg_type_deducer {
+  static_assert(util::is_callable_v<T>, "not a Callable type");
+  template <typename WrapOp>
+  static constexpr auto WrapFunc(T&& callable) {
+    using Call = decltype(&std::remove_reference_t<T>::operator());
+    return arg_type_deducer<T, Call>::template WrapFunc<WrapOp>(
+        std::forward<T>(callable));
+  }
+};
+
+template <typename Return, typename... Args>
+using func_t = Return(Args...);
+
+// Specializations to deduce argument types for vanilla function references.
+template <typename Return, typename... Args>
+struct arg_type_deducer<func_t<Return, Args...>&> {
+  template <typename WrapOp>
+  static constexpr auto WrapFunc(Return (&func)(Args...)) {
+    return WrapOp::template WrapFunc<Return, Args...>(func);
+  }
+};
+
+// Specializations to deduce argument types for vanilla function pointers.
+template <typename Return, typename... Args>
+struct arg_type_deducer<Return (*)(Args...)> {
+  template <typename WrapOp>
+  static constexpr auto WrapFunc(Return (*func)(Args...)) {
+    return WrapOp::template WrapFunc<Return, Args...>(*func);
+  }
+};
+
+// Specialization to deduce argument types for non-const operator().
+template <typename Callable, typename Return, typename... Args>
+struct arg_type_deducer<
+    Callable, Return (std::remove_reference_t<Callable>::*)(Args...)> {
+  template <typename WrapOp>
+  static constexpr auto WrapFunc(Callable&& callable) {
+    return WrapOp::template WrapFunc<Return, Args...>(
+        std::forward<Callable>(callable));
+  }
+};
+
+// Specialization to deduce argument types for const operator().
+template <typename Callable, typename Return, typename... Args>
+struct arg_type_deducer<
+    Callable, Return (std::remove_reference_t<Callable>::*)(Args...) const> {
+  template <typename WrapOp>
+  static constexpr auto WrapFunc(Callable&& callable) {
+    return WrapOp::template WrapFunc<Return, Args...>(
+        std::forward<Callable>(callable));
+  }
+};
+
+template <typename WrapOp, typename Callable>
+constexpr auto WrapFunc(Callable&& callable) {
+  return arg_type_deducer<Callable>::template WrapFunc<WrapOp>(
+      std::forward<Callable>(callable));
+}
+
+struct UnwrapArgs {
+  template <typename Return, typename... Args, typename Callable>
+  static constexpr auto WrapFunc(Callable&& callable) {
+    return [callable](wrapper_t<Args>... wrapped_args)
+              MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
+      return callable(wrapped<Args>::unwrap(wrapped_args)...);
+    };
+  }
+};
+
+template <bool OutArgProvided>
+struct ReturnArrayArg0 {
+  template <typename Return, typename OutArg, typename... InArgs,
+            typename Callable>
+  static constexpr auto WrapFunc(Callable&& callable) {
+    using OutArray = std::remove_reference_t<std::remove_pointer_t<OutArg>>;
+    using OutScalar = util::array_scalar_t<OutArray>;
+    static_assert(
+        std::is_array_v<OutArray> && std::is_arithmetic_v<OutScalar>,
+        "output is not an array of arithmetic type");
+    static_assert(
+        std::is_void_v<Return>,
+        "callable under ReturnArrayArg0 cannot return a value");
+
+    // MSVC has a bug with `if constexpr`, as a workaround we precompute the
+    // condition into a constexpr variable first.
+    // https://developercommunity.visualstudio.com/t/1509806
+    constexpr bool OutArgIsRef = std::is_same_v<OutArg, OutArray&>;
+
+    if constexpr (OutArgProvided) {
+      using EigenOutType = Eigen::Ref<decltype(util::MakeEigen<OutArray>())>;
+      return [callable](InArgs... args, EigenOutType eigen_out)
+          MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
+        if constexpr (OutArgIsRef) {
+          callable(*reinterpret_cast<OutArray*>(eigen_out.data()), args...);
+        } else {
+          callable(reinterpret_cast<OutArray*>(eigen_out.data()), args...);
+        }
+      };
+    } else {
+      return [callable](InArgs... args) MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
+        auto eigen_out = util::MakeEigen<OutArray>();
+        if constexpr (OutArgIsRef) {
+          callable(*reinterpret_cast<OutArray*>(eigen_out.data()), args...);
+        } else {
+          callable(reinterpret_cast<OutArray*>(eigen_out.data()), args...);
+        }
+        return eigen_out;
+      };
+    }
+  }
+};
+
+}  // namespace _impl
+
+// Makes a callable that unwraps each argument before passing it to the
+// given callable. Specifically, given f(T1 x1, T2 x2, ...) this function
+// returns a callable
+// g(wrapper_t<T1> w1, wrapper_t<T2> w2, ...) = f(unwrap(w1), unwrap(w2), ...).
+template <typename Callable>
+constexpr auto UnwrapArgs(Callable&& callable) {
+  return _impl::WrapFunc<_impl::UnwrapArgs>(std::forward<Callable>(callable));
+}
+
+template <bool OutArgProvided = false, typename Callable>
+constexpr auto ReturnArrayArg0(Callable&& callable) {
+  return _impl::WrapFunc<_impl::ReturnArrayArg0<OutArgProvided>>(
+      std::forward<Callable>(callable));
+}
+
+}  // namespace mujoco::util
+
+#endif  // MUJOCO_PYTHON_UTIL_FUNC_WRAP_H_

python/mujoco/util/func_wrap_test.cc

@@ -0,0 +1,102 @@
+// 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 "func_wrap.h"
+
+#include <string>
+#include <type_traits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <Eigen/Eigen>
+#include "func_traits.h"
+
+namespace {
+struct BoxedDouble {
+  double value;
+};
+struct BoxedInt {
+  int value;
+};
+}  // namespace
+
+namespace mujoco::util {
+template <> struct wrapped<BoxedDouble> {
+  static BoxedDouble unwrap(const std::string& wrapped) {
+    return BoxedDouble{std::stod(wrapped)};
+  }
+};
+template <> struct wrapped<BoxedInt> {
+  static BoxedInt unwrap(int wrapped) {
+    return BoxedInt{wrapped};
+  }
+};
+template <typename T, int N> struct wrapped<T(*)[N]> {
+  using Array = T[N];
+  static Array* unwrap(Eigen::Ref<Eigen::Vector<T, N>> wrapped) {
+    return reinterpret_cast<Array*>(wrapped.data());
+  }
+};
+template <typename T, int N> struct wrapped<const T(*)[N]> {
+  using Array = const T[N];
+  static Array* unwrap(const Eigen::Vector<T, N>& wrapped) {
+    return reinterpret_cast<Array*>(wrapped.data());
+  }
+};
+}  // namespace mujoco::util
+
+namespace {
+using ::mujoco::util::func_arg_t;
+using ::mujoco::util::UnwrapArgs;
+using ::mujoco::util::ReturnArrayArg0;
+
+double add(BoxedDouble x, float y, BoxedInt z) {
+  return x.value + y + z.value;
+}
+
+void add_array4(double (*out)[4], const double (*x)[4], const double (*y)[4]) {
+  for (int i = 0; i < 4; ++i) {
+    (*out)[i] = (*x)[i] + (*y)[i];
+  }
+}
+
+TEST(FuncWrapTest, UnwrapArgs) {
+  {
+    auto wrapped_add = UnwrapArgs(add);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(wrapped_add), 0>, const std::string&
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(wrapped_add), 1>, float
+    >);
+    static_assert(std::is_same_v<
+        func_arg_t<decltype(wrapped_add), 2>, int
+    >);
+    // Use binary powers so that we can do exact floating point comparison.
+    EXPECT_EQ(wrapped_add("1.6e+1", 5e-1, 2), 18.5);
+  }
+  {
+    Eigen::Vector4d out;
+    UnwrapArgs(add_array4)(out, {1, 3, 5, 7}, {2, 6, 9, 11});
+    EXPECT_THAT(out, ::testing::ElementsAre(3, 9, 14, 18));
+  }
+}
+
+TEST(FuncWrapTest, ReturnArrayArg0) {
+  auto out = UnwrapArgs(ReturnArrayArg0(add_array4))({1, 3, 5, 7},
+                                                     {2, 6, 9, 11});
+  static_assert(std::is_same_v<decltype(out), Eigen::Vector4d>);
+  EXPECT_THAT(out, ::testing::ElementsAre(3, 9, 14, 18));
+}
+}  // namespace

python/mujoco/util/tuple_tools.h

@@ -0,0 +1,159 @@
+// 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_PYTHON_UTIL_TUPLE_TOOLS_H_
+#define MUJOCO_PYTHON_UTIL_TUPLE_TOOLS_H_
+
+#include <optional>
+#include <string_view>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
+#include "crossplatform.h"
+
+namespace mujoco::util {
+
+// Forward declaration so that the public interface appears at the top of file.
+namespace _impl {
+template <int N> struct head_slicer;
+template <int N> struct tail_slicer;
+}  // namespace _impl
+
+// Removes the first `Begin` elements from the from a tuple.
+// If `Begin` is negative, the resulting tuple is obtained by removing the first
+// `N - Begin` elements, where N is the size of the input tuple.
+//
+// For consistency with the <Begin, End> form (see below), the second optional
+// template argument can be explicitly spelled out as void,
+// e.g. tuple_slice<2, void>(tuple), to indicate that the slice takes every
+// element up to the end of the tuple.
+//
+// This function has the same semantic as Python's list slicing x[Begin:].
+template <int Begin, typename End = void, typename Tuple>
+MUJOCO_ALWAYS_INLINE
+constexpr auto tuple_slice(Tuple&& tuple) {
+  static_assert(
+      std::is_void_v<End>,
+      "argument End must be either omitted or explicitly specified as void");
+  constexpr int Size = std::tuple_size_v<std::remove_reference_t<Tuple>>;
+  static_assert(Begin >= -Size && Begin <= Size);
+  constexpr int NHead = (Begin >= 0) ? Begin : (Size + Begin);
+  return std::apply(_impl::head_slicer<NHead>(), std::forward<Tuple>(tuple));
+}
+
+// Extracts a contiguous slice of elements from a tuple so that the resulting
+// tuple begins with the element at index `Begin` and ends with the element
+// immediately before the one at index `End`. A negative value of `Begin` or
+// `End` is interpreted as indexing an element from the end of the input tuple,
+// where -1 refers to the last element.
+//
+// This function has the same semantic as Python's list slicing x[Begin:End].
+template <int Begin, int End, typename Tuple>
+MUJOCO_ALWAYS_INLINE
+constexpr auto tuple_slice(Tuple&& tuple) {
+  constexpr int Size = std::tuple_size_v<std::remove_reference_t<Tuple>>;
+  static_assert(End >= -Size && End <= Size);
+  constexpr int NTail = (End >= 0) ? (Size - End) : (-End);
+  static_assert(
+      (Begin >= -Size && Begin <= -NTail) ||
+      (Begin >= 0 && Begin <= Size - NTail),
+      "Begin should refer to an element that comes before End");
+  constexpr int NHead = (Begin >= 0) ? Begin : (Size + Begin);
+  return tuple_slice<NHead>(
+      std::apply(_impl::tail_slicer<NTail>(), std::forward<Tuple>(tuple)));
+}
+
+// Compile-time function to check whether a string occurs in a tuple.
+// Should ideally be declared consteval if we switch to C++20.
+template <typename Str, typename Tuple>
+static constexpr bool string_is_in_tuple(Str str, Tuple&& tuple) {
+  if constexpr (std::tuple_size_v<std::remove_reference_t<Tuple>> == 0) {
+    return false;
+  } else if (std::string_view(str) == std::string_view(std::get<0>(tuple))) {
+    return true;
+  } else {
+    return string_is_in_tuple(str, util::tuple_slice<1, void>(tuple));
+  }
+}
+
+// Compile-time function to check whether the elements of one tuple is a subset
+// another. Should ideally be declared consteval if we switch to C++20.
+template <typename Tuple1, typename Tuple2>
+constexpr bool is_subset_strings(Tuple1 tuple1, Tuple2 tuple2) {
+  if constexpr (std::tuple_size_v<Tuple1> == 0) {
+    return true;
+  } else if (string_is_in_tuple(std::get<0>(tuple1), tuple2)) {
+    return is_subset_strings(util::tuple_slice<1, void>(tuple1), tuple2);
+  } else {
+    return false;
+  }
+}
+
+// =====================================================================
+// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
+// =====================================================================
+namespace _impl{
+template <int N>
+struct head_slicer {
+  template <typename T, typename... U>
+  MUJOCO_ALWAYS_INLINE
+  constexpr auto operator()(T&& t, U&&... u) const {
+    return head_slicer<N-1>()(std::forward<U>(u)...);
+  }
+};
+
+template <>
+struct head_slicer<0> {
+  template <typename... T>
+  MUJOCO_ALWAYS_INLINE
+  constexpr auto operator()(T&&... t) const {
+    return std::forward_as_tuple(t...);
+  }
+};
+
+template <int N>
+struct move_head_to_tail {
+  template <typename T, typename... U>
+  MUJOCO_ALWAYS_INLINE
+  static constexpr auto move(T&& t, U&&... u) {
+    return move_head_to_tail<N-1>::move(
+        std::forward<U>(u)..., std::forward<T>(t));
+  }
+};
+
+template <>
+struct move_head_to_tail<0> {
+  template <typename... T>
+  MUJOCO_ALWAYS_INLINE
+  static constexpr auto move(T&&... t) {
+    return std::forward_as_tuple(t...);
+  }
+};
+
+template <int N>
+struct tail_slicer {
+  template <typename... T>
+  MUJOCO_ALWAYS_INLINE
+  constexpr auto operator()(T&&... t) const {
+    constexpr int Size = std::tuple_size_v<std::tuple<T...>>;
+    constexpr int NHead = Size - N;
+    return std::apply(head_slicer<N>(),
+        move_head_to_tail<NHead>::move(std::forward<T>(t)...));
+  }
+};
+}  // namespace _impl
+}  // namespace mujoco::util
+
+#endif  // MUJOCO_PYTHON_UTIL_TUPLE_TOOLS_H_

python/mujoco/util/tuple_tools_test.cc

@@ -0,0 +1,54 @@
+// 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 "tuple_tools.h"
+
+#include <tuple>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mujoco::util {
+namespace {
+
+TEST(TupleToolsTest, Slice) {
+  auto tuple = std::make_tuple(2, 3, 5, 7, 11, 13);
+  EXPECT_EQ((tuple_slice<1, 5>(tuple)), std::make_tuple(3, 5, 7, 11));
+  EXPECT_EQ((tuple_slice<2, 3>(tuple)), std::make_tuple(5));
+
+  // negative indices
+  EXPECT_EQ((tuple_slice<1, -2>(tuple)), std::make_tuple(3, 5, 7));
+  EXPECT_EQ((tuple_slice<-4, 6>(tuple)), std::make_tuple(5, 7, 11, 13));
+  EXPECT_EQ((tuple_slice<-3, -1>(tuple)), std::make_tuple(7, 11));
+
+  // empty slices
+  EXPECT_EQ((tuple_slice<0, 0>(tuple)), std::make_tuple());
+  EXPECT_EQ((tuple_slice<3, 3>(tuple)), std::make_tuple());
+  EXPECT_EQ((tuple_slice<-2, -2>(tuple)), std::make_tuple());
+
+  // specify void as the End argument
+  EXPECT_EQ((tuple_slice<2, void>(tuple)), std::make_tuple(5, 7, 11, 13));
+  EXPECT_EQ((tuple_slice<-2, void>(tuple)), std::make_tuple(11, 13));
+
+  // omit the End argument
+  EXPECT_EQ((tuple_slice<2>(tuple)), std::make_tuple(5, 7, 11, 13));
+  EXPECT_EQ((tuple_slice<-2>(tuple)), std::make_tuple(11, 13));
+
+  // empty input tuples
+  EXPECT_EQ((tuple_slice<0>(std::make_tuple())), std::make_tuple());
+  EXPECT_EQ((tuple_slice<0, 0>(std::make_tuple())), std::make_tuple());
+}
+
+}  // namespace
+}  // namespace mujoco::util

python/setup.py

@@ -0,0 +1,301 @@
+# 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.
+# ==============================================================================
+"""Install script for MuJoCo."""
+
+import fnmatch
+import os
+import platform
+import random
+import re
+import shutil
+import string
+import subprocess
+import sys
+import sysconfig
+
+import setuptools
+from setuptools import find_packages
+from setuptools import setup
+from setuptools.command import build_ext
+
+__version__ = '2.1.2'
+
+MUJOCO_CMAKE = 'MUJOCO_CMAKE'
+MUJOCO_CMAKE_ARGS = 'MUJOCO_CMAKE_ARGS'
+MUJOCO_PATH = 'MUJOCO_PATH'
+
+EXT_PREFIX = 'mujoco.'
+
+
+def get_long_description():
+  """Creates a long description for the package from bundled markdown files."""
+  current_dir = os.path.dirname('__file__')
+  with open(os.path.join(current_dir, 'README.md')) as f:
+    description = f.read()
+  try:
+    with open(os.path.join(current_dir, 'LICENSES_THIRD_PARTY.md')) as f:
+      description = f'{description}\n{f.read()}'
+  except FileNotFoundError:
+    pass
+  return description
+
+
+def get_mujoco_lib_pattern():
+  if platform.system() == 'Windows':
+    return 'mujoco*.lib'
+  elif platform.system() == 'Darwin':
+    return 'libmujoco*.dylib'
+  else:
+    return 'libmujoco*.so*'
+
+
+def get_external_lib_patterns():
+  if platform.system() == 'Windows':
+    return ['mujoco*.dll']
+  elif platform.system() == 'Darwin':
+    return ['libmujoco*.dylib']
+  else:
+    return ['libmujoco*.so*', 'libglew*.so']
+
+
+def start_and_end(iterable):
+  it = iter(iterable)
+  while True:
+    try:
+      first = next(it)
+      second = next(it)
+      yield first, second
+    except StopIteration:
+      return
+
+
+def tokenize_quoted_substr(input_string, quote_char, placeholders=None):
+  """Replace quoted substrings with random text placeholders with no spaces."""
+  # Matches quote characters not proceded with a backslash.
+  pattern = re.compile(r'(?<!\\)' + quote_char)
+  quote_positions = [m.start() for m in pattern.finditer(input_string)]
+  if len(quote_positions) % 2:
+    raise ValueError(f'unbalanced quotes {quote_char}...{quote_char}')
+
+  output_string = ''
+  placeholders = placeholders if placeholders is not None else dict()
+  prev_end = -1
+  for start, end in start_and_end(quote_positions):
+    output_string += input_string[prev_end+1:start]
+    while True:
+      placeholder = ''.join(random.choices(string.ascii_lowercase, k=5))
+      if placeholder not in input_string and placeholder not in output_string:
+        break
+    output_string += placeholder
+    placeholders[placeholder] = input_string[start+1:end]
+    prev_end = end
+  output_string += input_string[prev_end+1:]
+
+  return output_string, placeholders
+
+
+def parse_cmake_args_from_environ(env_var_name=MUJOCO_CMAKE_ARGS):
+  """Parses CMake arguments from an environment variable."""
+  raw_args = os.environ.get(env_var_name, '').strip()
+  unquoted, placeholders = tokenize_quoted_substr(raw_args, '"')
+  unquoted, placeholders = tokenize_quoted_substr(unquoted, "'", placeholders)
+  parts = re.split(r'\s+', unquoted.strip())
+  out = []
+  for part in parts:
+    for k, v in placeholders.items():
+      part = part.replace(k, v)
+    part = part.replace('\\"', '"').replace("\\'", "'")
+    if part:
+      out.append(part)
+  return out
+
+
+class CMakeExtension(setuptools.Extension):
+  """A Python extension that has been prebuilt by CMake.
+
+  We do not want distutils to handle the build process for our extensions, so
+  so we pass an empty list to the super constructor.
+  """
+
+  def __init__(self, name):
+    super().__init__(name, sources=[])
+
+
+class BuildCMakeExtension(build_ext.build_ext):
+  """Uses CMake to build extensions."""
+
+  def run(self):
+    self._mujoco_library_path, self._mujoco_include_path = self._find_mujoco()
+    self._configure_cmake()
+    for ext in self.extensions:
+      assert ext.name.startswith(EXT_PREFIX)
+      assert '.' not in ext.name[len(EXT_PREFIX):]
+      self.build_extension(ext)
+    self._copy_external_libraries()
+    self._copy_mujoco_headers()
+
+  def _find_mujoco(self):
+    if MUJOCO_PATH not in os.environ:
+      raise RuntimeError(f'{MUJOCO_PATH} environment variable is not set')
+    library_path = None
+    include_path = None
+    for directory, _, filenames in os.walk(os.environ['MUJOCO_PATH']):
+      if fnmatch.filter(filenames, get_mujoco_lib_pattern()):
+        library_path = directory
+      if fnmatch.filter(filenames, 'mujoco.h'):
+        include_path = directory
+      if library_path and include_path:
+        return library_path, include_path
+    raise RuntimeError('Cannot find MuJoCo library and/or include paths')
+
+  def _copy_external_libraries(self):
+    dst = os.path.dirname(self.get_ext_fullpath(self.extensions[0].name))
+    for directory, _, filenames in os.walk(os.environ['MUJOCO_PATH']):
+      for pattern in get_external_lib_patterns():
+        for filename in fnmatch.filter(filenames, pattern):
+          shutil.copyfile(os.path.join(directory, filename),
+                          os.path.join(dst, filename))
+
+  def _copy_mujoco_headers(self):
+    dst = os.path.join(
+        os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
+        'include')
+    os.mkdir(dst)
+    for directory, _, filenames in os.walk(self._mujoco_include_path):
+      for filename in fnmatch.filter(filenames, '*.h'):
+        shutil.copyfile(os.path.join(directory, filename),
+                        os.path.join(dst, filename))
+
+  def _configure_cmake(self):
+    """Check for CMake."""
+    cmake = os.environ.get(MUJOCO_CMAKE, 'cmake')
+    build_cfg = 'Debug' if self.debug else 'Release'
+    cmake_module_path = os.path.join(os.path.dirname(__file__), 'cmake')
+    cmake_args = [
+        f'-DPython3_ROOT_DIR={sys.prefix}',
+        f'-DPython3_EXECUTABLE={sys.executable}',
+        f'-DMUJOCO_LIBRARY_DIR={self._mujoco_library_path}',
+        f'-DMUJOCO_INCLUDE_DIR={self._mujoco_include_path}',
+        f'-DCMAKE_MODULE_PATH={cmake_module_path}',
+        f'-DCMAKE_BUILD_TYPE={build_cfg}',
+        f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={self.build_temp}',
+        f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION={"OFF" if self.debug else "ON"}',
+        '-DCMAKE_Fortran_COMPILER:STRING=',
+        '-DCMAKE_VERBOSE_MAKEFILE=ON',
+        '-DBUILD_TESTING=OFF',
+    ]
+    if platform.system() != 'Windows':
+      cmake_args.extend([
+          f'-DPython3_LIBRARY={sysconfig.get_paths()["stdlib"]}',
+          f'-DPython3_INCLUDE_DIR={sysconfig.get_paths()["include"]}',
+      ])
+    if platform.system() == 'Darwin' and os.environ.get('ARCHFLAGS'):
+      osx_archs = []
+      if '-arch x86_64' in os.environ['ARCHFLAGS']:
+        osx_archs.append('x86_64')
+      if '-arch arm64' in os.environ['ARCHFLAGS']:
+        osx_archs.append('arm64')
+      cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={";".join(osx_archs)}')
+
+    cmake_args.extend(parse_cmake_args_from_environ())
+    os.makedirs(self.build_temp, exist_ok=True)
+
+    if platform.system() == 'Windows':
+      cmake_args = [arg.replace('\\', '/') for arg in cmake_args]
+
+    print('Configuring CMake with the following arguments:')
+    for arg in cmake_args:
+      print(f'    {arg}')
+    subprocess.check_call(
+        [cmake] + cmake_args +
+        [os.path.join(os.path.dirname(__file__), 'mujoco')],
+        cwd=self.build_temp)
+
+    print('Building all extensions with CMake')
+    subprocess.check_call(
+        [cmake, '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
+        cwd=self.build_temp)
+
+  def build_extension(self, ext):
+    dest_path = self.get_ext_fullpath(ext.name)
+    build_path = os.path.join(self.build_temp, os.path.basename(dest_path))
+    subprocess.check_call(['cp', build_path, dest_path])
+
+
+def find_data_files(package_dir, patterns):
+  """Recursively finds files whose names match the given shell patterns."""
+  paths = set()
+  for directory, _, filenames in os.walk(package_dir):
+    for pattern in patterns:
+      for filename in fnmatch.filter(filenames, pattern):
+        # NB: paths must be relative to the package directory.
+        relative_dirpath = os.path.relpath(directory, package_dir)
+        paths.add(os.path.join(relative_dirpath, filename))
+  return list(paths)
+
+
+setup(
+    name='mujoco',
+    version=__version__,
+    author='DeepMind',
+    author_email='mujoco@deepmind.com',
+    description='MuJoCo Physics Simulator',
+    long_description=get_long_description(),
+    long_description_content_type='text/markdown',
+    url='https://github.com/deepmind/mujoco',
+    license='Apache License 2.0',
+    classifiers=[
+        'License :: OSI Approved :: Apache Software License',
+    ],
+    cmdclass=dict(build_ext=BuildCMakeExtension),
+    ext_modules=[
+        CMakeExtension('mujoco._callbacks'),
+        CMakeExtension('mujoco._constants'),
+        CMakeExtension('mujoco._enums'),
+        CMakeExtension('mujoco._errors'),
+        CMakeExtension('mujoco._functions'),
+        CMakeExtension('mujoco._render'),
+        CMakeExtension('mujoco._rollout'),
+        CMakeExtension('mujoco._structs'),
+    ],
+    python_requires='>=3.7',
+    install_requires=[
+        'absl-py',
+        'glfw',
+        'numpy',
+        'pyopengl',
+    ],
+    tests_require=[
+        'absl-py',
+        'glfw',
+        'numpy',
+        'pyopengl',
+    ],
+    test_suite='mujoco',
+    packages=find_packages(),
+    package_data={
+        'mujoco':
+            find_data_files(
+                package_dir='mujoco',
+                patterns=[
+                    'libmujoco*.dylib',
+                    'libmujoco*.so*',
+                    'mujoco*.dll',
+                    'libglew*.so*',
+                    'mujoco.h',
+                    'mj*.h',
+                ]),
+    },
+)

sample/Makefile

@@ -1,12 +1,16 @@
-COMMON=-O2 -I../include -L../bin -std=c++11 -mavx -pthread -Wl,-rpath,'$$ORIGIN'
+# This Makefile assumes that you have GLFW libraries and headers installed on,
+# which is commonly available through your distro's package manager.
+# On Debian and Ubuntu, GLFW can be installed via `apt install libglfw3-dev`.
+
+COMMON=-O2 -I../include -L../lib -std=c++11 -pthread -Wl,-rpath,'$$ORIGIN'/../lib
 
 all:
-	g++ $(COMMON) testxml.cc    -lmujoco210nogl                             -o ../bin/testxml
+	$(CXX) $(COMMON) testxml.cc    -lmujoco_nogl               -o ../bin/testxml
-	g++ $(COMMON) testspeed.cc  -lmujoco210nogl                             -o ../bin/testspeed
+	$(CXX) $(COMMON) testspeed.cc  -lmujoco_nogl               -o ../bin/testspeed
-	g++ $(COMMON) compile.cc    -lmujoco210nogl                             -o ../bin/compile
+	$(CXX) $(COMMON) compile.cc    -lmujoco_nogl               -o ../bin/compile
-	g++ $(COMMON) derivative.cc -lmujoco210nogl -fopenmp                    -o ../bin/derivative
+	$(CXX) $(COMMON) derivative.cc -lmujoco_nogl -fopenmp      -o ../bin/derivative
-	g++ $(COMMON) basic.cc      -lmujoco210 -lGL -lglew ../bin/libglfw.so.3 -o ../bin/basic
+	$(CXX) $(COMMON) basic.cc      -lmujoco -lGL -lglew -lglfw -o ../bin/basic
-	g++ $(COMMON) record.cc     -lmujoco210 -lGL -lglew ../bin/libglfw.so.3 -o ../bin/record
+	$(CXX) $(COMMON) record.cc     -lmujoco -lGL -lglew -lglfw -o ../bin/record
-	gcc -c -O2 -mavx -I../include ../include/uitools.c
+	$(CC) -c -O2 -I../include uitools.c
-	g++ $(COMMON) uitools.o simulate.cc -lmujoco210 -lGL -lglew ../bin/libglfw.so.3 -o ../bin/simulate
+	$(CXX) $(COMMON) uitools.o simulate.cc -lmujoco -lGL -lglew -lglfw -o ../bin/simulate
 	rm *.o

sample/Makefile.macos

@@ -1,12 +1,21 @@
-COMMON=-O2 -I../include -L../bin -std=c++11 -stdlib=libc++ -mavx -pthread
+# This Makefile assumes that GLFW is installed via Homebrew, and that Homebrew
+# packages are installed in /usr/local. If your setup is different, you will
+# need to adjust the HOMEBREW variable accordingly.
+
+HOMEBREW?=/usr/local
+
+MUJOCO_LIBDIR=../MuJoCo.Framework/Versions/A
+MUJOCO_INCLUDEDIR=../MuJoCo.Framework/Versions/A/Headers
+CFLAGS=-O2 -I$(MUJOCO_INCLUDEDIR) -I$(HOMEBREW)/include -pthread
+ALLFLAGS=$(CFLAGS) -L$(MUJOCO_LIBDIR) -L$(HOMEBREW)/lib -std=c++11 -stdlib=libc++
 
 all:
-	clang++ $(COMMON) testxml.cc    -lmujoco210nogl      -o ../bin/testxml
+	clang++ $(ALLFLAGS) testxml.cc    -lmujoco.2.1.2        -o testxml
-	clang++ $(COMMON) testspeed.cc  -lmujoco210nogl      -o ../bin/testspeed
+	clang++ $(ALLFLAGS) testspeed.cc  -lmujoco.2.1.2        -o testspeed
-	clang++ $(COMMON) compile.cc    -lmujoco210nogl      -o ../bin/compile
+	clang++ $(ALLFLAGS) compile.cc    -lmujoco.2.1.2        -o compile
-	clang++ $(COMMON) derivative.cc -lmujoco210nogl      -o ../bin/derivative
+	clang++ $(ALLFLAGS) derivative.cc -lmujoco.2.1.2        -o derivative
-	clang++ $(COMMON) basic.cc      -lmujoco210 -lglfw.3 -o ../bin/basic
+	clang++ $(ALLFLAGS) basic.cc      -lmujoco.2.1.2 -lglfw -o basic
-	clang++ $(COMMON) record.cc     -lmujoco210 -lglfw.3 -o ../bin/record
+	clang++ $(ALLFLAGS) record.cc     -lmujoco.2.1.2 -lglfw -o record
-	clang -c -O2 -mavx -I../include ../include/uitools.c
+	clang -c $(CFLAGS) uitools.c
-	clang++ $(COMMON) uitools.o simulate.cc -lmujoco210 -lglfw.3 -o ../bin/simulate
+	clang++ $(ALLFLAGS) uitools.o simulate.cc -lmujoco.2.1.2 -lglfw -o simulate
 	rm *.o

sample/Makefile.windows

@@ -1,11 +1,19 @@
+# A copy of the GLFW library is required to build some sample programs.
+# If this is not already installed on your system, download the WIN64 archive
+# from https://github.com/glfw/glfw/releases, and copy files as follows:
+# - Copy the entire `include/GLFW` subdirectory to `mujoco/include/GLFW`.
+# - Copy glfw3dll.lib from the subdirectory corresponding into your compiler
+#   choice to `mujoco/lib/glfw3.lib`.
+# - Copy glfw3.dll from the same subdirectory into `mujoco/bin/glfw3.dll`.
+
 COMMON=/O2 /MT /EHsc /arch:AVX /I../include /Fe../bin/
 
 all:
-	cl $(COMMON) testxml.cc                 ../bin/mujoco210nogl.lib
+	cl $(COMMON) testxml.cc                 ../lib/mujoco_nogl.lib
-	cl $(COMMON) testspeed.cc               ../bin/mujoco210nogl.lib
+	cl $(COMMON) testspeed.cc               ../lib/mujoco_nogl.lib
-	cl $(COMMON) compile.cc                 ../bin/mujoco210nogl.lib
+	cl $(COMMON) compile.cc                 ../lib/mujoco_nogl.lib
-	cl $(COMMON) derivative.cc /openmp      ../bin/mujoco210nogl.lib
+	cl $(COMMON) derivative.cc /openmp      ../lib/mujoco_nogl.lib
-	cl $(COMMON) basic.cc ../bin/glfw3.lib  ../bin/mujoco210.lib
+	cl $(COMMON) basic.cc ../lib/glfw3.lib  ../lib/mujoco.lib
-	cl $(COMMON) record.cc ../bin/glfw3.lib ../bin/mujoco210.lib
+	cl $(COMMON) record.cc ../lib/glfw3.lib ../lib/mujoco.lib
-	cl $(COMMON) simulate.cc ../include/uitools.c ../bin/glfw3.lib ../bin/mujoco210.lib
+	cl $(COMMON) simulate.cc uitools.c ../lib/glfw3.lib ../lib/mujoco.lib
 	del *.obj

sample/basic.cc

@@ -16,7 +16,7 @@
 #include <cstring>
 
 #include "GLFW/glfw3.h"
-#include "mujoco.h"
+#include <mujoco.h>
 
 // MuJoCo data structures
 mjModel* m = NULL;                  // MuJoCo model

sample/compile.cc

@@ -17,7 +17,7 @@
 #include <cstdio>
 #include <cstring>
 
-#include "mujoco.h"
+#include <mujoco.h>
 
 // help
 const char helpstring[] =

sample/derivative.cc

@@ -16,7 +16,7 @@
 #include <cstdio>
 #include <cstring>
 
-#include "mujoco.h"
+#include <mujoco.h>
 
 // enable compilation with and without OpenMP support
 #if defined(_OPENMP)
@@ -61,7 +61,7 @@ void worker(const mjModel* m, const mjData* dmain, mjData* d, int id) {
   int nv = m->nv;
 
   // allocate stack space for result at center
-  mjMARKSTACK
+  mjMARKSTACK;
   mjtNum* center = mj_stackAlloc(d, nv);
   mjtNum* warmstart = mj_stackAlloc(d, nv);
 
@@ -188,7 +188,7 @@ void worker(const mjModel* m, const mjData* dmain, mjData* d, int id) {
     }
   }
 
-  mjFREESTACK
+  mjFREESTACK;
 }
 
 
@@ -222,7 +222,7 @@ void checkderiv(const mjModel* m, mjData* d, mjtNum error[7]) {
   int nv = m->nv;
 
   // allocate space
-  mjMARKSTACK
+  mjMARKSTACK;
   mjtNum* mat = mj_stackAlloc(d, nv*nv);
 
   // get pointers to derivative matrices
@@ -275,7 +275,7 @@ void checkderiv(const mjModel* m, mjData* d, mjtNum error[7]) {
   mju_addTo(mat, F0, nv*nv);
   error[7] = relnorm(mat, F0, nv*nv);
 
-  mjFREESTACK
+  mjFREESTACK;
 }
 
 

sample/record.cc

@@ -16,7 +16,7 @@
 #include <cstdlib>
 #include <cstring>
 
-#include "mujoco.h"
+#include <mujoco.h>
 
 // select EGL, OSMESA or GLFW
 #if defined(MJ_EGL)

sample/simulate.cc

@@ -19,7 +19,7 @@
 #include <string>
 #include <thread>
 
-#include "mjxmacro.h"
+#include <mjxmacro.h>
 #include "uitools.h"
 
 #include "array_safety.h"
@@ -34,15 +34,17 @@ const int maxgeom = 5000;           // preallocated geom array in mjvScene
 const double syncmisalign = 0.1;    // maximum time mis-alignment before re-sync
 const double refreshfactor = 0.7;   // fraction of refresh available for simulation
 const int max_slow_down = 128;      // maximum slow-down quotient
+const double zoom_increment = 0.02; // ratio of single click-wheel zoom increment to vertical extent
 
 // model and data
 mjModel* m = NULL;
 mjData* d = NULL;
 
 
-// filename strings
+// strings
 char filename[kBufSize] = "";
 char previous_filename[kBufSize] = "";
+char loadError[kBufSize] = "";
 
 
 // control noise variables
@@ -758,7 +760,7 @@ void makerendering(int oldstate) {
     },
     {
       mjITEM_BUTTON,
-      "Print camera",
+      "Copy camera",
       2,
       NULL,
       ""
@@ -1067,20 +1069,29 @@ void cleartimers(void) {
 
 
 
-// print current camera as MJCF specification
+// copy current camera to clipboard as MJCF specification
 void printcamera(mjvGLCamera* camera) {
+  char clipboard[500];
   mjtNum cam_right[3];
   mjtNum cam_forward[3];
   mjtNum cam_up[3];
+
+  // get camera spec from the GLCamera
   mju_f2n(cam_forward, camera[0].forward, 3);
   mju_f2n(cam_up, camera[0].up, 3);
   mju_cross(cam_right, cam_forward, cam_up);
-  std::printf("<camera pos=\"%.3f %.3f %.3f\" xyaxes=\"%.3f %.3f %.3f %.3f %.3f %.3f\"/>\n",
+
+  // make MJCF camera spec
+  mju::sprintf_arr(clipboard,
+               "<camera pos=\"%.3f %.3f %.3f\" xyaxes=\"%.3f %.3f %.3f %.3f %.3f %.3f\"/>\n",
               (camera[0].pos[0] + camera[1].pos[0]) / 2,
               (camera[0].pos[1] + camera[1].pos[1]) / 2,
               (camera[0].pos[2] + camera[1].pos[2]) / 2,
               cam_right[0], cam_right[1], cam_right[2],
               camera[0].up[0], camera[0].up[1], camera[0].up[2]);
+
+  // copy spec into clipboard
+  glfwSetClipboardString(window, clipboard);
 }
 
 
@@ -1134,28 +1145,34 @@ void loadmodel(void) {
   }
 
   // load and compile
-  char error[500] = "";
+  loadError[0] = '\0';
   mjModel* mnew = 0;
   if (mju::strlen_arr(filename)>4 &&
       !std::strncmp(filename+mju::strlen_arr(filename)-4, ".mjb",
                     mju::sizeof_arr(filename)-mju::strlen_arr(filename)+4)) {
     mnew = mj_loadModel(filename, NULL);
     if (!mnew) {
-      mju::strcpy_arr(error, "could not load binary model");
+      mju::strcpy_arr(loadError, "could not load binary model");
     }
   } else {
-    mnew = mj_loadXML(filename, NULL, error, 500);
+    mnew = mj_loadXML(filename, NULL, loadError, kBufSize);
+    // remove trailing newline character from loadError
+    if (loadError[0]) {
+      int error_length = mju::strlen_arr(loadError);
+      if (loadError[error_length-1] == '\n') {
+        loadError[error_length-1] = '\0';
+      }
+    }
   }
   if (!mnew) {
-    std::printf("%s\n", error);
+    std::printf("%s\n", loadError);
     return;
   }
 
   // compiler warning: print and pause
-  if (error[0]) {
+  if (loadError[0]) {
     // mj_forward() below will print the warning message
-    std::printf("Model compiled, but simulation warning (paused):\n  %s\n\n",
+    std::printf("Model compiled, but simulation warning (paused):\n  %s\n", loadError);
-                error);
     settings.run = 0;
   }
 
@@ -1594,8 +1611,8 @@ void uiEvent(mjuiState* state) {
 
   // 3D scroll
   if (state->type==mjEVENT_SCROLL && state->mouserect==3 && m) {
-    // emulate vertical mouse motion = 5% of window height
+    // emulate vertical mouse motion = 2% of window height
-    mjv_moveCamera(m, mjMOUSE_ZOOM, 0, -0.05*state->sy, &scn, &cam);
+    mjv_moveCamera(m, mjMOUSE_ZOOM, 0, -zoom_increment*state->sy, &scn, &cam);
 
     return;
   }
@@ -1797,12 +1814,19 @@ void render(GLFWwindow* window) {
     mjr_rectangle(rect, 0.2f, 0.3f, 0.4f, 1);
 
     // label
-    if (settings.loadrequest)
+    if (settings.loadrequest) {
-      mjr_overlay(mjFONT_BIG, mjGRID_TOPRIGHT, smallrect,
+      mjr_overlay(mjFONT_BIG, mjGRID_TOPRIGHT, smallrect, "loading", NULL, &con);
-                  "loading", NULL, &con);
+    } else {
-    else
+      char intro_message[kBufSize];
-      mjr_overlay(mjFONT_NORMAL, mjGRID_TOPLEFT, rect,
+      mju::sprintf_arr(intro_message,
-                  "Drag-and-drop model file here", 0, &con);
+                       "MuJoCo version %s\nDrag-and-drop model file here", mj_versionString());
+      mjr_overlay(mjFONT_NORMAL, mjGRID_TOPLEFT, rect, intro_message, 0, &con);
+    }
+
+    // show last loading error
+    if (loadError[0]) {
+      mjr_overlay(mjFONT_NORMAL, mjGRID_BOTTOMLEFT, rect, loadError, 0, &con);
+    }
 
     // render uis
     if (settings.ui0) {
@@ -1821,10 +1845,16 @@ void render(GLFWwindow* window) {
   // render scene
   mjr_render(rect, &scn, &con);
 
+  // show last loading error
+  if (loadError[0]) {
+    mjr_overlay(mjFONT_NORMAL, mjGRID_BOTTOMLEFT, rect, loadError, 0, &con);
+  }
+
   // show pause/loading label
-  if (!settings.run || settings.loadrequest)
+  if (!settings.run || settings.loadrequest) {
     mjr_overlay(mjFONT_BIG, mjGRID_TOPRIGHT, smallrect,
                 settings.loadrequest ? "loading" : "pause", NULL, &con);
+  }
 
   // show realtime label
   if (settings.run && settings.slow_down != 1) {

sample/testspeed.cc

@@ -18,18 +18,21 @@
 #include <string>
 #include <thread>
 
-#include "mujoco.h"
+#include <mujoco.h>
 
 
+// maximum number of threads
+const int maxthread = 512;
+
 // model and per-thread data
 mjModel* m = NULL;
-mjData* d[64];
+mjData* d[maxthread];
 
 
 // per-thread statistics
-int contacts[64];
+int contacts[maxthread];
-int constraints[64];
+int constraints[maxthread];
-double simtime[64];
+double simtime[maxthread];
 
 
 // timer
@@ -121,8 +124,8 @@ int main(int argc, const char** argv) {
   // clamp ctrlnoise to [0.0, 1.0]
   ctrlnoise = mjMAX(0.0, mjMIN(ctrlnoise, 1.0));
 
-  // clamp nthread to [1, 64]
+  // clamp nthread to [1, maxthread]
-  nthread = mjMAX(1, mjMIN(64, nthread));
+  nthread = mjMAX(1, mjMIN(maxthread, nthread));
 
   // get filename, determine file type
   std::string filename(argv[1]);
@@ -169,7 +172,7 @@ int main(int argc, const char** argv) {
   }
 
   // run simulation, record total time
-  std::thread th[64];
+  std::thread th[maxthread];
   double starttime = gettm();
   for (int id=0; id<nthread; id++) {
     th[id] = std::thread(simulate, id, nstep, ctrlnoise);

sample/testxml.cc

@@ -17,8 +17,8 @@
 #include <cstring>
 #include <string>
 
-#include "mjxmacro.h"
+#include <mjxmacro.h>
-#include "mujoco.h"
+#include <mujoco.h>
 
 #include "array_safety.h"
 namespace mju = ::mujoco::sample_util;

sample/uitools.h

@@ -17,7 +17,7 @@
 
 
 #include "GLFW/glfw3.h"
-#include "mujoco.h"
+#include <mujoco.h>
 
 // this is a C-API
 #if defined(__cplusplus)

unity/CHANGELOG.md

@@ -1,4 +1,18 @@
-# Changelog
+# MuJoCo Plug-in for Unity
+
+## [2.1.2] - 2022-03-15
+
+### Minor modifications
+
+When converting a world geom to a free-moving body, the parent body is created
+where the geom is, rather than at the root.
+
+File name changes:
+
+- Added `.Runtime` to `Mujoco.asmdef` and the corresponding test file.
+- Reorganized the files under `Runtime/Componenets` into different
+  subdirectories.
+- File directories for Editor tests are now named with a capital letter.
 
 ## [2.1.1] - 2022-02-02
 

unity/Editor/Components/MjGeomEditor.cs

@@ -42,15 +42,18 @@ public class MjGeomEditor : MjShapeComponentEditor {
       Debug.LogError("This geom already has a Body parent.", geom.GetComponentInParent<MjBody>());
       return;
     }
-    var parent = new GameObject(geom.gameObject.name + " Body").transform;
+    var parent = new GameObject(geom.gameObject.name + " Body").AddComponent<MjBody>().transform;
+    parent.position = geom.transform.position;
+    geom.transform.parent = parent;
+
     var root = geom.transform.root;
     if (root != geom.transform) {
       parent.parent = geom.transform.parent;
     }
-    parent.gameObject.AddComponent<MjBody>();
+
-    var joint = new GameObject("Free Joint").AddComponent<MjFreeJoint>();
+    var joint = new GameObject("Free Joint").AddComponent<MjFreeJoint>().transform;
-    joint.transform.parent = parent;
+    joint.position = geom.transform.position;
-    geom.transform.parent = parent;
+    joint.parent = parent;
   }
 
   [MenuItem("CONTEXT/Collider/Add a matching MuJoCo geom")]

unity/Editor/Mujoco.Editor.asmdef

@@ -2,7 +2,7 @@
     "name": "Mujoco.Editor",
     "rootNamespace": "",
     "references": [
-        "Mujoco"
+        "Mujoco.Runtime"
     ],
     "includePlatforms": [],
     "excludePlatforms": [],

unity/Runtime/Bindings/MujocoBindings.cs

@@ -23,10 +23,6 @@ public static class MujocoLib {
 
 // ----------------------------------Constants----------------------------------
 
-public const bool THIRD_PARTY_MUJOCO_MJDATA_H_ = true;
-public const bool THIRD_PARTY_MUJOCO_MJMODEL_H_ = true;
-public const bool mjUSEDOUBLE = true;
-public const double mjMINVAL = 1e-15;
 public const double mjPI = 3.141592653589793;
 public const double mjMAXVAL = 10000000000.0;
 public const double mjMINMU = 1e-05;
@@ -39,13 +35,14 @@ public const int mjNEQDATA = 7;
 public const int mjNDYN = 10;
 public const int mjNGAIN = 10;
 public const int mjNBIAS = 10;
+public const int mjNFLUID = 12;
 public const int mjNREF = 2;
 public const int mjNIMP = 5;
 public const int mjNSOLVER = 1000;
-public const bool THIRD_PARTY_MUJOCO_MJRENDER_H_ = true;
 public const int mjNAUX = 10;
 public const int mjMAXTEXTURE = 1000;
-public const bool THIRD_PARTY_MUJOCO_MJUI_H_ = true;
+public const bool mjUSEDOUBLE = true;
+public const double mjMINVAL = 1e-15;
 public const int mjMAXUISECT = 10;
 public const int mjMAXUIITEM = 80;
 public const int mjMAXUITEXT = 300;
@@ -80,21 +77,13 @@ public const int mjKEY_F9 = 298;
 public const int mjKEY_F10 = 299;
 public const int mjKEY_F11 = 300;
 public const int mjKEY_F12 = 301;
-public const bool THIRD_PARTY_MUJOCO_MJVISUALIZE_H_ = true;
 public const int mjNGROUP = 6;
+public const int mjMAXLIGHT = 100;
 public const int mjMAXOVERLAY = 500;
 public const int mjMAXLINE = 100;
 public const int mjMAXLINEPNT = 1000;
 public const int mjMAXPLANEGRID = 200;
-public const bool THIRD_PARTY_MUJOCO_MJXMACRO_H_ = true;
+public const int mjVERSION_HEADER = 212;
-public const bool THIRD_PARTY_MUJOCO_MUJOCO_H_ = true;
-public const int mjVERSION_HEADER = 211;
-public const bool THIRD_PARTY_MUJOCO_MUJOCO_EXPORT_H_ = true;
-public const bool MUJOCO_HELPER_DLL_LOCAL = true;
-public const bool MUJOCO_HELPER_DLL_IMPORT = true;
-public const bool MUJOCO_HELPER_DLL_EXPORT = true;
-public const bool MJAPI = true;
-public const bool MJLOCAL = true;
 
 
 // ------------------------------------Enums------------------------------------
@@ -1950,6 +1939,7 @@ public unsafe struct mjModel_ {
   public double* geom_friction;
   public double* geom_margin;
   public double* geom_gap;
+  public double* geom_fluid;
   public double* geom_user;
   public float* geom_rgba;
   public int* site_type;
@@ -2103,6 +2093,8 @@ public unsafe struct mjModel_ {
   public int* sensor_needstage;
   public int* sensor_objtype;
   public int* sensor_objid;
+  public int* sensor_reftype;
+  public int* sensor_refid;
   public int* sensor_dim;
   public int* sensor_adr;
   public double* sensor_cutoff;
@@ -2497,6 +2489,98 @@ public unsafe struct mjvScene_ {
   public mjvLight_ lights5;
   public mjvLight_ lights6;
   public mjvLight_ lights7;
+  public mjvLight_ lights8;
+  public mjvLight_ lights9;
+  public mjvLight_ lights10;
+  public mjvLight_ lights11;
+  public mjvLight_ lights12;
+  public mjvLight_ lights13;
+  public mjvLight_ lights14;
+  public mjvLight_ lights15;
+  public mjvLight_ lights16;
+  public mjvLight_ lights17;
+  public mjvLight_ lights18;
+  public mjvLight_ lights19;
+  public mjvLight_ lights20;
+  public mjvLight_ lights21;
+  public mjvLight_ lights22;
+  public mjvLight_ lights23;
+  public mjvLight_ lights24;
+  public mjvLight_ lights25;
+  public mjvLight_ lights26;
+  public mjvLight_ lights27;
+  public mjvLight_ lights28;
+  public mjvLight_ lights29;
+  public mjvLight_ lights30;
+  public mjvLight_ lights31;
+  public mjvLight_ lights32;
+  public mjvLight_ lights33;
+  public mjvLight_ lights34;
+  public mjvLight_ lights35;
+  public mjvLight_ lights36;
+  public mjvLight_ lights37;
+  public mjvLight_ lights38;
+  public mjvLight_ lights39;
+  public mjvLight_ lights40;
+  public mjvLight_ lights41;
+  public mjvLight_ lights42;
+  public mjvLight_ lights43;
+  public mjvLight_ lights44;
+  public mjvLight_ lights45;
+  public mjvLight_ lights46;
+  public mjvLight_ lights47;
+  public mjvLight_ lights48;
+  public mjvLight_ lights49;
+  public mjvLight_ lights50;
+  public mjvLight_ lights51;
+  public mjvLight_ lights52;
+  public mjvLight_ lights53;
+  public mjvLight_ lights54;
+  public mjvLight_ lights55;
+  public mjvLight_ lights56;
+  public mjvLight_ lights57;
+  public mjvLight_ lights58;
+  public mjvLight_ lights59;
+  public mjvLight_ lights60;
+  public mjvLight_ lights61;
+  public mjvLight_ lights62;
+  public mjvLight_ lights63;
+  public mjvLight_ lights64;
+  public mjvLight_ lights65;
+  public mjvLight_ lights66;
+  public mjvLight_ lights67;
+  public mjvLight_ lights68;
+  public mjvLight_ lights69;
+  public mjvLight_ lights70;
+  public mjvLight_ lights71;
+  public mjvLight_ lights72;
+  public mjvLight_ lights73;
+  public mjvLight_ lights74;
+  public mjvLight_ lights75;
+  public mjvLight_ lights76;
+  public mjvLight_ lights77;
+  public mjvLight_ lights78;
+  public mjvLight_ lights79;
+  public mjvLight_ lights80;
+  public mjvLight_ lights81;
+  public mjvLight_ lights82;
+  public mjvLight_ lights83;
+  public mjvLight_ lights84;
+  public mjvLight_ lights85;
+  public mjvLight_ lights86;
+  public mjvLight_ lights87;
+  public mjvLight_ lights88;
+  public mjvLight_ lights89;
+  public mjvLight_ lights90;
+  public mjvLight_ lights91;
+  public mjvLight_ lights92;
+  public mjvLight_ lights93;
+  public mjvLight_ lights94;
+  public mjvLight_ lights95;
+  public mjvLight_ lights96;
+  public mjvLight_ lights97;
+  public mjvLight_ lights98;
+  public mjvLight_ lights99;
   public mjvGLCamera_ camera0;
   public mjvGLCamera_ camera1;
   public byte enabletransform;
@@ -2663,11 +2747,14 @@ public static unsafe extern void mj_setConst(mjModel_* m, mjData_* d);
 public static unsafe extern int mj_setLengthRange(mjModel_* m, mjData_* d, int index, mjLROpt_* opt, StringBuilder error, int error_sz);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mj_printFormattedModel(mjModel_* m, [MarshalAs(UnmanagedType.LPStr)]string filename, [MarshalAs(UnmanagedType.LPStr)]string float_format_str);
+public static unsafe extern void mj_printFormattedModel(mjModel_* m, [MarshalAs(UnmanagedType.LPStr)]string filename, [MarshalAs(UnmanagedType.LPStr)]string float_format);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_printModel(mjModel_* m, [MarshalAs(UnmanagedType.LPStr)]string filename);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_printFormattedData(mjModel_* m, mjData_* d, [MarshalAs(UnmanagedType.LPStr)]string filename, [MarshalAs(UnmanagedType.LPStr)]string float_format);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_printData(mjModel_* m, mjData_* d, [MarshalAs(UnmanagedType.LPStr)]string filename);
 

unity/Runtime/Components/Bodies.meta

@@ -0,0 +1,8 @@
+fileFormatVersion: 2
+guid: 7d450c02683794ca6998862eebc97e43
+folderAsset: yes
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

unity/Runtime/Components/Equality.meta

@@ -0,0 +1,8 @@
+fileFormatVersion: 2
+guid: f98312e1fa23c4aa0a44050a7d76fc98
+folderAsset: yes
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: