Use mocap pos/quat in smooth.kinematics.

PiperOrigin-RevId: 681136974 Change-Id: I0ece009ed1e68b61b0b96da2e685e86000f7f6b8
2024-10-01 13:10:38 -07:00
parent befe4af8f0
commit 1a9d307042
5 changed files with 25 additions and 2 deletions
@@ -14,6 +14,11 @@ General
  that previously required these plugins.
 - Replaced the function ``mjs_setActivePlugins`` with :ref:`mjs_activatePlugin`.

+MJX
+^^^
+
+- Added ``mocap_pos`` and ``mocap_quat`` in kinematics.
+
 Bug fixes
 ^^^^^^^^^
 - Fixed a bug where ``actuator_force`` was not set in MJX (:github:issue:`2068`).
@@ -98,6 +98,14 @@ def kinematics(m: Model, d: Data) -> Data:
      m.body_quat,
  )

+  if m.nmocap:
+    xpos = xpos.at[m.body_mocapid >= 0].set(d.mocap_pos)
+    mocap_quat = jax.vmap(math.normalize)(d.mocap_quat)
+    xquat = xquat.at[m.body_mocapid >= 0].set(mocap_quat)
+    xmat = xmat.at[m.body_mocapid >= 0].set(
+        jax.vmap(math.quat_to_mat)(mocap_quat)
+    )
+
  v_local_to_global = jax.vmap(support.local_to_global)

  # TODO(erikfrey): confirm that quats are more performant for mjx than mats
@@ -55,6 +55,9 @@ class SmoothTest(absltest.TestCase):
    # give the system a little kick to ensure we have non-identity rotations
    d.qvel = np.random.random(m.nv)
    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    # randomize mocap
+    d.mocap_pos = np.random.random(d.mocap_pos.shape)
+    d.mocap_quat = np.random.random(d.mocap_quat.shape)
    mujoco.mj_forward(m, d)
    mx = mjx.put_model(m)

@@ -1265,8 +1265,8 @@ class Data(PyTreeNode):
  xfrc_applied: jax.Array
  eq_active: jax.Array
  # mocap data:
-  mocap_pos: jax.Array = _restricted_to('mujoco')
-  mocap_quat: jax.Array = _restricted_to('mujoco')
+  mocap_pos: jax.Array
+  mocap_quat: jax.Array
  # dynamics:
  qacc: jax.Array
  act_dot: jax.Array
@@ -124,6 +124,13 @@
        </body>
      </body>
    </body>
+
+    <body mocap="true">
+      <geom type="sphere" size="0.05"/>
+    </body>
+    <body mocap="true" pos="1.0 2.0 0.5" quat="1 0 1 0">
+      <geom type="sphere" size="0.05"/>
+    </body>
  </worldbody>

  <tendon>