Add sensors for potential and kinetic energy.

PiperOrigin-RevId: 716775375
Change-Id: Ic8ab7f1a51df970ab551dabe7c797cde13cd97fd

doc/XMLreference.rst

@@ -601,9 +601,17 @@ from its default.
 .. _option-flag-energy:
 
 :at:`energy`: :at-val:`[disable, enable], "disable"`
-   This flag enables the computation of kinetic and potential energy, stored in mjData.energy and displayed in the GUI.
+   This flag enables the computation of potential and kinetic energy in ``mjData.energy[0, 1]`` respectively,
-   This feature adds some CPU time but it is usually negligible. Monitoring energy for a system that is supposed to be
+   and displayed in the simulate GUI info overlay. Potential energy includes the gravitational component summed over
-   energy-conserving is one of the best ways to assess the accuracy of a complex simulation.
+   all bodies :math:`\sum_b m_b g h` and energy stored in passive springs in joints, tendons and flexes
+   :math:`\tfrac{1}{2} k x^2`, where :math:`x` is the displacement and and :math:`k` is the spring constant. Kinetic
+   energy is given by :math:`\tfrac{1}{2} v^T M v`, where :math:`v` is the velocity and :math:`M` is the
+   mass matrix. Note that potential and kinetic energy in constraints is not accounted for.
+
+   The extra computation (also triggered by :ref:`potential<sensor-e_potential>` and
+   :ref:`kinetic<sensor-e_kinetic>` energy sensors) adds some CPU time but it is usually negligible. Monitoring energy
+   for a system that is supposed to be energy-conserving is one of the best ways to assess the accuracy of a complex
+   simulation.
 
 .. _option-flag-fwdinv:
 
@@ -7245,6 +7253,45 @@ See :ref:`collision-sensors` for more details about sensors of this type.
 :at:`name`, :at:`noise`, :at:`user`
    See :ref:`CSensor`.
 
+
+.. _sensor-e_potential:
+
+:el-prefix:`sensor/` |-| **e_potential** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This element creates sensor that returns the potential energy.
+
+.. _sensor-e_potential-name:
+
+.. _sensor-e_potential-noise:
+
+.. _sensor-e_potential-cutoff:
+
+.. _sensor-e_potential-user:
+
+:at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
+   See :ref:`CSensor`.
+
+
+.. _sensor-e_kinetic:
+
+:el-prefix:`sensor/` |-| **e_kinetic** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This element creates sensor that returns the kinetic energy.
+
+.. _sensor-e_kinetic-name:
+
+.. _sensor-e_kinetic-noise:
+
+.. _sensor-e_kinetic-cutoff:
+
+.. _sensor-e_kinetic-user:
+
+:at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
+   See :ref:`CSensor`.
+
+
 .. _sensor-clock:
 
 :el-prefix:`sensor/` |-| **clock** (*)

doc/XMLschema.rst

@@ -1231,6 +1231,20 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`e_potential                  | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <sensor-e_potential>`              |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<sensor-e_potential-name>`                            | :ref:`cutoff<sensor-e_potential-cutoff>`                        | :ref:`noise<sensor-e_potential-noise>`                          | :ref:`user<sensor-e_potential-user>`                            | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`e_kinetic                    | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <sensor-e_kinetic>`                |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<sensor-e_kinetic-name>`                              | :ref:`cutoff<sensor-e_kinetic-cutoff>`                          | :ref:`noise<sensor-e_kinetic-noise>`                            | :ref:`user<sensor-e_kinetic-user>`                              | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`clock                        | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <sensor-clock>`                    |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |

doc/changelog.rst

@@ -5,10 +5,13 @@ Changelog
 Upcoming version (not yet released)
 -----------------------------------
 
-- Added ``mjs_setDeepCopy`` API function. When the deep copy flag is 0, attaching a model will not copy it to the
+General
-  parent, so the original references to the child allow to modify the parent as well. The default behavior is to perform
+^^^^^^^
-  such a shallow copy. The old behavioud of creating a deep copy of the child model while attaching can be restored by
+- Added :ref:`mjs_setDeepCopy` API function. When the deep copy flag is 0, attaching a model will not copy it to the
-  setting the deep copy flag to 1.
+  parent, so the original references to the child can be used to modify the parent after attachment. The default
+  behavior is to perform such a shallow copy. The old behavior of creating a deep copy of the child model while
+  attaching can be restored by setting the deep copy flag to 1.
+- Added :ref:`potential<sensor-e_potential>` and :ref:`kinetic<sensor-e_kinetic>` energy sensors.
 
 Version 3.2.7 (Jan 14, 2025)
 ----------------------------

doc/includes/references.h

@@ -674,6 +674,8 @@ typedef enum mjtSensor_ {         // type of sensor
   mjSENS_GEOMFROMTO,              // segment between two geoms
 
   // global sensors
+  mjSENS_E_POTENTIAL,             // potential energy
+  mjSENS_E_KINETIC,               // kinetic energy
   mjSENS_CLOCK,                   // simulation time
 
   // plugin-controlled sensors

include/mujoco/mjmodel.h

@@ -347,6 +347,8 @@ typedef enum mjtSensor_ {         // type of sensor
   mjSENS_GEOMFROMTO,              // segment between two geoms
 
   // global sensors
+  mjSENS_E_POTENTIAL,             // potential energy
+  mjSENS_E_KINETIC,               // kinetic energy
   mjSENS_CLOCK,                   // simulation time
 
   // plugin-controlled sensors

introspect/enums.py

@@ -361,9 +361,11 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjSENS_GEOMDIST', 37),
              ('mjSENS_GEOMNORMAL', 38),
              ('mjSENS_GEOMFROMTO', 39),
-             ('mjSENS_CLOCK', 40),
+             ('mjSENS_E_POTENTIAL', 40),
-             ('mjSENS_PLUGIN', 41),
+             ('mjSENS_E_KINETIC', 41),
-             ('mjSENS_USER', 42),
+             ('mjSENS_CLOCK', 42),
+             ('mjSENS_PLUGIN', 43),
+             ('mjSENS_USER', 44),
          ]),
      )),
     ('mjtStage',

src/engine/engine_forward.c

@@ -1013,6 +1013,38 @@ void mj_implicit(const mjModel* m, mjData* d) {
 
 
 
+// return 1 if potential energy was computed by sensor, 0 otherwise
+static int energyPosSensor(const mjModel* m) {
+  if (mjDISABLED(mjDSBL_SENSOR)) {
+    return 0;
+  }
+
+  for (int i=0; i < m->nsensor; i++) {
+    if (m->sensor_type[i] == mjSENS_E_POTENTIAL) {
+      return 1;
+    }
+  }
+  return 0;
+}
+
+
+
+// return 1 if kinetic energy was computed by sensor, 0 otherwise
+static int energyVelSensor(const mjModel* m) {
+  if (mjDISABLED(mjDSBL_SENSOR)) {
+    return 0;
+  }
+
+  for (int i=0; i < m->nsensor; i++) {
+    if (m->sensor_type[i] == mjSENS_E_KINETIC) {
+      return 1;
+    }
+  }
+  return 0;
+}
+
+
+
 //-------------------------- top-level API ---------------------------------------------------------
 
 // forward dynamics with skip; skipstage is mjtStage
@@ -1022,21 +1054,33 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
   // position-dependent
   if (skipstage < mjSTAGE_POS) {
     mj_fwdPosition(m, d);
+
+    int energyPos = 0;
     if (!skipsensor) {
       mj_sensorPos(m, d);
+      energyPos = energyPosSensor(m);
     }
-    if (mjENABLED(mjENBL_ENERGY)) {
+
-      mj_energyPos(m, d);
+    if (!energyPos) {
+      if (mjENABLED(mjENBL_ENERGY)) {
+        mj_energyPos(m, d);
+      } else {
+        d->energy[0] = d->energy[1] = 0;
+      }
     }
   }
 
   // velocity-dependent
   if (skipstage < mjSTAGE_VEL) {
     mj_fwdVelocity(m, d);
+
+    int energyVel = 0;
     if (!skipsensor) {
       mj_sensorVel(m, d);
+      energyVel = energyVelSensor(m);
     }
-    if (mjENABLED(mjENBL_ENERGY)) {
+
+    if (mjENABLED(mjENBL_ENERGY) && !energyVel) {
       mj_energyVel(m, d);
     }
   }
@@ -1111,10 +1155,18 @@ void mj_step1(const mjModel* m, mjData* d) {
   mj_checkVel(m, d);
   mj_fwdPosition(m, d);
   mj_sensorPos(m, d);
-  mj_energyPos(m, d);
+  if (!energyPosSensor(m)) {
+    if (mjENABLED(mjENBL_ENERGY)) {
+      mj_energyPos(m, d);
+    } else {
+      d->energy[0] = d->energy[1] = 0;
+    }
+  }
   mj_fwdVelocity(m, d);
   mj_sensorVel(m, d);
-  mj_energyVel(m, d);
+  if (mjENABLED(mjENBL_ENERGY) && !energyVelSensor(m)) {
+    mj_energyVel(m, d);
+  }
   if (mjcb_control) {
     mjcb_control(m, d);
   }

src/engine/engine_io.c

@@ -2056,6 +2056,8 @@ static int sensorSize(mjtSensor sensor_type, int sensor_dim) {
   case mjSENS_TENDONLIMITVEL:
   case mjSENS_TENDONLIMITFRC:
   case mjSENS_GEOMDIST:
+  case mjSENS_E_POTENTIAL:
+  case mjSENS_E_KINETIC:
   case mjSENS_CLOCK:
     return 1;
 

src/engine/engine_sensor.c

@@ -458,6 +458,16 @@ void mj_sensorPos(const mjModel* m, mjData* d) {
         }
         break;
 
+      case mjSENS_E_POTENTIAL:                            // potential energy
+        mj_energyPos(m, d);
+        d->sensordata[adr] = d->energy[0];
+        break;
+
+      case mjSENS_E_KINETIC:                              // kinetic energy
+        mj_energyVel(m, d);
+        d->sensordata[adr] = d->energy[1];
+        break;
+
       case mjSENS_CLOCK:                                  // clock
         d->sensordata[adr] = d->time;
         break;
@@ -902,12 +912,6 @@ void mj_energyPos(const mjModel* m, mjData* d) {
   int padr;
   mjtNum dif[3], quat[4], stiffness;
 
-  // disabled: clear and return
-  if (!mjENABLED(mjENBL_ENERGY)) {
-    d->energy[0] = d->energy[1] = 0;
-    return;
-  }
-
   // init potential energy:  -sum_i body(i).mass * mju_dot(body(i).pos, gravity)
   d->energy[0] = 0;
   if (!mjDISABLED(mjDSBL_GRAVITY)) {
@@ -996,11 +1000,6 @@ void mj_energyPos(const mjModel* m, mjData* d) {
 
 // velocity-dependent energy (kinetic)
 void mj_energyVel(const mjModel* m, mjData* d) {
-  // return if disabled (already cleared in potential)
-  if (!mjENABLED(mjENBL_ENERGY)) {
-    return;
-  }
-
   mj_markStack(d);
   mjtNum *vec = mjSTACKALLOC(d, m->nv, mjtNum);
 

src/user/user_objects.cc

@@ -6206,7 +6206,11 @@ void mjCSensor::ResolveReferences(const mjCModel* m) {
       ((mjCGeom*)obj)->SetNotVisual();
     }
 
-  } else if (type != mjSENS_CLOCK && type != mjSENS_PLUGIN && type != mjSENS_USER) {
+  } else if (type != mjSENS_E_POTENTIAL &&
+             type != mjSENS_E_KINETIC   &&
+             type != mjSENS_CLOCK       &&
+             type != mjSENS_PLUGIN      &&
+             type != mjSENS_USER) {
     throw mjCError(this, "invalid type in sensor");
   }
 
@@ -6539,6 +6543,8 @@ void mjCSensor::Compile(void) {
     }
     break;
 
+  case mjSENS_E_POTENTIAL:
+  case mjSENS_E_KINETIC:
   case mjSENS_CLOCK:
     dim = 1;
     needstage = mjSTAGE_POS;

src/xml/xml_native_reader.cc

@@ -485,6 +485,8 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
         {"distance", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
         {"normal", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
         {"fromto", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
+        {"e_potential", "*", "4", "name", "cutoff", "noise", "user"},
+        {"e_kinetic", "*", "4", "name", "cutoff", "noise", "user"},
         {"clock", "*", "4", "name", "cutoff", "noise", "user"},
         {"user", "*", "9", "name", "objtype", "objname", "datatype", "needstage",
             "dim", "cutoff", "noise", "user"},
@@ -4172,7 +4174,13 @@ void mjXReader::Sensor(XMLElement* section) {
     }
 
     // global sensors
-    else if (type=="clock") {
+    else if (type=="e_potential") {
+      sensor->type = mjSENS_E_POTENTIAL;
+      sensor->objtype = mjOBJ_UNKNOWN;
+    } else if (type=="e_kinetic") {
+      sensor->type = mjSENS_E_KINETIC;
+      sensor->objtype = mjOBJ_UNKNOWN;
+    } else if (type=="clock") {
       sensor->type = mjSENS_CLOCK;
       sensor->objtype = mjOBJ_UNKNOWN;
     }

src/xml/xml_native_reader.h

@@ -101,7 +101,7 @@ class mjXReader : public mjXBase {
 };
 
 // MJCF schema
-#define nMJCF 237
+#define nMJCF 239
 extern const char* MJCF[nMJCF][mjXATTRNUM];
 
 #endif  // MUJOCO_SRC_XML_XML_NATIVE_READER_H_

src/xml/xml_native_writer.cc

@@ -2156,6 +2156,12 @@ void mjXWriter::Sensor(XMLElement* root) {
       break;
 
     // global sensors
+    case mjSENS_E_POTENTIAL:
+      elem = InsertEnd(section, "potential");
+      break;
+    case mjSENS_E_KINETIC:
+      elem = InsertEnd(section, "kinetic");
+      break;
     case mjSENS_CLOCK:
       elem = InsertEnd(section, "clock");
       break;

test/engine/engine_sensor_test.cc

@@ -21,7 +21,6 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mujoco.h>
-#include "src/engine/engine_support.h"
 #include "src/engine/engine_util_blas.h"
 #include "src/engine/engine_util_spatial.h"
 #include "test/fixture.h"
@@ -398,6 +397,93 @@ TEST_F(RelativeFrameSensorTest, FrameVelGeneral) {
 // ------------------------- general sensor tests  -----------------------------
 using SensorTest = MujocoTest;
 
+TEST_F(SensorTest, EnableEnergy) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <option gravity="0 0 -5">
+     <flag energy="enable"/>
+    </option>
+    <worldbody>
+      <body pos="0 0 2">
+        <geom size="1" mass="3"/>
+        <freejoint/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  mjModel* model = LoadModelFromString(xml);
+  mjData* data = mj_makeData(model);
+
+  mj_forward(model, data);
+  EXPECT_EQ(data->energy[0], 2*3*5);
+
+  model->opt.enableflags &= ~mjENBL_ENERGY;
+  mj_forward(model, data);
+  EXPECT_EQ(data->energy[0], 0);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+TEST_F(SensorTest, PotentialEnergy) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <option gravity="0 0 -5"/>
+    <worldbody>
+      <body pos="0 0 2">
+        <geom size="1" mass="3"/>
+        <freejoint/>
+      </body>
+    </worldbody>
+    <sensor>
+      <e_potential/>
+    </sensor>
+  </mujoco>
+  )";
+  mjModel* model = LoadModelFromString(xml);
+  mjData* data = mj_makeData(model);
+
+  mj_forward(model, data);
+  EXPECT_EQ(data->sensordata[0], 2*3*5);
+
+  data->qpos[2] = 7;
+  mj_forward(model, data);
+  EXPECT_EQ(data->sensordata[0], 7*3*5);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+TEST_F(SensorTest, KineticEnergy) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body pos="0 0 2">
+        <geom size="1" mass="3"/>
+        <freejoint/>
+      </body>
+    </worldbody>
+    <sensor>
+      <e_kinetic/>
+    </sensor>
+  </mujoco>
+  )";
+  mjModel* model = LoadModelFromString(xml);
+  mjData* data = mj_makeData(model);
+
+  while (data->time < 1.5) {
+    mj_step(model, data);
+  }
+  mj_forward(model, data);
+
+  mjtNum mass = 3;
+  mjtNum speed = data->time * mju_norm3(model->opt.gravity);
+  EXPECT_FLOAT_EQ(data->sensordata[0], 0.5 * mass * speed * speed);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 // test clock sensor
 TEST_F(SensorTest, Clock) {
   constexpr char xml[] = R"(

unity/Runtime/Bindings/MjBindings.cs

@@ -377,9 +377,11 @@ public enum mjtSensor : int{
   mjSENS_GEOMDIST = 37,
   mjSENS_GEOMNORMAL = 38,
   mjSENS_GEOMFROMTO = 39,
-  mjSENS_CLOCK = 40,
+  mjSENS_E_POTENTIAL = 40,
-  mjSENS_PLUGIN = 41,
+  mjSENS_E_KINETIC = 41,
-  mjSENS_USER = 42,
+  mjSENS_CLOCK = 42,
+  mjSENS_PLUGIN = 43,
+  mjSENS_USER = 44,
 }
 public enum mjtStage : int{
   mjSTAGE_NONE = 0,