Merge pull request #25 from eric-heiden/fwd-actuation

fwd_actuation, (transmission)

Author: erikfrey

.gitignore

@@ -1,6 +1,7 @@
 # Exclude temporary folders
 *.egg-info/
 env/
+.history
 
 # Python byte-compiled / optimized / DLL files
 __pycache__/

mujoco/mjx/__init__.py

@@ -17,6 +17,7 @@
 
 from ._src.forward import euler
 from ._src.forward import forward
+from ._src.forward import fwd_actuation
 from ._src.forward import fwd_acceleration
 from ._src.forward import fwd_position
 from ._src.forward import fwd_velocity
@@ -31,6 +32,7 @@
 from ._src.smooth import kinematics
 from ._src.smooth import rne
 from ._src.smooth import solve_m
+from ._src.smooth import transmission
 from ._src.support import is_sparse
 from ._src.test_util import benchmark
 from ._src.types import *

mujoco/mjx/_src/forward.py

@@ -21,7 +21,7 @@
 from . import passive
 from . import smooth
 
-from .types import array2df
+from .types import array2df, array3df
 from .types import Model
 from .types import Data
 from .types import MJ_MINVAL
@@ -218,7 +218,7 @@ def fwd_position(m: Model, d: Data):
   smooth.factor_m(m, d)
   # TODO(team): collision_driver.collision
   # TODO(team): constraint.make_constraint
-  # TODO(team): smooth.transmission
+  smooth.transmission(m, d)
 
 
 def fwd_velocity(m: Model, d: Data):
@@ -241,6 +241,64 @@ def _actuator_velocity(d: Data):
   smooth.rne(m, d)
 
 
+def fwd_actuation(m: Model, d: Data):
+  """Actuation-dependent computations."""
+  if not m.nu:
+    return
+
+  # TODO support stateful actuators
+
+  @wp.kernel
+  def _force(
+    m: Model,
+    ctrl: array2df,
+    # outputs
+    force: array2df,
+  ):
+    worldid, dofid = wp.tid()
+    gain = m.actuator_gainprm[dofid, 0]
+    bias = m.actuator_biasprm[dofid, 0]
+    # TODO support gain types other than FIXED
+    c = ctrl[worldid, dofid]
+    if m.actuator_ctrllimited[dofid]:
+      r = m.actuator_ctrlrange[dofid]
+      c = wp.clamp(c, r[0], r[1])
+    f = gain * c + bias
+    if m.actuator_forcelimited[dofid]:
+      r = m.actuator_forcerange[dofid]
+      f = wp.clamp(f, r[0], r[1])
+    force[worldid, dofid] = f
+
+  wp.launch(
+    _force, dim=[d.nworld, m.nu], inputs=[m, d.ctrl], outputs=[d.actuator_force]
+  )
+
+  @wp.kernel
+  def _qfrc(m: Model, moment: array3df, force: array2df, qfrc: array2df):
+    worldid, vid = wp.tid()
+
+    s = float(0.0)
+    for uid in range(m.nu):
+      # TODO consider using Tile API or transpose moment for better access pattern
+      s += moment[worldid, uid, vid] * force[worldid, uid]
+    jntid = m.dof_jntid[vid]
+    if m.jnt_actfrclimited[jntid]:
+      r = m.jnt_actfrcrange[jntid]
+      s = wp.clamp(s, r[0], r[1])
+    qfrc[worldid, vid] = s
+
+  wp.launch(
+    _qfrc,
+    dim=(d.nworld, m.nv),
+    inputs=[m, d.actuator_moment, d.actuator_force],
+    outputs=[d.qfrc_actuator],
+  )
+
+  # TODO actuator-level gravity compensation, skip if added as passive force
+
+  return d
+
+
 def fwd_acceleration(m: Model, d: Data):
   """Add up all non-constraint forces, compute qacc_smooth."""
 
@@ -269,7 +327,7 @@ def forward(m: Model, d: Data):
   # TODO(team): sensor.sensor_pos
   fwd_velocity(m, d)
   # TODO(team): sensor.sensor_vel
-  # TODO(team): fwd_actuation
+  fwd_actuation(m, d)
   fwd_acceleration(m, d)
   # TODO(team): sensor.sensor_acc
 

mujoco/mjx/_src/forward_test.py

@@ -20,6 +20,8 @@
 import numpy as np
 import warp as wp
 
+wp.config.verify_cuda = True
+
 import mujoco
 from mujoco import mjx
 
@@ -42,6 +44,8 @@ def _load(self, fname: str, is_sparse: bool = True):
     mjd = mujoco.MjData(mjm)
     mujoco.mj_resetDataKeyframe(mjm, mjd, 1)  # reset to stand_on_left_leg
     mjd.qvel = np.random.uniform(low=-0.01, high=0.01, size=mjd.qvel.shape)
+    mjd.ctrl = np.random.normal(scale=10, size=mjd.ctrl.shape)
+    mjd.act = np.random.normal(scale=10, size=mjd.act.shape)
     mujoco.mj_forward(mjm, mjd)
     m = mjx.put_model(mjm)
     d = mjx.put_data(mjm, mjd)
@@ -59,6 +63,21 @@ def test_fwd_velocity(self):
     )
     _assert_eq(d.qfrc_bias.numpy()[0], mjd.qfrc_bias, "qfrc_bias")
 
+  def test_fwd_actuation(self):
+    """Tests MJX fwd_actuation."""
+    mjm, mjd, m, d = self._load("humanoid/humanoid.xml", is_sparse=False)
+
+    mujoco.mj_fwdActuation(mjm, mjd)
+
+    for arr in (d.actuator_force, d.qfrc_actuator):
+      arr.zero_()
+
+    mjx.fwd_actuation(m, d)
+
+    _assert_eq(d.ctrl.numpy()[0], mjd.ctrl, "ctrl")
+    _assert_eq(d.actuator_force.numpy()[0], mjd.actuator_force, "actuator_force")
+    _assert_eq(d.qfrc_actuator.numpy()[0], mjd.qfrc_actuator, "qfrc_actuator")
+
   def test_fwd_acceleration(self):
     """Tests MJX fwd_acceleration."""
     _, mjd, m, d = self._load("humanoid/humanoid.xml", is_sparse=False)

mujoco/mjx/_src/io.py

@@ -122,6 +122,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.jnt_axis = wp.array(mjm.jnt_axis, dtype=wp.vec3, ndim=1)
   m.jnt_pos = wp.array(mjm.jnt_pos, dtype=wp.vec3, ndim=1)
   m.jnt_stiffness = wp.array(mjm.jnt_stiffness, dtype=wp.float32, ndim=1)
+  m.jnt_actfrclimited = wp.array(mjm.jnt_actfrclimited, dtype=wp.bool, ndim=1)
+  m.jnt_actfrcrange = wp.array(mjm.jnt_actfrcrange, dtype=wp.vec2, ndim=1)
   m.geom_pos = wp.array(mjm.geom_pos, dtype=wp.vec3, ndim=1)
   m.geom_quat = wp.array(mjm.geom_quat, dtype=wp.quat, ndim=1)
   m.site_pos = wp.array(mjm.site_pos, dtype=wp.vec3, ndim=1)
@@ -132,8 +134,17 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.dof_Madr = wp.array(mjm.dof_Madr, dtype=wp.int32, ndim=1)
   m.dof_armature = wp.array(mjm.dof_armature, dtype=wp.float32, ndim=1)
   m.dof_damping = wp.array(mjm.dof_damping, dtype=wp.float32, ndim=1)
-  m.actuator_actlimited = wp.array(mjm.actuator_actlimited, dtype=wp.int32, ndim=1)
-  m.actuator_actrange = wp.array(mjm.actuator_actrange, dtype=wp.vec2f, ndim=1)
+  m.actuator_trntype = wp.array(mjm.actuator_trntype, dtype=wp.int32, ndim=1)
+  m.actuator_trnid = wp.array(mjm.actuator_trnid, dtype=wp.int32, ndim=2)
+  m.actuator_ctrllimited = wp.array(mjm.actuator_ctrllimited, dtype=wp.bool, ndim=1)
+  m.actuator_ctrlrange = wp.array(mjm.actuator_ctrlrange, dtype=wp.vec2, ndim=1)
+  m.actuator_forcelimited = wp.array(mjm.actuator_forcelimited, dtype=wp.bool, ndim=1)
+  m.actuator_forcerange = wp.array(mjm.actuator_forcerange, dtype=wp.vec2, ndim=1)
+  m.actuator_gainprm = wp.array(mjm.actuator_gainprm, dtype=wp.float32, ndim=2)
+  m.actuator_biasprm = wp.array(mjm.actuator_biasprm, dtype=wp.float32, ndim=2)
+  m.actuator_gear = wp.array(mjm.actuator_gear, dtype=wp.spatial_vector, ndim=1)
+  m.actuator_actlimited = wp.array(mjm.actuator_actlimited, dtype=wp.bool, ndim=1)
+  m.actuator_actrange = wp.array(mjm.actuator_actrange, dtype=wp.vec2, ndim=1)
   m.actuator_actadr = wp.array(mjm.actuator_actadr, dtype=wp.int32, ndim=1)
   m.actuator_dyntype = wp.array(mjm.actuator_dyntype, dtype=wp.int32, ndim=1)
   m.actuator_dynprm = wp.array(mjm.actuator_dynprm, dtype=types.vec10f, ndim=1)
@@ -167,6 +178,10 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1) -> types.Data:
   d.site_xmat = wp.zeros((nworld, mjm.nsite), dtype=wp.mat33)
   d.cinert = wp.zeros((nworld, mjm.nbody), dtype=types.vec10)
   d.cdof = wp.zeros((nworld, mjm.nv), dtype=wp.spatial_vector)
+  d.ctrl = wp.zeros((nworld, mjm.nu), dtype=wp.float32)
+  d.actuator_velocity = wp.zeros((nworld, mjm.nu), dtype=wp.float32)
+  d.actuator_force = wp.zeros((nworld, mjm.nu), dtype=wp.float32)
+  d.actuator_length = wp.zeros((nworld, mjm.nu), dtype=wp.float32)
   d.actuator_moment = wp.zeros((nworld, mjm.nu, mjm.nv), dtype=wp.float32)
   d.crb = wp.zeros((nworld, mjm.nbody), dtype=types.vec10)
   if support.is_sparse(mjm):
@@ -178,7 +193,6 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1) -> types.Data:
   d.act_dot = wp.zeros((nworld, mjm.na), dtype=wp.float32)
   d.act = wp.zeros((nworld, mjm.na), dtype=wp.float32)
   d.qLDiagInv = wp.zeros((nworld, mjm.nv), dtype=wp.float32)
-  d.actuator_velocity = wp.zeros((nworld, mjm.nu), dtype=wp.float32)
   d.cvel = wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector)
   d.cdof_dot = wp.zeros((nworld, mjm.nv), dtype=wp.spatial_vector)
   d.qfrc_bias = wp.zeros((nworld, mjm.nv), dtype=wp.float32)
@@ -247,12 +261,15 @@ def tile(x):
   d.site_xmat = wp.array(tile(mjd.site_xmat), dtype=wp.mat33, ndim=2)
   d.cinert = wp.array(tile(mjd.cinert), dtype=types.vec10, ndim=2)
   d.cdof = wp.array(tile(mjd.cdof), dtype=wp.spatial_vector, ndim=2)
-  d.actuator_moment = wp.array(tile(actuator_moment), dtype=wp.float32, ndim=3)
   d.crb = wp.array(tile(mjd.crb), dtype=types.vec10, ndim=2)
   d.qM = wp.array(tile(qM), dtype=wp.float32, ndim=3)
   d.qLD = wp.array(tile(qLD), dtype=wp.float32, ndim=3)
   d.qLDiagInv = wp.array(tile(mjd.qLDiagInv), dtype=wp.float32, ndim=2)
+  d.ctrl = wp.array(tile(mjd.ctrl), dtype=wp.float32, ndim=2)
   d.actuator_velocity = wp.array(tile(mjd.actuator_velocity), dtype=wp.float32, ndim=2)
+  d.actuator_force = wp.array(tile(mjd.actuator_force), dtype=wp.float32, ndim=2)
+  d.actuator_length = wp.array(tile(mjd.actuator_length), dtype=wp.float32, ndim=2)
+  d.actuator_moment = wp.array(tile(actuator_moment), dtype=wp.float32, ndim=3)
   d.cvel = wp.array(tile(mjd.cvel), dtype=wp.spatial_vector, ndim=2)
   d.cdof_dot = wp.array(tile(mjd.cdof_dot), dtype=wp.spatial_vector, ndim=2)
   d.qfrc_bias = wp.array(tile(mjd.qfrc_bias), dtype=wp.float32, ndim=2)

mujoco/mjx/_src/math.py

@@ -62,6 +62,11 @@ def quat_to_mat(quat: wp.quat) -> wp.mat33:
   )
 
 
+@wp.func
+def quat_inv(quat: wp.quat) -> wp.quat:
+  return wp.quat(quat[0], -quat[1], -quat[2], -quat[3])
+
+
 @wp.func
 def inert_vec(i: types.vec10, v: wp.spatial_vector) -> wp.spatial_vector:
   """mju_mulInertVec: multiply 6D vector (rotation, translation) by 6D inertia matrix."""

mujoco/mjx/_src/smooth.py

@@ -14,14 +14,14 @@
 # ==============================================================================
 
 import warp as wp
-
 from . import math
 
 from .types import Model
 from .types import Data
 from .types import array2df
 from .types import array3df
 from .types import vec10
+from .types import JointType, TrnType
 
 
 def kinematics(m: Model, d: Data):
@@ -433,6 +433,75 @@ def qfrc_bias(m: Model, d: Data, cfrc: wp.array(dtype=wp.spatial_vector, ndim=2)
   wp.launch(qfrc_bias, dim=[d.nworld, m.nv], inputs=[m, d, cfrc])
 
 
+def transmission(m: Model, d: Data):
+  """Computes actuator/transmission lengths and moments."""
+  if not m.nu:
+    return d
+
+  @wp.kernel
+  def _transmission(
+    m: Model,
+    d: Data,
+    # outputs
+    length: array2df,
+    moment: array3df,
+  ):
+    worldid, actid = wp.tid()
+    qpos = d.qpos[worldid]
+    jntid = m.actuator_trnid[actid, 0]
+    jnt_typ = m.jnt_type[jntid]
+    qadr = m.jnt_qposadr[jntid]
+    vadr = m.jnt_dofadr[jntid]
+    trntype = m.actuator_trntype[actid]
+    gear = m.actuator_gear[actid]
+    if trntype == wp.static(TrnType.JOINT.value) or trntype == wp.static(
+      TrnType.JOINTINPARENT.value
+    ):
+      if jnt_typ == wp.static(JointType.FREE.value):
+        length[worldid, actid] = 0.0
+        if trntype == wp.static(TrnType.JOINTINPARENT.value):
+          quat_neg = math.quat_inv(
+            wp.quat(qpos[qadr + 3], qpos[qadr + 4], qpos[qadr + 5], qpos[qadr + 6])
+          )
+          gearaxis = math.rot_vec_quat(wp.spatial_bottom(gear), quat_neg)
+          moment[worldid, actid, vadr + 0] = gear[0]
+          moment[worldid, actid, vadr + 1] = gear[1]
+          moment[worldid, actid, vadr + 2] = gear[2]
+          moment[worldid, actid, vadr + 3] = gearaxis[0]
+          moment[worldid, actid, vadr + 4] = gearaxis[1]
+          moment[worldid, actid, vadr + 5] = gearaxis[2]
+        else:
+          for i in range(6):
+            moment[worldid, actid, vadr + i] = gear[i]
+      elif jnt_typ == wp.static(JointType.BALL.value):
+        q = wp.quat(qpos[qadr + 0], qpos[qadr + 1], qpos[qadr + 2], qpos[qadr + 3])
+        axis_angle = math.quat_to_vel(q)
+        gearaxis = wp.spatial_top(gear)  # [:3]
+        if trntype == wp.static(TrnType.JOINTINPARENT.value):
+          quat_neg = math.quat_inv(q)
+          gearaxis = math.rot_vec_quat(gearaxis, quat_neg)
+        length[worldid, actid] = wp.dot(axis_angle, gearaxis)
+        for i in range(3):
+          moment[worldid, actid, vadr + i] = gearaxis[i]
+      elif jnt_typ == wp.static(JointType.SLIDE.value) or jnt_typ == wp.static(
+        JointType.HINGE.value
+      ):
+        length[worldid, actid] = qpos[qadr] * gear[0]
+        moment[worldid, actid, vadr] = gear[0]
+      else:
+        wp.printf("unrecognized joint type")
+    else:
+      # TODO handle site, tendon transmission types
+      wp.printf("unhandled transmission type %d\n", trntype)
+
+  wp.launch(
+    _transmission,
+    dim=[d.nworld, m.nu],
+    inputs=[m, d],
+    outputs=[d.actuator_length, d.actuator_moment],
+  )
+
+
 def com_vel(m: Model, d: Data):
   """Computes cvel, cdof_dot."""
 

mujoco/mjx/_src/smooth_test.py

@@ -22,8 +22,7 @@
 import numpy as np
 import warp as wp
 
-import mujoco
-from mujoco import mjx
+wp.config.verify_cuda = True
 
 from . import test_util
 
@@ -141,6 +140,26 @@ def test_com_vel(self):
     _assert_eq(d.cvel.numpy()[0], mjd.cvel, "cvel")
     _assert_eq(d.cdof_dot.numpy()[0], mjd.cdof_dot, "cdof_dot")
 
+  def test_transmission(self):
+    """Tests transmission."""
+    mjm, mjd, m, d = test_util.fixture("pendula.xml")
+
+    for arr in (d.actuator_length, d.actuator_moment):
+      arr.zero_()
+
+    actuator_moment = np.zeros((mjm.nu, mjm.nv))
+    mujoco.mju_sparse2dense(
+      actuator_moment,
+      mjd.actuator_moment,
+      mjd.moment_rownnz,
+      mjd.moment_rowadr,
+      mjd.moment_colind,
+    )
+
+    mjx._src.smooth.transmission(m, d)
+    _assert_eq(d.actuator_length.numpy()[0], mjd.actuator_length, "actuator_length")
+    _assert_eq(d.actuator_moment.numpy()[0], actuator_moment, "actuator_moment")
+
 
 if __name__ == "__main__":
   wp.init()