From a0945e4b5a65e9197a658ee1e23b858cdc27a01f Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Thu, 4 Jul 2024 07:23:05 -0700 Subject: [PATCH] Add non-tree element lists to python bindings. Fixes #1752. PiperOrigin-RevId: 649405039 Change-Id: Ie0b575a3ec912b773141694397ac2f394bff1227 --- doc/APIreference/functions.rst | 162 ++++++++++++++++++ doc/includes/references.h | 18 ++ include/mujoco/mujoco.h | 54 ++++++ introspect/functions.py | 298 +++++++++++++++++++++++++++++++++ python/mujoco/specs.cc | 192 +++++++++++++++++++++ python/mujoco/specs_test.py | 13 ++ src/user/user_api.cc | 128 ++++++++++++++ src/user/user_api.h | 54 ++++++ src/user/user_model.cc | 64 +++++++ src/user/user_model.h | 1 + 10 files changed, 984 insertions(+) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 238d65b8..1e543270 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -4176,6 +4176,24 @@ mjs_nextChild Return body's next child of the same type; return NULL if child is last. +.. _mjs_firstElement: + +mjs_firstElement +~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_firstElement + +Return spec's first element of selected type. + +.. _mjs_nextElement: + +mjs_nextElement +~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_nextElement + +Return spec's next element; return NULL if element is last. + .. _mjs_asBody: mjs_asBody @@ -4239,6 +4257,150 @@ mjs_asFrame Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame. +.. _mjs_asActuator: + +mjs_asActuator +~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asActuator + +Safely cast an element as mjsActuator, or return NULL if the element is not an mjsActuator. + +.. _mjs_asSensor: + +mjs_asSensor +~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asSensor + +Safely cast an element as mjsSensor, or return NULL if the element is not an mjsSensor. + +.. _mjs_asFlex: + +mjs_asFlex +~~~~~~~~~~ + +.. mujoco-include:: mjs_asFlex + +Safely cast an element as mjsFlex, or return NULL if the element is not an mjsFlex. + +.. _mjs_asPair: + +mjs_asPair +~~~~~~~~~~ + +.. mujoco-include:: mjs_asPair + +Safely cast an element as mjsPair, or return NULL if the element is not an mjsPair. + +.. _mjs_asEquality: + +mjs_asEquality +~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asEquality + +Safely cast an element as mjsEquality, or return NULL if the element is not an mjsEquality. + +.. _mjs_asExclude: + +mjs_asExclude +~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asExclude + +Safely cast an element as mjsExclude, or return NULL if the element is not an mjsExclude. + +.. _mjs_asTendon: + +mjs_asTendon +~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asTendon + +Safely cast an element as mjsTendon, or return NULL if the element is not an mjsTendon. + +.. _mjs_asNumeric: + +mjs_asNumeric +~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asNumeric + +Safely cast an element as mjsNumeric, or return NULL if the element is not an mjsNumeric. + +.. _mjs_asText: + +mjs_asText +~~~~~~~~~~ + +.. mujoco-include:: mjs_asText + +Safely cast an element as mjsText, or return NULL if the element is not an mjsText. + +.. _mjs_asTuple: + +mjs_asTuple +~~~~~~~~~~~ + +.. mujoco-include:: mjs_asTuple + +Safely cast an element as mjsTuple, or return NULL if the element is not an mjsTuple. + +.. _mjs_asKey: + +mjs_asKey +~~~~~~~~~ + +.. mujoco-include:: mjs_asKey + +Safely cast an element as mjsKey, or return NULL if the element is not an mjsKey. + +.. _mjs_asMesh: + +mjs_asMesh +~~~~~~~~~~ + +.. mujoco-include:: mjs_asMesh + +Safely cast an element as mjsMesh, or return NULL if the element is not an mjsMesh. + +.. _mjs_asHField: + +mjs_asHField +~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asHField + +Safely cast an element as mjsHField, or return NULL if the element is not an mjsHField. + +.. _mjs_asSkin: + +mjs_asSkin +~~~~~~~~~~ + +.. mujoco-include:: mjs_asSkin + +Safely cast an element as mjsSkin, or return NULL if the element is not an mjsSkin. + +.. _mjs_asTexture: + +mjs_asTexture +~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asTexture + +Safely cast an element as mjsTexture, or return NULL if the element is not an mjsTexture. + +.. _mjs_asMaterial: + +mjs_asMaterial +~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_asMaterial + +Safely cast an element as mjsMaterial, or return NULL if the element is not an mjsMaterial. + .. _AttributeSetters: Attribute setters diff --git a/doc/includes/references.h b/doc/includes/references.h index 1d2c724e..2e3fe36b 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3551,6 +3551,8 @@ mjsDefault* mjs_getSpecDefault(mjSpec* s); int mjs_getId(mjsElement* element); mjsElement* mjs_firstChild(mjsBody* body, mjtObj type); mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child); +mjsElement* mjs_firstElement(mjSpec* s, mjtObj type); +mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element); mjsBody* mjs_asBody(mjsElement* element); mjsGeom* mjs_asGeom(mjsElement* element); mjsJoint* mjs_asJoint(mjsElement* element); @@ -3558,6 +3560,22 @@ mjsSite* mjs_asSite(mjsElement* element); mjsCamera* mjs_asCamera(mjsElement* element); mjsLight* mjs_asLight(mjsElement* element); mjsFrame* mjs_asFrame(mjsElement* element); +mjsActuator* mjs_asActuator(mjsElement* element); +mjsSensor* mjs_asSensor(mjsElement* element); +mjsFlex* mjs_asFlex(mjsElement* element); +mjsPair* mjs_asPair(mjsElement* element); +mjsEquality* mjs_asEquality(mjsElement* element); +mjsExclude* mjs_asExclude(mjsElement* element); +mjsTendon* mjs_asTendon(mjsElement* element); +mjsNumeric* mjs_asNumeric(mjsElement* element); +mjsText* mjs_asText(mjsElement* element); +mjsTuple* mjs_asTuple(mjsElement* element); +mjsKey* mjs_asKey(mjsElement* element); +mjsMesh* mjs_asMesh(mjsElement* element); +mjsHField* mjs_asHField(mjsElement* element); +mjsSkin* mjs_asSkin(mjsElement* element); +mjsTexture* mjs_asTexture(mjsElement* element); +mjsMaterial* mjs_asMaterial(mjsElement* element); void mjs_setString(mjString* dest, const char* text); void mjs_setStringVec(mjStringVec* dest, const char* text); mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 120518ff..520ddc40 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1551,6 +1551,12 @@ MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type); // Return body's next child of the same type; return NULL if child is last. MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child); +// Return spec's first element of selected type. +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); + // Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody. MJAPI mjsBody* mjs_asBody(mjsElement* element); @@ -1572,6 +1578,54 @@ MJAPI mjsLight* mjs_asLight(mjsElement* element); // Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame. MJAPI mjsFrame* mjs_asFrame(mjsElement* element); +// Safely cast an element as mjsActuator, or return NULL if the element is not an mjsActuator. +MJAPI mjsActuator* mjs_asActuator(mjsElement* element); + +// Safely cast an element as mjsSensor, or return NULL if the element is not an mjsSensor. +MJAPI mjsSensor* mjs_asSensor(mjsElement* element); + +// Safely cast an element as mjsFlex, or return NULL if the element is not an mjsFlex. +MJAPI mjsFlex* mjs_asFlex(mjsElement* element); + +// Safely cast an element as mjsPair, or return NULL if the element is not an mjsPair. +MJAPI mjsPair* mjs_asPair(mjsElement* element); + +// Safely cast an element as mjsEquality, or return NULL if the element is not an mjsEquality. +MJAPI mjsEquality* mjs_asEquality(mjsElement* element); + +// Safely cast an element as mjsExclude, or return NULL if the element is not an mjsExclude. +MJAPI mjsExclude* mjs_asExclude(mjsElement* element); + +// Safely cast an element as mjsTendon, or return NULL if the element is not an mjsTendon. +MJAPI mjsTendon* mjs_asTendon(mjsElement* element); + +// Safely cast an element as mjsNumeric, or return NULL if the element is not an mjsNumeric. +MJAPI mjsNumeric* mjs_asNumeric(mjsElement* element); + +// Safely cast an element as mjsText, or return NULL if the element is not an mjsText. +MJAPI mjsText* mjs_asText(mjsElement* element); + +// Safely cast an element as mjsTuple, or return NULL if the element is not an mjsTuple. +MJAPI mjsTuple* mjs_asTuple(mjsElement* element); + +// Safely cast an element as mjsKey, or return NULL if the element is not an mjsKey. +MJAPI mjsKey* mjs_asKey(mjsElement* element); + +// Safely cast an element as mjsMesh, or return NULL if the element is not an mjsMesh. +MJAPI mjsMesh* mjs_asMesh(mjsElement* element); + +// Safely cast an element as mjsHField, or return NULL if the element is not an mjsHField. +MJAPI mjsHField* mjs_asHField(mjsElement* element); + +// Safely cast an element as mjsSkin, or return NULL if the element is not an mjsSkin. +MJAPI mjsSkin* mjs_asSkin(mjsElement* element); + +// Safely cast an element as mjsTexture, or return NULL if the element is not an mjsTexture. +MJAPI mjsTexture* mjs_asTexture(mjsElement* element); + +// Safely cast an element as mjsMaterial, or return NULL if the element is not an mjsMaterial. +MJAPI mjsMaterial* mjs_asMaterial(mjsElement* element); + //---------------------------------- Attribute setters --------------------------------------------- diff --git a/introspect/functions.py b/introspect/functions.py index 2aa42c88..3141750f 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -9835,6 +9835,48 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc="Return body's next child of the same type; return NULL if child is last.", # pylint: disable=line-too-long )), + ('mjs_firstElement', + FunctionDecl( + name='mjs_firstElement', + return_type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + parameters=( + FunctionParameterDecl( + name='s', + type=PointerType( + inner_type=ValueType(name='mjSpec'), + ), + ), + FunctionParameterDecl( + name='type', + type=ValueType(name='mjtObj'), + ), + ), + doc="Return spec's first element of selected type.", + )), + ('mjs_nextElement', + FunctionDecl( + name='mjs_nextElement', + return_type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + parameters=( + FunctionParameterDecl( + name='s', + type=PointerType( + inner_type=ValueType(name='mjSpec'), + ), + ), + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc="Return spec's next element; return NULL if element is last.", + )), ('mjs_asBody', FunctionDecl( name='mjs_asBody', @@ -9947,6 +9989,262 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame.', # pylint: disable=line-too-long )), + ('mjs_asActuator', + FunctionDecl( + name='mjs_asActuator', + return_type=PointerType( + inner_type=ValueType(name='mjsActuator'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsActuator, or return NULL if the element is not an mjsActuator.', # pylint: disable=line-too-long + )), + ('mjs_asSensor', + FunctionDecl( + name='mjs_asSensor', + return_type=PointerType( + inner_type=ValueType(name='mjsSensor'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsSensor, or return NULL if the element is not an mjsSensor.', # pylint: disable=line-too-long + )), + ('mjs_asFlex', + FunctionDecl( + name='mjs_asFlex', + return_type=PointerType( + inner_type=ValueType(name='mjsFlex'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsFlex, or return NULL if the element is not an mjsFlex.', # pylint: disable=line-too-long + )), + ('mjs_asPair', + FunctionDecl( + name='mjs_asPair', + return_type=PointerType( + inner_type=ValueType(name='mjsPair'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsPair, or return NULL if the element is not an mjsPair.', # pylint: disable=line-too-long + )), + ('mjs_asEquality', + FunctionDecl( + name='mjs_asEquality', + return_type=PointerType( + inner_type=ValueType(name='mjsEquality'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsEquality, or return NULL if the element is not an mjsEquality.', # pylint: disable=line-too-long + )), + ('mjs_asExclude', + FunctionDecl( + name='mjs_asExclude', + return_type=PointerType( + inner_type=ValueType(name='mjsExclude'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsExclude, or return NULL if the element is not an mjsExclude.', # pylint: disable=line-too-long + )), + ('mjs_asTendon', + FunctionDecl( + name='mjs_asTendon', + return_type=PointerType( + inner_type=ValueType(name='mjsTendon'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsTendon, or return NULL if the element is not an mjsTendon.', # pylint: disable=line-too-long + )), + ('mjs_asNumeric', + FunctionDecl( + name='mjs_asNumeric', + return_type=PointerType( + inner_type=ValueType(name='mjsNumeric'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsNumeric, or return NULL if the element is not an mjsNumeric.', # pylint: disable=line-too-long + )), + ('mjs_asText', + FunctionDecl( + name='mjs_asText', + return_type=PointerType( + inner_type=ValueType(name='mjsText'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsText, or return NULL if the element is not an mjsText.', # pylint: disable=line-too-long + )), + ('mjs_asTuple', + FunctionDecl( + name='mjs_asTuple', + return_type=PointerType( + inner_type=ValueType(name='mjsTuple'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsTuple, or return NULL if the element is not an mjsTuple.', # pylint: disable=line-too-long + )), + ('mjs_asKey', + FunctionDecl( + name='mjs_asKey', + return_type=PointerType( + inner_type=ValueType(name='mjsKey'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsKey, or return NULL if the element is not an mjsKey.', # pylint: disable=line-too-long + )), + ('mjs_asMesh', + FunctionDecl( + name='mjs_asMesh', + return_type=PointerType( + inner_type=ValueType(name='mjsMesh'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsMesh, or return NULL if the element is not an mjsMesh.', # pylint: disable=line-too-long + )), + ('mjs_asHField', + FunctionDecl( + name='mjs_asHField', + return_type=PointerType( + inner_type=ValueType(name='mjsHField'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsHField, or return NULL if the element is not an mjsHField.', # pylint: disable=line-too-long + )), + ('mjs_asSkin', + FunctionDecl( + name='mjs_asSkin', + return_type=PointerType( + inner_type=ValueType(name='mjsSkin'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsSkin, or return NULL if the element is not an mjsSkin.', # pylint: disable=line-too-long + )), + ('mjs_asTexture', + FunctionDecl( + name='mjs_asTexture', + return_type=PointerType( + inner_type=ValueType(name='mjsTexture'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsTexture, or return NULL if the element is not an mjsTexture.', # pylint: disable=line-too-long + )), + ('mjs_asMaterial', + FunctionDecl( + name='mjs_asMaterial', + return_type=PointerType( + inner_type=ValueType(name='mjsMaterial'), + ), + parameters=( + FunctionParameterDecl( + name='element', + type=PointerType( + inner_type=ValueType(name='mjsElement'), + ), + ), + ), + doc='Safely cast an element as mjsMaterial, or return NULL if the element is not an mjsMaterial.', # pylint: disable=line-too-long + )), ('mjs_setString', FunctionDecl( name='mjs_setString', diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index ce5ece39..227f8630 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -297,6 +297,198 @@ PYBIND11_MODULE(_specs, m) { mjSpec.def("detach_body", [](MjSpec& self, raw::MjsBody& body) { mjs_detachBody(self.ptr, &body); }); + mjSpec.def_property_readonly( + "actuators", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_ACTUATOR); + while (el) { + list.append(mjs_asActuator(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "sensors", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_SENSOR); + while (el) { + list.append(mjs_asSensor(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "flexes", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_FLEX); + while (el) { + list.append(mjs_asFlex(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "pairs", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_PAIR); + while (el) { + list.append(mjs_asPair(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "equality", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_EQUALITY); + while (el) { + list.append(mjs_asEquality(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "excludes", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_EXCLUDE); + while (el) { + list.append(mjs_asExclude(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "tendons", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TENDON); + while (el) { + list.append(mjs_asTendon(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "numeric", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_NUMERIC); + while (el) { + list.append(mjs_asNumeric(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "text", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TEXT); + while (el) { + list.append(mjs_asText(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "tuple", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TUPLE); + while (el) { + list.append(mjs_asTuple(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "key", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_KEY); + while (el) { + list.append(mjs_asKey(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "mesh", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_MESH); + while (el) { + list.append(mjs_asMesh(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "hfield", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_HFIELD); + while (el) { + list.append(mjs_asHField(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "skin", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_SKIN); + while (el) { + list.append(mjs_asSkin(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "texture", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TEXTURE); + while (el) { + list.append(mjs_asTexture(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); + mjSpec.def_property_readonly( + "material", + [](MjSpec& self) -> py::list { + py::list list; + raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_MATERIAL); + while (el) { + list.append(mjs_asMaterial(el)); + el = mjs_nextElement(self.ptr, el); + } + return list; + }, + py::return_value_policy::reference_internal); // ============================= MJSBODY ==================================== mjsBody.def("id", [](raw::MjsBody& self) -> int { diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 9bf3685e..b8e2f2ac 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -217,6 +217,19 @@ class SpecsTest(absltest.TestCase): """)) + def test_element_list(self): + spec = mujoco.MjSpec() + sensor1 = spec.add_sensor() + sensor2 = spec.add_sensor() + sensor3 = spec.add_sensor() + sensor1.name = 'sensor1' + sensor2.name = 'sensor2' + sensor3.name = 'sensor3' + self.assertLen(spec.sensors, 3) + self.assertEqual(spec.sensors[0].name, 'sensor1') + self.assertEqual(spec.sensors[1].name, 'sensor2') + self.assertEqual(spec.sensors[2].name, 'sensor3') + if __name__ == '__main__': absltest.main() diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 704acd53..dfef94a5 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -590,6 +590,22 @@ mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child) { +// return spec's first element of selected type +mjsElement* mjs_firstElement(mjSpec* s, mjtObj type) { + mjCModel* modelC = static_cast(s->element); + return modelC->NextObject(NULL, type); +} + + + +// return spec's next element; return NULL if element is last +mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element) { + mjCModel* modelC = static_cast(s->element); + return modelC->NextObject(element); +} + + + // return body given mjsElement mjsBody* mjs_asBody(mjsElement* element) { return element->elemtype == mjOBJ_BODY ? &(static_cast(element)->spec) : nullptr; @@ -639,6 +655,118 @@ mjsFrame* mjs_asFrame(mjsElement* element) { +// return actuator given mjsElement +mjsActuator* mjs_asActuator(mjsElement* element) { + return element->elemtype == mjOBJ_ACTUATOR ? &(static_cast(element)->spec) : nullptr; +} + + + +// return sensor given mjsElement +mjsSensor* mjs_asSensor(mjsElement* element) { + return element->elemtype == mjOBJ_SENSOR ? &(static_cast(element)->spec) : nullptr; +} + + + +// return flex given mjsElement +mjsFlex* mjs_asFlex(mjsElement* element) { + return element->elemtype == mjOBJ_FLEX ? &(static_cast(element)->spec) : nullptr; +} + + + +// return pair given mjsElement +mjsPair* mjs_asPair(mjsElement* element) { + return element->elemtype == mjOBJ_PAIR ? &(static_cast(element)->spec) : nullptr; +} + + + +// return equality given mjsElement +mjsEquality* mjs_asEquality(mjsElement* element) { + return element->elemtype == mjOBJ_EQUALITY ? &(static_cast(element)->spec) : nullptr; +} + + + +// return exclude given mjsElement +mjsExclude* mjs_asExclude(mjsElement* element) { + return element->elemtype == mjOBJ_EXCLUDE ? &(static_cast(element)->spec) : nullptr; +} + + + +// return tendon given mjsElement +mjsTendon* mjs_asTendon(mjsElement* element) { + return element->elemtype == mjOBJ_TENDON ? &(static_cast(element)->spec) : nullptr; +} + + + +// return numeric given mjsElement +mjsNumeric* mjs_asNumeric(mjsElement* element) { + return element->elemtype == mjOBJ_NUMERIC ? &(static_cast(element)->spec) : nullptr; +} + + + +// return text given mjsElement +mjsText* mjs_asText(mjsElement* element) { + return element->elemtype == mjOBJ_TEXT ? &(static_cast(element)->spec) : nullptr; +} + + + +// return tuple given mjsElement +mjsTuple* mjs_asTuple(mjsElement* element) { + return element->elemtype == mjOBJ_TUPLE ? &(static_cast(element)->spec) : nullptr; +} + + + +// return key given mjsElement +mjsKey* mjs_asKey(mjsElement* element) { + return element->elemtype == mjOBJ_KEY ? &(static_cast(element)->spec) : nullptr; +} + + + +// return mesh given mjsElement +mjsMesh* mjs_asMesh(mjsElement* element) { + return element->elemtype == mjOBJ_MESH ? &(static_cast(element)->spec) : nullptr; +} + + + +// return hfield given mjsElement +mjsHField* mjs_asHField(mjsElement* element) { + return element->elemtype == mjOBJ_HFIELD ? &(static_cast(element)->spec) : nullptr; +} + + + +// return skin given mjsElement +mjsSkin* mjs_asSkin(mjsElement* element) { + return element->elemtype == mjOBJ_SKIN ? &(static_cast(element)->spec) : nullptr; +} + + + +// return texture given mjsElement +mjsTexture* mjs_asTexture(mjsElement* element) { + return element->elemtype == mjOBJ_TEXTURE ? &(static_cast(element)->spec) : nullptr; +} + + + +// return material given mjsElement +mjsMaterial* mjs_asMaterial(mjsElement* element) { + return element->elemtype == mjOBJ_MATERIAL ? &(static_cast(element)->spec) : nullptr; +} + + + // set string void mjs_setString(mjString* dest, const char* text) { std::string* str = static_cast(dest); diff --git a/src/user/user_api.h b/src/user/user_api.h index 000f0154..7c870ae2 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -218,6 +218,12 @@ MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type); // Return body's next child of the same type; return NULL if child is last. MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child); +// Return spec's first element of selected type. +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); + // Safely cast an element as mjsBody, or return NULL if the element is not an mjsBody. MJAPI mjsBody* mjs_asBody(mjsElement* element); @@ -239,6 +245,54 @@ MJAPI mjsLight* mjs_asLight(mjsElement* element); // Safely cast an element as mjsFrame, or return NULL if the element is not an mjsFrame. MJAPI mjsFrame* mjs_asFrame(mjsElement* element); +// Safely cast an element as mjsActuator, or return NULL if the element is not an mjsActuator. +MJAPI mjsActuator* mjs_asActuator(mjsElement* element); + +// Safely cast an element as mjsSensor, or return NULL if the element is not an mjsSensor. +MJAPI mjsSensor* mjs_asSensor(mjsElement* element); + +// Safely cast an element as mjsFlex, or return NULL if the element is not an mjsFlex. +MJAPI mjsFlex* mjs_asFlex(mjsElement* element); + +// Safely cast an element as mjsPair, or return NULL if the element is not an mjsPair. +MJAPI mjsPair* mjs_asPair(mjsElement* element); + +// Safely cast an element as mjsEquality, or return NULL if the element is not an mjsEquality. +MJAPI mjsEquality* mjs_asEquality(mjsElement* element); + +// Safely cast an element as mjsExclude, or return NULL if the element is not an mjsExclude. +MJAPI mjsExclude* mjs_asExclude(mjsElement* element); + +// Safely cast an element as mjsTendon, or return NULL if the element is not an mjsTendon. +MJAPI mjsTendon* mjs_asTendon(mjsElement* element); + +// Safely cast an element as mjsNumeric, or return NULL if the element is not an mjsNumeric. +MJAPI mjsNumeric* mjs_asNumeric(mjsElement* element); + +// Safely cast an element as mjsText, or return NULL if the element is not an mjsText. +MJAPI mjsText* mjs_asText(mjsElement* element); + +// Safely cast an element as mjsTuple, or return NULL if the element is not an mjsTuple. +MJAPI mjsTuple* mjs_asTuple(mjsElement* element); + +// Safely cast an element as mjsKey, or return NULL if the element is not an mjsKey. +MJAPI mjsKey* mjs_asKey(mjsElement* element); + +// Safely cast an element as mjsMesh, or return NULL if the element is not an mjsMesh. +MJAPI mjsMesh* mjs_asMesh(mjsElement* element); + +// Safely cast an element as mjsHField, or return NULL if the element is not an mjsHField. +MJAPI mjsHField* mjs_asHField(mjsElement* element); + +// Safely cast an element as mjsSkin, or return NULL if the element is not an mjsSkin. +MJAPI mjsSkin* mjs_asSkin(mjsElement* element); + +// Safely cast an element as mjsTexture, or return NULL if the element is not an mjsTexture. +MJAPI mjsTexture* mjs_asTexture(mjsElement* element); + +// Safely cast an element as mjsMaterial, or return NULL if the element is not an mjsMaterial. +MJAPI mjsMaterial* mjs_asMaterial(mjsElement* element); + //---------------------------------- Attribute setters --------------------------------------------- diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 7159001e..50f37bf8 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -706,6 +706,70 @@ mjCBase* mjCModel::GetObject(mjtObj type, int id) { +template +static mjsElement* GetNext(std::vector& list, mjsElement* child) { + // TODO: use id for direct indexing instead of a loop + for (unsigned int i = 0; i < list.size()-1; i++) { + if (list[i]->spec.element == child) { + return list[i+1]->spec.element; + } + } + return nullptr; +} + + + +// next object of specified type +mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) { + if (type == mjOBJ_UNKNOWN) { + if (!object) { + throw mjCError(NULL, "type must be specified if no element is given"); + } else { + type = object->elemtype; + } + } else if (object && object->elemtype != type) { + throw mjCError(NULL, "element is not of requested type"); + } + + switch (type) { + case mjOBJ_ACTUATOR: + return object ? GetNext(actuators_, object) : actuators_[0]; + case mjOBJ_SENSOR: + return object ? GetNext(sensors_, object) : sensors_[0]; + case mjOBJ_FLEX: + return object ? GetNext(flexes_, object) : flexes_[0]; + case mjOBJ_PAIR: + return object ? GetNext(pairs_, object) : pairs_[0]; + case mjOBJ_EXCLUDE: + return object ? GetNext(excludes_, object) : excludes_[0]; + case mjOBJ_EQUALITY: + return object ? GetNext(equalities_, object) : equalities_[0]; + case mjOBJ_TENDON: + return object ? GetNext(tendons_, object) : tendons_[0]; + case mjOBJ_NUMERIC: + return object ? GetNext(numerics_, object) : numerics_[0]; + case mjOBJ_TEXT: + return object ? GetNext(texts_, object) : texts_[0]; + case mjOBJ_TUPLE: + return object ? GetNext(tuples_, object) : tuples_[0]; + case mjOBJ_KEY: + return object ? GetNext(keys_, object) : keys_[0]; + case mjOBJ_MESH: + return object ? GetNext(meshes_, object) : meshes_[0]; + case mjOBJ_HFIELD: + return object ? GetNext(hfields_, object) : hfields_[0]; + case mjOBJ_SKIN: + return object ? GetNext(skins_, object) : skins_[0]; + case mjOBJ_TEXTURE: + return object ? GetNext(textures_, object) : textures_[0]; + case mjOBJ_MATERIAL: + return object ? GetNext(materials_, object) : materials_[0]; + default: + return nullptr; + } +} + + //------------------------ API FOR ACCESS TO PRIVATE VARIABLES ------------------------------------- // compiled flag diff --git a/src/user/user_model.h b/src/user/user_model.h index 95adb202..06a52b8f 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -203,6 +203,7 @@ class mjCModel : public mjCModel_, private mjSpec { // API for access to model elements (outside tree) int NumObjects(mjtObj type); // number of objects in specified list mjCBase* GetObject(mjtObj type, int id); // pointer to specified object + mjsElement* NextObject(mjsElement* object, mjtObj type = mjOBJ_UNKNOWN); // next object of specified type // API for access to other variables bool IsCompiled() const; // is model already compiled