Add sensors for potential and kinetic energy.

PiperOrigin-RevId: 716775375
Change-Id: Ic8ab7f1a51df970ab551dabe7c797cde13cd97fd

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);