add energy_pos and energy_vel

mujoco_warp/__init__.py

@@ -34,6 +34,8 @@
 from ._src.io import put_data as put_data
 from ._src.io import put_model as put_model
 from ._src.passive import passive as passive
+from ._src.sensor import energy_pos as energy_pos
+from ._src.sensor import energy_vel as energy_vel
 from ._src.sensor import sensor_acc as sensor_acc
 from ._src.sensor import sensor_pos as sensor_pos
 from ._src.sensor import sensor_vel as sensor_vel

mujoco_warp/_src/io.py

@@ -469,6 +469,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     tendon_solimp_lim=wp.array(mjm.tendon_solimp_lim, dtype=types.vec5),
     tendon_range=wp.array(mjm.tendon_range, dtype=wp.vec2f),
     tendon_margin=wp.array(mjm.tendon_margin, dtype=float),
+    tendon_stiffness=wp.array(mjm.tendon_stiffness, dtype=float),
+    tendon_lengthspring=wp.array(mjm.tendon_lengthspring, dtype=wp.vec2),
     tendon_length0=wp.array(mjm.tendon_length0, dtype=float),
     tendon_invweight0=wp.array(mjm.tendon_invweight0, dtype=float),
     wrap_objid=wp.array(mjm.wrap_objid, dtype=int),
@@ -550,6 +552,7 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in
     nl=wp.zeros(1, dtype=int),
     nefc=wp.zeros(1, dtype=int),
     time=wp.zeros(nworld, dtype=float),
+    energy=wp.zeros(nworld, dtype=wp.vec2),
     qpos=wp.zeros((nworld, mjm.nq), dtype=float),
     qvel=wp.zeros((nworld, mjm.nv), dtype=float),
     act=wp.zeros((nworld, mjm.na), dtype=float),
@@ -836,6 +839,7 @@ def padtile(x, length, dtype=None):
     nl=arr([mjd.nl * nworld]),
     nefc=arr([mjd.nefc * nworld]),
     time=arr(mjd.time * np.ones(nworld)),
+    energy=tile(mjd.energy, dtype=wp.vec2),
     qpos=tile(mjd.qpos),
     qvel=tile(mjd.qvel),
     act=tile(mjd.act),
@@ -1017,6 +1021,7 @@ def get_data_into(
     mujoco._functions._realloc_con_efc(result, ncon=ncon, nefc=nefc)
 
   result.time = d.time.numpy()[0]
+  result.energy = d.energy.numpy()[0]
   result.ne = d.ne.numpy()[0]
   result.qpos[:] = d.qpos.numpy()[0]
   result.qvel[:] = d.qvel.numpy()[0]

mujoco_warp/_src/sensor.py

@@ -15,19 +15,21 @@
 
 from typing import Any, Tuple
 
-import numpy as np
 import warp as wp
 
 from . import math
 from . import smooth
+from . import support
 from .types import MJ_MINVAL
 from .types import Data
 from .types import DataType
 from .types import DisableBit
+from .types import JointType
 from .types import Model
 from .types import ObjType
 from .types import SensorType
 from .warp_util import event_scope
+from .warp_util import kernel as nested_kernel
 
 
 @wp.func
@@ -1254,3 +1256,227 @@ def sensor_acc(m: Model, d: Data):
     ],
     outputs=[d.sensordata],
   )
+
+
+@wp.kernel
+def _energy_pos_gravity(
+  # Model:
+  opt_gravity: wp.vec3,
+  body_mass: wp.array(dtype=float),
+  # Data in:
+  xipos_in: wp.array2d(dtype=wp.vec3),
+  # Data out:
+  energy_out: wp.array(dtype=wp.vec2),
+):
+  worldid, bodyid = wp.tid()
+  bodyid += 1  # skip world body
+
+  energy = wp.vec2(
+    body_mass[bodyid] * wp.dot(opt_gravity, xipos_in[worldid, bodyid]),
+    0.0,
+  )
+
+  wp.atomic_sub(energy_out, worldid, energy)
+
+
+@wp.kernel
+def _energy_pos_passive_joint(
+  # Model:
+  qpos_spring: wp.array(dtype=float),
+  jnt_type: wp.array(dtype=int),
+  jnt_qposadr: wp.array(dtype=int),
+  jnt_stiffness: wp.array(dtype=float),
+  # Data in:
+  qpos_in: wp.array2d(dtype=float),
+  # Data out:
+  energy_out: wp.array(dtype=wp.vec2),
+):
+  worldid, jntid = wp.tid()
+  stiffness = jnt_stiffness[jntid]
+
+  if stiffness == 0.0:
+    return
+
+  padr = jnt_qposadr[jntid]
+  jnttype = jnt_type[jntid]
+
+  if jnttype == int(JointType.FREE.value):
+    dif0 = wp.vec3(
+      qpos_in[worldid, padr + 0] - qpos_spring[padr + 0],
+      qpos_in[worldid, padr + 1] - qpos_spring[padr + 1],
+      qpos_in[worldid, padr + 2] - qpos_spring[padr + 2],
+    )
+
+    # convert quaternion difference into angular "velocity"
+    quat1 = wp.quat(
+      qpos_in[worldid, padr + 3],
+      qpos_in[worldid, padr + 4],
+      qpos_in[worldid, padr + 5],
+      qpos_in[worldid, padr + 6],
+    )
+    quat1 = wp.normalize(quat1)
+
+    quat_spring = wp.quat(
+      qpos_spring[padr + 3],
+      qpos_spring[padr + 4],
+      qpos_spring[padr + 5],
+      qpos_spring[padr + 6],
+    )
+
+    dif1 = math.quat_sub(quat1, quat_spring)
+
+    energy = wp.vec2(
+      0.5 * stiffness * (wp.dot(dif0, dif0) + wp.dot(dif1, dif1)),
+      0.0,
+    )
+
+    wp.atomic_add(energy_out, worldid, energy)
+
+  elif jnttype == int(JointType.BALL.value):
+    quat = wp.quat(
+      qpos_in[worldid, padr + 0],
+      qpos_in[worldid, padr + 1],
+      qpos_in[worldid, padr + 2],
+      qpos_in[worldid, padr + 3],
+    )
+    quat = wp.normalize(quat)
+
+    quat_spring = wp.quat(
+      qpos_spring[padr + 0],
+      qpos_spring[padr + 1],
+      qpos_spring[padr + 2],
+      qpos_spring[padr + 3],
+    )
+
+    dif = math.quat_sub(quat, quat_spring)
+    energy = wp.vec2(
+      0.5 * stiffness * wp.dot(dif, dif),
+      0.0,
+    )
+    wp.atomic_add(energy_out, worldid, energy)
+  elif jnttype == int(JointType.SLIDE.value) or jnttype == int(JointType.HINGE.value):
+    dif_ = qpos_in[worldid, padr] - qpos_spring[padr]
+    energy = wp.vec2(
+      0.5 * stiffness * dif_ * dif_,
+      0.0,
+    )
+    wp.atomic_add(energy_out, worldid, energy)
+
+
+@wp.kernel
+def _energy_pos_passive_tendon(
+  # Model:
+  tendon_stiffness: wp.array(dtype=float),
+  tendon_lengthspring: wp.array(dtype=wp.vec2),
+  # Data in:
+  ten_length_in: wp.array2d(dtype=float),
+  # Data out:
+  energy_out: wp.array(dtype=wp.vec2),
+):
+  worldid, tenid = wp.tid()
+
+  stiffness = tendon_stiffness[tenid]
+
+  if stiffness == 0.0:
+    return
+
+  length = ten_length_in[worldid, tenid]
+
+  # compute spring displacement
+  lengthspring = tendon_lengthspring[tenid]
+  lower = lengthspring[0]
+  upper = lengthspring[1]
+
+  if length > upper:
+    displacement = upper - length
+  elif length < lower:
+    displacement = lower - length
+  else:
+    displacement = 0.0
+
+  energy = wp.vec2(0.5 * stiffness * displacement * displacement, 0.0)
+  wp.atomic_add(energy_out, worldid, energy)
+
+
+def energy_pos(m: Model, d: Data):
+  """Position-dependent energy (potential)."""
+
+  # init potential energy: -sum_i(body_i.mass * dot(gravity, body_i.pos))
+  if not m.opt.disableflags & DisableBit.GRAVITY:
+    wp.launch(
+      _energy_pos_gravity, dim=(d.nworld, m.nbody - 1), inputs=[m.opt.gravity, m.body_mass, d.xipos], outputs=[d.energy]
+    )
+
+  if not m.opt.disableflags & DisableBit.PASSIVE:
+    # add joint-level springs
+    wp.launch(
+      _energy_pos_passive_joint,
+      dim=(d.nworld, m.njnt),
+      inputs=[
+        m.qpos_spring,
+        m.jnt_type,
+        m.jnt_qposadr,
+        m.jnt_stiffness,
+        d.qpos,
+      ],
+      outputs=[d.energy],
+    )
+
+    # add tendon-level springs
+    if m.ntendon:
+      wp.launch(
+        _energy_pos_passive_tendon,
+        dim=(d.nworld, m.ntendon),
+        inputs=[
+          m.tendon_stiffness,
+          m.tendon_lengthspring,
+          d.ten_length,
+        ],
+        outputs=[d.energy],
+      )
+
+    # TODO(team): flex
+
+
+def _energy_vel_kinetic(nv: int):
+  @nested_kernel
+  def energy_vel_kinetic(
+    # Data in:
+    qvel_in: wp.array2d(dtype=float),
+    # In:
+    Mqvel: wp.array2d(dtype=float),
+    # Out:
+    energy_out: wp.array(dtype=wp.vec2),
+  ):
+    worldid = wp.tid()
+
+    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
+    Mqvel_tile = wp.tile_load(Mqvel[worldid], shape=wp.static(nv))
+
+    # qvel * (M @ qvel)
+    qvelMqvel_tile = wp.tile_map(wp.mul, qvel_tile, Mqvel_tile)
+
+    # sum(qvel * (M @ qvel))
+    quadratic_tile = wp.tile_reduce(wp.add, qvelMqvel_tile)
+
+    energy_out[worldid][1] = 0.5 * quadratic_tile[0]
+
+  return energy_vel_kinetic
+
+
+def energy_vel(m: Model, d: Data):
+  """Velocity-dependent energy (kinetic)."""
+
+  # kinetic energy: 0.5 * qvel.T @ M @ qvel
+
+  # M @ qvel
+  skip = wp.zeros(d.nworld, dtype=bool)
+  support.mul_m(m, d, d.efc.mv, d.qvel, skip)
+
+  wp.launch_tiled(
+    _energy_vel_kinetic(m.nv),
+    dim=(d.nworld,),
+    inputs=[d.qvel, d.efc.mv],
+    outputs=[d.energy],
+    block_dim=32,
+  )

mujoco_warp/_src/sensor_test.py

@@ -211,6 +211,22 @@ def test_tendon_sensor(self):
 
     _assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
 
+  @parameterized.parameters("humanoid/humanoid.xml", "constraints.xml")
+  def test_energy(self, xml):
+    mjm, mjd, m, d = test_util.fixture(xml, constraint=False, kick=True)
+
+    d.energy.zero_()
+
+    mujoco.mj_energyPos(mjm, mjd)
+    mjwarp.energy_pos(m, d)
+
+    _assert_eq(d.energy.numpy()[0][0], mjd.energy[0], "potential energy")
+
+    mujoco.mj_energyVel(mjm, mjd)
+    mjwarp.energy_vel(m, d)
+
+    _assert_eq(d.energy.numpy()[0][1], mjd.energy[1], "kinetic energy")
+
 
 if __name__ == "__main__":
   wp.init()

mujoco_warp/_src/types.py

@@ -748,6 +748,8 @@ class Model:
     tendon_solimp_lim: constraint solver impedance: limit    (ntendon, mjNIMP)
     tendon_range: tendon length limits                       (ntendon, 2)
     tendon_margin: min distance for limit detection          (ntendon,)
+    tendon_stiffness: stiffness coefficient                  (ntendon,)
+    tendon_lengthspring: spring resting length range         (ntendon, 2)
     tendon_length0: tendon length in qpos0                   (ntendon,)
     tendon_invweight0: inv. weight in qpos0                  (ntendon,)
     wrap_objid: object id: geom, site, joint                 (nwrap,)
@@ -961,6 +963,8 @@ class Model:
   tendon_solimp_lim: wp.array(dtype=vec5)
   tendon_range: wp.array(dtype=wp.vec2)
   tendon_margin: wp.array(dtype=float)
+  tendon_stiffness: wp.array(dtype=float)
+  tendon_lengthspring: wp.array(dtype=wp.vec2)
   tendon_length0: wp.array(dtype=float)
   tendon_invweight0: wp.array(dtype=float)
   wrap_objid: wp.array(dtype=int)
@@ -1042,6 +1046,7 @@ class Data:
     nl: number of limit constraints                             ()
     nefc: number of constraints                                 (1,)
     time: simulation time                                       (nworld,)
+    energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
     qvel: velocity                                              (nworld, nv)
     act: actuator activation                                    (nworld, na)
@@ -1148,6 +1153,7 @@ class Data:
   nl: wp.array(dtype=int)
   nefc: wp.array(dtype=int)
   time: wp.array(dtype=float)
+  energy: wp.array(dtype=wp.vec2)
   qpos: wp.array2d(dtype=float)
   qvel: wp.array2d(dtype=float)
   act: wp.array2d(dtype=float)