diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index b43db087..3bc2bc31 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -19,6 +19,7 @@ from typing import Optional, Tuple, Union 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 @@ class BindModel(object): 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 @@ class BindData(object): 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 @@ class BindData(object): 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.""" diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 5fb1dcd1..5ae5af9f 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -305,11 +305,11 @@ class SupportTest(parameterized.TestCase): ): 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 @@ class SupportTest(parameterized.TestCase): 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 = """ diff --git a/python/mujoco/introspect/mjxmacro.py b/python/mujoco/introspect/mjxmacro.py new file mode 100644 index 00000000..9222c2c1 --- /dev/null +++ b/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]))