Add batch support to MJX data bind.

PiperOrigin-RevId: 729033713
Change-Id: I331317a54a7882dc67bf329da8aeaff7246754e4

Author: quagla

mjx/mujoco/mjx/_src/support.py

@@ -19,6 +19,7 @@
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.introspect import mjxmacro
 from mujoco.mjx._src import math
 from mujoco.mjx._src import scan
 # pylint: disable=g-importing-member
@@ -368,8 +369,17 @@ def __init__(self, model: Model, specs: Sequence[mujoco.MjStruct]):
     else:
       self.id = ids
 
+  def _slice(self, name: str, idx: Union[int, slice, Sequence[int]]):
+    _, expected_dim = mjxmacro.MJMODEL[name]
+    var = getattr(self.model, name)
+    if expected_dim == '1':
+      return var[..., idx]
+    elif expected_dim == '9':
+      return var[..., idx, :, :]
+    return var[..., idx, :]
+
   def __getattr__(self, name: str):
-    return getattr(self.model, self.prefix + name)[self.id, ...]
+    return self._slice(self.prefix + name, self.id)
 
 
 def _bind_model(
@@ -453,6 +463,15 @@ def __getname(self, name: str):
     else:
       return self.prefix + name
 
+  def _slice(self, name: str, idx: Union[int, slice, Sequence[int]]):
+    _, expected_dim = mjxmacro.MJDATA[name]
+    var = getattr(self.data, name)
+    if expected_dim == '1':
+      return var[..., idx]
+    elif expected_dim == '9':
+      return var[..., idx, :, :]
+    return var[..., idx, :]
+
   def __getattr__(self, name: str):
     if name in ('sensordata', 'qpos', 'qvel', 'qacc'):
       adr = num = 0
@@ -471,12 +490,12 @@ def __getattr__(self, name: str):
         idx = []
         for a, n in zip(adr, num):
           idx.extend(a + j for j in range(n))
-        return getattr(self.data, self.__getname(name))[idx, ...]
+        return self._slice(self.__getname(name), idx)
       elif num > 1:
-        return getattr(self.data, self.__getname(name))[adr : adr + num, ...]
+        return self._slice(self.__getname(name), slice(adr, adr + num))
       else:
-        return getattr(self.data, self.__getname(name))[adr, ...]
-    return getattr(self.data, self.__getname(name))[self.id, ...]
+        return self._slice(self.__getname(name), adr)
+    return self._slice(self.__getname(name), self.id)
 
   def set(self, name: str, value: jax.Array) -> Data:
     """Set the value of an array in an MJX Data."""

mjx/mujoco/mjx/_src/support_test.py

@@ -305,11 +305,11 @@ def test_bind(self):
     ):
       print(dx.bind(mx, s.geoms).ctrl)
     with self.assertRaises(
-        AttributeError, msg='ctrl is not available for this type'
+        KeyError, msg='actuator_actuator_ctrl'
     ):
       print(dx.bind(mx, s.actuators).actuator_ctrl)
     with self.assertRaises(
-        AttributeError, msg='ctrl is not available for this type'
+        AttributeError, msg='actuator_actuator_ctrl'
     ):
       print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
     with self.assertRaises(
@@ -323,6 +323,16 @@ def test_bind(self):
       s.geoms[0].name = 'invalid_geom_name'
       print(mx.bind(s.geoms).pos)
 
+    # test batched data
+    batch_size = 16
+    ds = [d for _ in range(batch_size)]
+    vdx = jax.vmap(lambda xpos: dx.replace(xpos=xpos))(
+        jp.array([d.xpos for d in ds], device=jax.devices('cpu')[0]))
+    for i in range(m.nbody):
+      np.testing.assert_array_equal(
+          vdx.bind(mx, s.bodies[i]).xpos, [d.xpos[i, :]] * batch_size
+      )
+
   _CONTACTS = """
     <mujoco>
       <worldbody>

python/mujoco/introspect/mjxmacro.py

@@ -0,0 +1,43 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Generate X macros for Mujoco structs."""
+
+from . import structs
+
+MJMODEL_S = structs.STRUCTS['mjModel']
+MJDATA_S = structs.STRUCTS['mjData']
+
+MJMODEL = dict()
+MJDATA = dict()
+
+for field in MJMODEL_S.fields:
+  if not isinstance(field, structs.StructFieldDecl):
+    continue
+  if field.array_extent is None:
+    continue
+  if len(field.array_extent) == 1:
+    MJMODEL[field.name] = (field.array_extent[0], '1')
+  else:
+    MJMODEL[field.name] = (field.array_extent[0], str(field.array_extent[1]))
+
+for field in MJDATA_S.fields:
+  if not isinstance(field, structs.StructFieldDecl):
+    continue
+  if field.array_extent is None:
+    continue
+  if len(field.array_extent) == 1:
+    MJDATA[field.name] = (field.array_extent[0], '1')
+  else:
+    MJDATA[field.name] = (field.array_extent[0], str(field.array_extent[1]))