Add support for querying of MjsTendon path from Python.
Fixes #2670 PiperOrigin-RevId: 822215254 Change-Id: Ice5a264e54f963776093511c6e4ac830aae6d4f8
This commit is contained in:
Sam Haves
committed by
Copybara-Service
Copybara-Service
parent
2f65e23779
commit
ac2cd5dfd6
@@ -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',
|
||||
|
||||
+79
-1
@@ -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) {
|
||||
|
||||
+54
-10
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user