chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for querying of MjsTendon path from Python.

Browse Source 
Fixes #2670

PiperOrigin-RevId: 822215254
Change-Id: Ice5a264e54f963776093511c6e4ac830aae6d4f8
This commit is contained in:
Sam Haves
2025-10-21 11:50:24 -07:00
committed by Copybara-Service
parent 2f65e23779
commit ac2cd5dfd6
9 changed files with 411 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+54
View File
@@ -4388,6 +4388,42 @@ Return spec's first element of selected type.
Return spec's next element; return NULL if element is last.
.. _mjs_getWrapTarget:
`mjs_getWrapTarget <#mjs_getWrapTarget>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrapTarget
Get wrapped element in tendon path.
.. _mjs_getWrapSideSite:
`mjs_getWrapSideSite <#mjs_getWrapSideSite>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrapSideSite
Get wrapped element side site in tendon path if it has one, nullptr otherwise.
.. _mjs_getWrapDivisor:
`mjs_getWrapDivisor <#mjs_getWrapDivisor>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrapDivisor
Get divisor of mjsWrap wrapping a puller.
.. _mjs_getWrapCoef:
`mjs_getWrapCoef <#mjs_getWrapCoef>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrapCoef
Get coefficient of mjsWrap wrapping a joint.
.. _AttributeSetters:
Attribute setters
@@ -4533,6 +4569,24 @@ Get double array contents and optionally its size.
*Nullable:* ``size``
.. _mjs_getWrapNum:
`mjs_getWrapNum <#mjs_getWrapNum>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrapNum
Get number of elements a tendon wraps.
.. _mjs_getWrap:
`mjs_getWrap <#mjs_getWrap>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getWrap
Get mjsWrap element at position i in the tendon path.
.. _mjs_getPluginAttributes:
`mjs_getPluginAttributes <#mjs_getPluginAttributes>`__
doc/changelog.rst
+9
View File
@@ -12,6 +12,15 @@ General
- Increase Windows stack size to 16MB to enable models with deep nested body hierarchies.
- Added a new :ref:`mj_extractState` function that allows a subset of a state that was previously returned by
:ref:`mj_getState` to be extracted without having to be written back into ``mjData`` first.
- Tendon paths can now be queried from Python via ``MjsTendon.path``, the returned object
is iterable and indexing it will give the ``MjsWrap`` at the given index in the path.
- ``MjsWrap`` now exposes:
- ``type -> mujoco.mjtWrap``
- ``target -> MjsSite|MjsJoint|MjsGeom|None``
- ``sidesite -> MjsSite|None``
- ``coef -> real``
- ``divisor -> real``
Version 3.3.7 (October 13, 2025)
-----------------------------------
doc/includes/references.h
+6
View File
@@ -3500,6 +3500,10 @@ mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
mjsElement* mjs_getWrapTarget(mjsWrap* wrap);
mjsSite* mjs_getWrapSideSite(mjsWrap* wrap);
double mjs_getWrapDivisor(mjsWrap* wrap);
double mjs_getWrapCoef(mjsWrap* wrap);
int mjs_setName(mjsElement* element, const char* name);
void mjs_setBuffer(mjByteVec* dest, const void* array, int size);
void mjs_setString(mjString* dest, const char* text);
@@ -3515,6 +3519,8 @@ void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes);
mjString* mjs_getName(mjsElement* element);
const char* mjs_getString(const mjString* source);
const double* mjs_getDouble(const mjDoubleVec* source, int* size);
int mjs_getWrapNum(const mjsTendon* tendonspec);
mjsWrap* mjs_getWrap(const mjsTendon* tendonspec, int i);
const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
void mjs_setDefault(mjsElement* element, const mjsDefault* def);
int mjs_setFrame(mjsElement* dest, mjsFrame* frame);
include/mujoco/mujoco.h
+17
View File
@@ -1678,6 +1678,17 @@ MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
// Return spec's next element; return NULL if element is last.
MJAPI mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
// Get wrapped element in tendon path.
MJAPI mjsElement* mjs_getWrapTarget(mjsWrap* wrap);
// Get wrapped element side site in tendon path if it has one, nullptr otherwise.
MJAPI mjsSite* mjs_getWrapSideSite(mjsWrap* wrap);
// Get divisor of mjsWrap wrapping a puller.
MJAPI double mjs_getWrapDivisor(mjsWrap* wrap);
// Get coefficient of mjsWrap wrapping a joint.
MJAPI double mjs_getWrapCoef(mjsWrap* wrap);
//---------------------------------- Attribute setters ---------------------------------------------
@@ -1730,6 +1741,12 @@ MJAPI const char* mjs_getString(const mjString* source);
// Nullable: size
MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size);
// Get number of elements a tendon wraps.
MJAPI int mjs_getWrapNum(const mjsTendon* tendonspec);
// Get mjsWrap element at position i in the tendon path.
MJAPI mjsWrap* mjs_getWrap(const mjsTendon* tendonspec, int i);
// Get plugin attributes.
MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
python/mujoco/introspect/functions.py
+94
View File
@@ -10471,6 +10471,66 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc="Return spec's next element; return NULL if element is last.",
)),
('mjs_getWrapTarget',
FunctionDecl(
name='mjs_getWrapTarget',
return_type=PointerType(
inner_type=ValueType(name='mjsElement'),
),
parameters=(
FunctionParameterDecl(
name='wrap',
type=PointerType(
inner_type=ValueType(name='mjsWrap'),
),
),
),
doc='Get wrapped element in tendon path.',
)),
('mjs_getWrapSideSite',
FunctionDecl(
name='mjs_getWrapSideSite',
return_type=PointerType(
inner_type=ValueType(name='mjsSite'),
),
parameters=(
FunctionParameterDecl(
name='wrap',
type=PointerType(
inner_type=ValueType(name='mjsWrap'),
),
),
),
doc='Get wrapped element side site in tendon path if it has one, nullptr otherwise.', # pylint: disable=line-too-long
)),
('mjs_getWrapDivisor',
FunctionDecl(
name='mjs_getWrapDivisor',
return_type=ValueType(name='double'),
parameters=(
FunctionParameterDecl(
name='wrap',
type=PointerType(
inner_type=ValueType(name='mjsWrap'),
),
),
),
doc='Get divisor of mjsWrap wrapping a puller.',
)),
('mjs_getWrapCoef',
FunctionDecl(
name='mjs_getWrapCoef',
return_type=ValueType(name='double'),
parameters=(
FunctionParameterDecl(
name='wrap',
type=PointerType(
inner_type=ValueType(name='mjsWrap'),
),
),
),
doc='Get coefficient of mjsWrap wrapping a joint.',
)),
('mjs_setName',
FunctionDecl(
name='mjs_setName',
@@ -10794,6 +10854,40 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Get double array contents and optionally its size.',
)),
('mjs_getWrapNum',
FunctionDecl(
name='mjs_getWrapNum',
return_type=ValueType(name='int'),
parameters=(
FunctionParameterDecl(
name='tendonspec',
type=PointerType(
inner_type=ValueType(name='mjsTendon', is_const=True),
),
),
),
doc='Get number of elements a tendon wraps.',
)),
('mjs_getWrap',
FunctionDecl(
name='mjs_getWrap',
return_type=PointerType(
inner_type=ValueType(name='mjsWrap'),
),
parameters=(
FunctionParameterDecl(
name='tendonspec',
type=PointerType(
inner_type=ValueType(name='mjsTendon', is_const=True),
),
),
FunctionParameterDecl(
name='i',
type=ValueType(name='int'),
),
),
doc='Get mjsWrap element at position i in the tendon path.',
)),
('mjs_getPluginAttributes',
FunctionDecl(
name='mjs_getPluginAttributes',
python/mujoco/specs.cc
+79 -1
View File
@@ -1304,16 +1304,94 @@ PYBIND11_MODULE(_specs, m) {
},
py::arg("gain"));
// ============================= MJSTENDON ===================================
// ============================= MJSTENDONPATH ===============================
// helper struct for tendon path indexing
struct MjsTendonPath {
raw::MjsTendon* tendon;
};
py::class_<MjsTendonPath>(m, "MjsTendonPath")
.def("__getitem__",
[](MjsTendonPath& self, int i) -> const raw::MjsWrap* {
int num_wrap = mjs_getWrapNum(self.tendon);
if (i < 0 || i >= num_wrap) {
throw py::index_error("Index out of range.");
}
return mjs_getWrap(self.tendon, i);
},
py::return_value_policy::reference_internal)
.def("__len__", [](MjsTendonPath& self) {
return mjs_getWrapNum(self.tendon);
});
// ============================= MJSWRAP =====================================
mjsWrap.def_property_readonly(
"target",
[](raw::MjsWrap& self) -> py::object {
raw::MjsElement* target = mjs_getWrapTarget(&self);
if (!target) {
return py::none();
}
switch (target->elemtype) {
case mjOBJ_SITE:
return py::cast(mjs_asSite(target));
case mjOBJ_GEOM:
return py::cast(mjs_asGeom(target));
case mjOBJ_JOINT:
return py::cast(mjs_asJoint(target));
default:
throw pybind11::value_error("Unsupported wrap target type: " +
std::to_string(target->elemtype));
}
return py::none();
},
py::return_value_policy::reference_internal);
mjsWrap.def_property_readonly(
"sidesite",
[](raw::MjsWrap& self) -> raw::MjsSite* {
return mjs_getWrapSideSite(&self);
},
py::return_value_policy::reference_internal);
mjsWrap.def_property_readonly(
"divisor",
[](raw::MjsWrap& self) -> py::object {
if (self.type != mjWRAP_PULLEY) {
return py::none();
}
return py::cast(mjs_getWrapDivisor(&self));
},
py::return_value_policy::reference_internal);
mjsWrap.def_property_readonly(
"coef",
[](raw::MjsWrap& self) -> py::object {
if (self.type != mjWRAP_JOINT) {
return py::none();
}
return py::cast(mjs_getWrapCoef(&self));
},
py::return_value_policy::reference_internal);
mjSpec.def("delete", [](MjSpec& self, raw::MjsTendon& obj) {
mjs_delete(self.ptr, obj.element);
});
// ============================= MJSTENDON ===================================
mjsTendon.def(
"default",
[](raw::MjsTendon& self) -> raw::MjsDefault* {
return mjs_getDefault(self.element);
},
py::return_value_policy::reference_internal);
mjsTendon.def_property_readonly(
"path",
[](raw::MjsTendon& self) {
return MjsTendonPath{&self};
},
py::return_value_policy::reference_internal);
mjsTendon.def(
"wrap_site",
[](raw::MjsTendon& self, std::string& name) {
python/mujoco/specs_test.py
+54 -10
View File
@@ -1511,27 +1511,28 @@ class SpecsTest(absltest.TestCase):
body = spec.worldbody.add_body(name='body')
body.add_geom(name='body_geom', pos=[0, 0, 0], size=[.1, 0, 0])
body.add_site(name='site1', pos=[0, 0, 0])
body.add_site(name='site2', pos=[0, 0, -1])
body.add_site(name='site3', pos=[0, 0, -4])
body.add_site(name='site4', pos=[0, 1, -6])
site1 = body.add_site(name='site1', pos=[0, 0, 0])
site2 = body.add_site(name='site2', pos=[0, 0, -1])
site3 = body.add_site(name='site3', pos=[0, 0, -4])
sidesite = body.add_site(name='sidesite', pos=[2, 0, -5])
site4 = body.add_site(name='site4', pos=[0, 1, -6])
spec.worldbody.add_geom(name='sphere', size=[.2, 0, 0], pos=[0, 0, -2])
sphere = spec.worldbody.add_geom(name='sphere', size=[.2, 0, 0], pos=[0, 0, -2])
spec.worldbody.add_geom(
cylinder = spec.worldbody.add_geom(
name='cylinder',
type=mujoco.mjtGeom.mjGEOM_CYLINDER,
size=[0.1, 0.2, 0.3],
pos=[0, 0, -5]
)
body.add_joint(
joint1 = body.add_joint(
name='joint1', type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0]
)
body2 = spec.worldbody.add_body(name='body2', pos=[2, 0, 0])
body2.add_geom(name='body2_geom', pos=[0, 0, 0], size=[.1, 0, 0])
body2.add_joint(
joint2 = body2.add_joint(
name='joint2', type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0]
)
@@ -1546,11 +1547,50 @@ class SpecsTest(absltest.TestCase):
wrap_site4_1 = spatial_tendon.wrap_site('site4')
wrap_pulley2 = spatial_tendon.wrap_pulley(2.0)
wrap_site3_2 = spatial_tendon.wrap_site('site3')
wrap_cylinder = spatial_tendon.wrap_geom('cylinder', '')
wrap_cylinder = spatial_tendon.wrap_geom('cylinder', 'sidesite')
wrap_site4_2 = spatial_tendon.wrap_site('site4')
wrap_joint1 = fixed_tendon.wrap_joint('joint1', 1.0)
wrap_joint2 = fixed_tendon.wrap_joint('joint2', 1.0)
wrap_joint2 = fixed_tendon.wrap_joint('joint2', 2.0)
self.assertListEqual(
list(spatial_tendon.path),
[
wrap_site1,
wrap_site2,
wrap_pulley1,
wrap_site3_1,
wrap_sphere,
wrap_site4_1,
wrap_pulley2,
wrap_site3_2,
wrap_cylinder,
wrap_site4_2,
],
)
self.assertListEqual(
[w.target for w in spatial_tendon.path],
[
site1,
site2,
None, # Pulley wraps have no targets
site3,
sphere,
site4,
None, # Pulley wraps have no targets
site3,
cylinder,
site4,
],
)
self.assertEqual(spatial_tendon.path[8].sidesite, sidesite)
self.assertIsNone(spatial_tendon.path[7].sidesite)
self.assertListEqual(list(fixed_tendon.path), [wrap_joint1, wrap_joint2])
self.assertListEqual(
[w.target for w in fixed_tendon.path],
[joint1, joint2]
)
# Wrap type for geom is only set during compilation.
spec.compile()
@@ -1562,12 +1602,16 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(wrap_site3_2.type, mujoco.mjtWrap.mjWRAP_SITE)
self.assertEqual(wrap_site4_2.type, mujoco.mjtWrap.mjWRAP_SITE)
self.assertEqual(wrap_pulley1.type, mujoco.mjtWrap.mjWRAP_PULLEY)
self.assertEqual(wrap_pulley1.divisor, 2.0)
self.assertEqual(wrap_sphere.type, mujoco.mjtWrap.mjWRAP_SPHERE)
self.assertEqual(wrap_cylinder.type, mujoco.mjtWrap.mjWRAP_CYLINDER)
self.assertEqual(wrap_pulley2.type, mujoco.mjtWrap.mjWRAP_PULLEY)
self.assertEqual(wrap_pulley2.divisor, 2.0)
self.assertEqual(wrap_joint1.type, mujoco.mjtWrap.mjWRAP_JOINT)
self.assertEqual(wrap_joint1.coef, 1.0)
self.assertEqual(wrap_joint2.type, mujoco.mjtWrap.mjWRAP_JOINT)
self.assertEqual(wrap_joint2.coef, 2.0)
if __name__ == '__main__':
absltest.main()
src/user/user_api.cc
+81
View File
@@ -1338,6 +1338,76 @@ mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element) {
mjsElement* mjs_getWrapTarget(mjsWrap* wrap) {
mjCWrap* cwrap = static_cast<mjCWrap*>(wrap->element);
mjtObj type = mjOBJ_UNKNOWN;
switch (cwrap->Type()) {
case mjWRAP_SPHERE:
case mjWRAP_CYLINDER:
type = mjOBJ_GEOM;
break;
case mjWRAP_SITE:
type = mjOBJ_SITE;
break;
case mjWRAP_JOINT:
type = mjOBJ_JOINT;
break;
case mjWRAP_PULLEY:
// Pulleys have no target.
return nullptr;
default:
return nullptr;
}
mjSpec* spec = mjs_getSpec(wrap->element);
mjsElement* target = mjs_findElement(spec, type, cwrap->name.c_str());
return target;
}
mjsSite* mjs_getWrapSideSite(mjsWrap* wrap) {
mjCWrap* cwrap = static_cast<mjCWrap*>(wrap->element);
// only sphere and cylinder (geoms) have side sites
if ((cwrap->Type() != mjWRAP_SPHERE &&
cwrap->Type() != mjWRAP_CYLINDER) ||
cwrap->sidesite.empty()) {
return nullptr;
}
mjSpec* spec = mjs_getSpec(wrap->element);
mjsElement* site = mjs_findElement(spec, mjOBJ_SITE, cwrap->sidesite.c_str());
if (site == nullptr) {
mju_warning("Could not find side site %s for wrap %s in spec",
cwrap->sidesite.c_str(), cwrap->name.c_str());
return nullptr;
}
return mjs_asSite(site);
}
double mjs_getWrapDivisor(mjsWrap* wrap) {
mjCWrap* cwrap = static_cast<mjCWrap*>(wrap->element);
if (cwrap->Type() != mjWRAP_PULLEY) {
mju_warning("Querying divisor attribute of non-pulley wrap: %s", cwrap->name.c_str());
return 1.0;
}
return cwrap->prm;
}
double mjs_getWrapCoef(mjsWrap* wrap) {
mjCWrap* cwrap = static_cast<mjCWrap*>(wrap->element);
if (cwrap->Type() != mjWRAP_JOINT) {
mju_warning("Querying coef attribute of non-joint wrap: %s", cwrap->name.c_str());
return 1.0;
}
return cwrap->prm;
}
// return body given mjsElement
mjsBody* mjs_asBody(mjsElement* element) {
if (element && element->elemtype == mjOBJ_BODY) {
@@ -1712,7 +1782,18 @@ const double* mjs_getDouble(const mjDoubleVec* source, int* size) {
return source->data();
}
int mjs_getWrapNum(const mjsTendon* tendonspec) {
mjCTendon* tendon = static_cast<mjCTendon*>(tendonspec->element);
return tendon->NumWraps();
}
mjsWrap* mjs_getWrap(const mjsTendon* tendonspec, int i) {
mjCTendon* tendon = static_cast<mjCTendon*>(tendonspec->element);
if (i < 0 || i >= tendon->NumWraps()) {
mju_error("Wrap index out of range (0, %d)", tendon->NumWraps());
}
return &const_cast<mjCWrap*>(tendon->GetWrap(i))->spec;
}
// set plugin attributes
void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes) {
src/user/user_api.h
+17
View File
@@ -266,6 +266,18 @@ MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
// Return spec's next element; return NULL if element is last.
MJAPI mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
// Get wrapped element in tendon path.
MJAPI mjsElement* mjs_getWrapTarget(mjsWrap* wrap);
// Get wrapped element in tendon path.
MJAPI mjsSite* mjs_getWrapSideSite(mjsWrap* wrap);
// Get divisor of mjsWrap wrapping a puller.
MJAPI double mjs_getWrapDivisor(mjsWrap* wrap);
// Get coefficient of mjsWrap wrapping a joint.
MJAPI double mjs_getWrapCoef(mjsWrap* wrap);
// Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody.
MJAPI mjsBody* mjs_asBody(mjsElement* element);
@@ -389,6 +401,11 @@ MJAPI const char* mjs_getString(const mjString* source);
// Get double array contents and optionally its size.
MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size);
// Get number of elements a tendon wraps.
MJAPI int mjs_getWrapNum(const mjsTendon* tendonspec);
MJAPI mjsWrap* mjs_getWrap(const mjsTendon* tendonspec, int i);
// Get plugin attributes.
MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin);