Merge pull request #50 from thowell/actuator_velocity_tile

compute actuator_velocity with tile operations

Author: erikfrey

mujoco/mjx/_src/forward.py

@@ -286,11 +286,11 @@ def qderiv_actuator_fused_kernel(
         m: Model, d: Data, damping: wp.array(dtype=wp.float32), leveladr: int
       ):
         worldid, nodeid = wp.tid()
-        offset_nv = m.qderiv_implicit_offset_nv[leveladr + nodeid]
+        offset_nv = m.actuator_moment_offset_nv[leveladr + nodeid]
 
         # skip tree with no actuators.
         if wp.static(actuation_enabled and tilesize_nu != 0):
-          offset_nu = m.qderiv_implicit_offset_nu[leveladr + nodeid]
+          offset_nu = m.actuator_moment_offset_nu[leveladr + nodeid]
           actuator_moment_tile = wp.tile_load(
             d.actuator_moment[worldid],
             shape=(tilesize_nu, tilesize_nv),
@@ -340,9 +340,9 @@ def qderiv_actuator_fused_kernel(
         block_dim=block_dim,
       )
 
-    qderiv_tilesize_nv = m.qderiv_implicit_tilesize_nv.numpy()
-    qderiv_tilesize_nu = m.qderiv_implicit_tilesize_nu.numpy()
-    qderiv_tileadr = m.qderiv_implicit_tileadr.numpy()
+    qderiv_tilesize_nv = m.actuator_moment_tilesize_nv.numpy()
+    qderiv_tilesize_nu = m.actuator_moment_tilesize_nu.numpy()
+    qderiv_tileadr = m.actuator_moment_tileadr.numpy()
 
     for i in range(len(qderiv_tileadr)):
       beg = qderiv_tileadr[i]
@@ -389,17 +389,75 @@ def fwd_position(m: Model, d: Data):
 def fwd_velocity(m: Model, d: Data):
   """Velocity-dependent computations."""
 
-  # TODO(team): tile operations?
-  d.actuator_velocity.zero_()
+  if m.opt.is_sparse:
+    # TODO(team): sparse version
+    d.actuator_velocity.zero_()
 
-  @wp.kernel
-  def _actuator_velocity(d: Data):
-    worldid, actid, dofid = wp.tid()
-    moment = d.actuator_moment[worldid, actid]
-    qvel = d.qvel[worldid]
-    wp.atomic_add(d.actuator_velocity[worldid], actid, moment[dofid] * qvel[dofid])
+    @wp.kernel
+    def _actuator_velocity(d: Data):
+      worldid, actid, dofid = wp.tid()
+      moment = d.actuator_moment[worldid, actid]
+      qvel = d.qvel[worldid]
+      wp.atomic_add(d.actuator_velocity[worldid], actid, moment[dofid] * qvel[dofid])
+
+    wp.launch(_actuator_velocity, dim=(d.nworld, m.nu, m.nv), inputs=[d])
+  else:
+
+    def actuator_velocity(
+      adr: int,
+      size: int,
+      tilesize_nu: int,
+      tilesize_nv: int,
+    ):
+      @wp.kernel
+      def _actuator_velocity(
+        m: Model, d: Data, leveladr: int, velocity: array3df, qvel: array3df
+      ):
+        worldid, nodeid = wp.tid()
+        offset_nu = m.actuator_moment_offset_nu[leveladr + nodeid]
+        offset_nv = m.actuator_moment_offset_nv[leveladr + nodeid]
+        actuator_moment_tile = wp.tile_load(
+          d.actuator_moment[worldid],
+          shape=(tilesize_nu, tilesize_nv),
+          offset=(offset_nu, offset_nv),
+        )
+        qvel_tile = wp.tile_load(
+          qvel[worldid], shape=(tilesize_nv, 1), offset=(offset_nv, 0)
+        )
+        velocity_tile = wp.tile_matmul(actuator_moment_tile, qvel_tile)
+
+        wp.tile_store(velocity[worldid], velocity_tile, offset=(offset_nu, 0))
+
+      wp.launch_tiled(
+        _actuator_velocity,
+        dim=(d.nworld, size),
+        inputs=[
+          m,
+          d,
+          adr,
+          d.actuator_velocity.reshape(d.actuator_velocity.shape + (1,)),
+          d.qvel.reshape(d.qvel.shape + (1,)),
+        ],
+        block_dim=32,
+      )
+
+    actuator_moment_tilesize_nu = m.actuator_moment_tilesize_nu.numpy()
+    actuator_moment_tilesize_nv = m.actuator_moment_tilesize_nv.numpy()
+    actuator_moment_tileadr = m.actuator_moment_tileadr.numpy()
 
-  wp.launch(_actuator_velocity, dim=(d.nworld, m.nu, m.nv), inputs=[d])
+    for i in range(len(actuator_moment_tileadr)):
+      beg = actuator_moment_tileadr[i]
+      end = (
+        m.actuator_moment_tileadr.shape[0]
+        if i == len(actuator_moment_tileadr) - 1
+        else actuator_moment_tileadr[i + 1]
+      )
+      actuator_velocity(
+        beg,
+        end - beg,
+        int(actuator_moment_tilesize_nu[i]),
+        int(actuator_moment_tilesize_nv[i]),
+      )
 
   smooth.com_vel(m, d)
   passive.passive(m, d)

mujoco/mjx/_src/forward_test.py

@@ -57,7 +57,7 @@ def _load(self, fname: str, is_sparse: bool = True):
 
   def test_fwd_velocity(self):
     """Tests MJX fwd_velocity."""
-    _, mjd, m, d = self._load("humanoid/humanoid.xml")
+    _, mjd, m, d = self._load("humanoid/humanoid.xml", is_sparse=False)
 
     d.actuator_velocity.zero_()
     mjx.fwd_velocity(m, d)

mujoco/mjx/_src/io.py

@@ -139,12 +139,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     qLD_tileadr = np.cumsum(tile_off)[:-1]
     qLD_tilesize = np.array(sorted(tiles.keys()))
 
-  # tiles for implicit integration - needs nu + nv tile size and offset
-  qderiv_implicit_offset_nv = np.empty(shape=(0,), dtype=int)
-  qderiv_implicit_offset_nu = np.empty(shape=(0,), dtype=int)
-  qderiv_implicit_tileadr = np.empty(shape=(0,), dtype=int)
-  qderiv_implicit_tilesize_nv = np.empty(shape=(0,), dtype=int)
-  qderiv_implicit_tilesize_nu = np.empty(shape=(0,), dtype=int)
+  # tiles for actuator_moment - needs nu + nv tile size and offset
+  actuator_moment_offset_nv = np.empty(shape=(0,), dtype=int)
+  actuator_moment_offset_nu = np.empty(shape=(0,), dtype=int)
+  actuator_moment_tileadr = np.empty(shape=(0,), dtype=int)
+  actuator_moment_tilesize_nv = np.empty(shape=(0,), dtype=int)
+  actuator_moment_tilesize_nu = np.empty(shape=(0,), dtype=int)
 
   if not support.is_sparse(mjm):
     # how many actuators for each tree
@@ -166,18 +166,18 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       act_beg += act_num
 
     sorted_keys = sorted(tiles.keys())
-    qderiv_implicit_offset_nv = [
+    actuator_moment_offset_nv = [
       t[0] for key in sorted_keys for t in tiles.get(key, [])
     ]
-    qderiv_implicit_offset_nu = [
+    actuator_moment_offset_nu = [
       t[1] for key in sorted_keys for t in tiles.get(key, [])
     ]
     tile_off = [0] + [len(tiles[sz]) for sz in sorted(tiles.keys())]
-    qderiv_implicit_tileadr = np.cumsum(tile_off)[:-1]  # offset
-    qderiv_implicit_tilesize_nv = np.array(
+    actuator_moment_tileadr = np.cumsum(tile_off)[:-1]  # offset
+    actuator_moment_tilesize_nv = np.array(
       [a[0] for a in sorted_keys]
     )  # for this level
-    qderiv_implicit_tilesize_nu = np.array(
+    actuator_moment_tilesize_nu = np.array(
       [int(a[1]) for a in sorted_keys]
     )  # for this level
 
@@ -193,20 +193,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.qLD_tile = wp.array(qLD_tile, dtype=wp.int32, ndim=1)
   m.qLD_tileadr = wp.array(qLD_tileadr, dtype=wp.int32, ndim=1, device="cpu")
   m.qLD_tilesize = wp.array(qLD_tilesize, dtype=wp.int32, ndim=1, device="cpu")
-  m.qderiv_implicit_offset_nv = wp.array(
-    qderiv_implicit_offset_nv, dtype=wp.int32, ndim=1
+  m.actuator_moment_offset_nv = wp.array(
+    actuator_moment_offset_nv, dtype=wp.int32, ndim=1
   )
-  m.qderiv_implicit_offset_nu = wp.array(
-    qderiv_implicit_offset_nu, dtype=wp.int32, ndim=1
+  m.actuator_moment_offset_nu = wp.array(
+    actuator_moment_offset_nu, dtype=wp.int32, ndim=1
   )
-  m.qderiv_implicit_tileadr = wp.array(
-    qderiv_implicit_tileadr, dtype=wp.int32, ndim=1, device="cpu"
+  m.actuator_moment_tileadr = wp.array(
+    actuator_moment_tileadr, dtype=wp.int32, ndim=1, device="cpu"
   )
-  m.qderiv_implicit_tilesize_nv = wp.array(
-    qderiv_implicit_tilesize_nv, dtype=wp.int32, ndim=1, device="cpu"
+  m.actuator_moment_tilesize_nv = wp.array(
+    actuator_moment_tilesize_nv, dtype=wp.int32, ndim=1, device="cpu"
   )
-  m.qderiv_implicit_tilesize_nu = wp.array(
-    qderiv_implicit_tilesize_nu, dtype=wp.int32, ndim=1, device="cpu"
+  m.actuator_moment_tilesize_nu = wp.array(
+    actuator_moment_tilesize_nu, dtype=wp.int32, ndim=1, device="cpu"
   )
   m.body_dofadr = wp.array(mjm.body_dofadr, dtype=wp.int32, ndim=1)
   m.body_dofnum = wp.array(mjm.body_dofnum, dtype=wp.int32, ndim=1)

mujoco/mjx/_src/types.py

@@ -239,11 +239,11 @@ class Model:
   qpos_spring: wp.array(dtype=wp.float32, ndim=1)
   body_tree: wp.array(dtype=wp.int32, ndim=1)  # warp only
   body_treeadr: wp.array(dtype=wp.int32, ndim=1)  # warp only
-  qderiv_implicit_offset_nv: wp.array(dtype=wp.int32, ndim=1)  # warp only
-  qderiv_implicit_offset_nu: wp.array(dtype=wp.int32, ndim=1)  # warp only
-  qderiv_implicit_tileadr: wp.array(dtype=wp.int32, ndim=1)  # warp only
-  qderiv_implicit_tilesize_nv: wp.array(dtype=wp.int32, ndim=1)  # warp only
-  qderiv_implicit_tilesize_nu: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  actuator_moment_offset_nv: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  actuator_moment_offset_nu: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  actuator_moment_tileadr: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  actuator_moment_tilesize_nv: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  actuator_moment_tilesize_nu: wp.array(dtype=wp.int32, ndim=1)  # warp only
   qM_fullm_i: wp.array(dtype=wp.int32, ndim=1)  # warp only
   qM_fullm_j: wp.array(dtype=wp.int32, ndim=1)  # warp only
   qM_mulm_i: wp.array(dtype=wp.int32, ndim=1)  # warp only