Add support for binding to arrays of mjs element.
Fixes #2402. PiperOrigin-RevId: 726449592 Change-Id: I0cde44889837a71a222b20cf91fa56568b4af689
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
1b4258d274
commit
5b924fec98
@@ -19,7 +19,7 @@ import ctypes.util
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
from typing import IO, Union
|
||||
from typing import Any, IO, Union, Sequence
|
||||
from typing_extensions import TypeAlias
|
||||
import warnings
|
||||
import zipfile
|
||||
@@ -50,6 +50,7 @@ elif _SYSTEM == 'Darwin':
|
||||
'native, arm64 build of Python.')
|
||||
|
||||
from mujoco import _specs
|
||||
from mujoco import _structs
|
||||
from mujoco._callbacks import *
|
||||
from mujoco._constants import *
|
||||
from mujoco._enums import *
|
||||
@@ -88,6 +89,7 @@ MjStruct: TypeAlias = Union[
|
||||
_specs.MjsPlugin,
|
||||
]
|
||||
|
||||
|
||||
def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
|
||||
"""Converts a spec to a zip file.
|
||||
|
||||
@@ -106,7 +108,68 @@ def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
|
||||
zip_info = zipfile.ZipInfo(os.path.join(spec.modelname, filename))
|
||||
zip_file.writestr(zip_info, contents)
|
||||
|
||||
|
||||
class _MjBindModel:
|
||||
def __init__(self, elements: Sequence[Any]):
|
||||
self.elements = elements
|
||||
|
||||
def __getattr__(self, key: str):
|
||||
items = []
|
||||
for e in self.elements:
|
||||
items.extend(getattr(e, key))
|
||||
return items
|
||||
|
||||
|
||||
class _MjBindData:
|
||||
def __init__(self, elements: Sequence[Any]):
|
||||
self.elements = elements
|
||||
|
||||
def __getattr__(self, key: str):
|
||||
items = []
|
||||
for e in self.elements:
|
||||
items.extend(getattr(e, key))
|
||||
return items
|
||||
|
||||
|
||||
def _bind_model(
|
||||
model: _structs.MjModel, specs: Union[Sequence[MjStruct], MjStruct]
|
||||
):
|
||||
"""Bind a Mujoco spec to a mjModel.
|
||||
|
||||
Args:
|
||||
model: The mjModel to bind to.
|
||||
specs: The mjSpec elements to use for binding, can be a single element or a
|
||||
sequence.
|
||||
Returns:
|
||||
A MjModelGroupedViews object or a list of the same type.
|
||||
"""
|
||||
if isinstance(specs, Sequence):
|
||||
return _MjBindModel([model.bind_scalar(s) for s in specs])
|
||||
else:
|
||||
return model.bind_scalar(specs)
|
||||
|
||||
|
||||
def _bind_data(
|
||||
data: _structs.MjData, specs: Union[Sequence[MjStruct], MjStruct]
|
||||
):
|
||||
"""Bind a Mujoco spec to a mjData.
|
||||
|
||||
Args:
|
||||
data: The mjData to bind to.
|
||||
specs: The mjSpec elements to use for binding, can be a single element or a
|
||||
sequence.
|
||||
Returns:
|
||||
A MjDataGroupedViews object or a list of the same type.
|
||||
"""
|
||||
if isinstance(specs, Sequence):
|
||||
return _MjBindData([data.bind_scalar(s) for s in specs])
|
||||
else:
|
||||
return data.bind_scalar(specs)
|
||||
|
||||
|
||||
_specs.MjSpec.to_zip = to_zip
|
||||
_structs.MjData.bind = _bind_data
|
||||
_structs.MjModel.bind = _bind_model
|
||||
|
||||
HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
|
||||
PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')
|
||||
|
||||
@@ -1100,5 +1100,40 @@ class SpecsTest(absltest.TestCase):
|
||||
with self.assertRaisesRegex(ValueError, 'Frame not found.'):
|
||||
parent.attach(child4, frame='invalid_frame', prefix='child3-')
|
||||
|
||||
def test_bind(self):
|
||||
spec = mujoco.MjSpec.from_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="main">
|
||||
<geom name="main" size="0.15 0.15 0.15" mass="1" type="box"/>
|
||||
<freejoint/>
|
||||
<body name="box">
|
||||
<joint name="box" type="hinge" range="-1 +1"/>
|
||||
<geom name="box" size="0.15 0.15 0.15" mass="1" type="box"/>
|
||||
</body>
|
||||
<body name="sphere">
|
||||
<joint name="sphere" type="hinge" range="-1 +1"/>
|
||||
<geom name="sphere" size="0.15 0.15 0.15" mass="1" type="box"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
""")
|
||||
joint_box = spec.joint('box')
|
||||
joint_sphere = spec.joint('sphere')
|
||||
joints = [joint_box, joint_sphere]
|
||||
mj_model = spec.compile()
|
||||
mj_data = mujoco.MjData(mj_model)
|
||||
np.testing.assert_array_equal(mj_data.bind(joint_box).qpos, 0)
|
||||
np.testing.assert_array_equal(mj_model.bind(joint_box).qposadr, 7)
|
||||
np.testing.assert_array_equal(mj_data.bind(joints).qpos, [0, 0])
|
||||
np.testing.assert_array_equal(mj_model.bind(joints).qposadr, [7, 8])
|
||||
np.testing.assert_array_equal(mj_data.bind([]).qpos, [])
|
||||
np.testing.assert_array_equal(mj_model.bind([]).qposadr, [])
|
||||
with self.assertRaisesRegex(
|
||||
AttributeError, "object has no attribute 'invalid'"
|
||||
):
|
||||
print(mj_model.bind(joints).invalid)
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
|
||||
@@ -1720,7 +1720,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
|
||||
#define XGROUP(spectype, field) \
|
||||
mjModel.def( \
|
||||
"bind", \
|
||||
"bind_scalar", \
|
||||
[](MjModelWrapper& m, spectype& spec) -> auto& { \
|
||||
return m.indexer().field##_by_name(mjs_getString(spec.name)); \
|
||||
}, \
|
||||
@@ -2066,7 +2066,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
|
||||
#define XGROUP(spectype, field) \
|
||||
mjData.def( \
|
||||
"bind", \
|
||||
"bind_scalar", \
|
||||
[](MjDataWrapper& d, spectype& spec) -> auto& { \
|
||||
return d.indexer().field##_by_name(mjs_getString(spec.name)); \
|
||||
}, \
|
||||
|
||||
Reference in New Issue
Block a user