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

Polynomial stiffness and damping https://youtu.be/aKa3ZlEF9_Y

Browse Source 
PiperOrigin-RevId: 884607673
Change-Id: If8088dbf37fed1055304778a7eb84dec52cba920
This commit is contained in:
Yuval Tassa
2026-03-16 13:24:44 -07:00
committed by Copybara-Service
parent aec1b45dce
commit efae9157a7
38 changed files with 1093 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/bindings_test.py
+23
View File
@@ -1881,6 +1881,29 @@ Euler integrator, semi-implicit in velocity.
mujoco.mj_readSensor(model, data, 0, delay, result, interp=0)
np.testing.assert_array_equal(result, [1, 0, 0, 0])
def test_poly_stiffness(self):
xml = r"""
<mujoco>
<worldbody>
<body>
<geom type="sphere" size="0.1"/>
<joint name="slide" type="slide" stiffness="1 2 3"/>
</body>
</worldbody>
</mujoco>
"""
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)
self.assertEqual(model.jnt_stiffness[0], 1)
np.testing.assert_array_equal(
model.jnt_stiffnesspoly[0], [2, 3])
x = 0.5
data.qpos[0] = x
mujoco.mj_forward(model, data)
expected = -(1*x + 2*x*abs(x) + 3*x**3)
np.testing.assert_allclose(data.qfrc_spring[0], expected)
def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
for name in attr_to_compare:
actual_value = getattr(actual_obj, name)
python/mujoco/codegen/generate_spec_bindings.py
+40 -19
View File
@@ -486,13 +486,20 @@ def generate_add() -> None:
f'py::object& {f.name}',
f'py::arg("{f.name}") = py::none()',
)
if f.name == 'size' and f.type.extents[0] == 3:
if (
f.name == 'size'
and f.type.extents[0] == 3
or f.name in ('stiffness', 'damping')
):
return (
f'set_array_size(out->{f.name}, {f.name});',
'array_size',
(
f'set_array_padded(out->{f.name}, {f.name},'
f' {f.type.extents[0]}, "{f.name}");'
),
'array_padded',
f.name,
'list[float]',
f'std::optional<std::vector<double>>& {f.name}',
'Optional[list[float]]',
f'std::optional<py::object>& {f.name}',
f'py::arg("{f.name}") = py::none()',
)
return (
@@ -782,6 +789,34 @@ def generate_add() -> None:
}
};
"""
elif t == 'array_padded':
code += """\n
auto set_array_padded = [](auto&& des, const std::optional<py::object>& obj, int size, const char* name) {
if (obj.has_value() && !obj->is_none()) {
std::vector<double> array;
if (py::isinstance<py::int_>(*obj) || py::isinstance<py::float_>(*obj)) {
array.push_back(py::cast<double>(*obj));
} else if (py::isinstance<py::str>(*obj)) {
throw pybind11::type_error(
std::string(name) + " should be a numeric scalar or list.");
} else {
try {
array = py::cast<std::vector<double>>(*obj);
} catch (const py::cast_error&) {
throw pybind11::type_error(
std::string(name) + " should be a numeric scalar or list.");
}
}
if (array.empty() || array.size() > static_cast<size_t>(size)) {
std::string msg = std::string(name) + " should be a list/array of size 1 to " + std::to_string(size) + ".";
throw pybind11::value_error(msg);
}
for (int i = 0; i < size; i++) {
des[i] = (i < static_cast<int>(array.size())) ? array[i] : 0;
}
}
};
"""
elif t == 'char_array':
code += """\n
auto set_char_array = [](auto&& des, py::object& obj, int size, const char* name) {
@@ -814,20 +849,6 @@ def generate_add() -> None:
}
};
"""
elif t == 'array_size':
code += """\n
auto set_array_size = [](auto&& des, const std::optional<std::vector<double>>& array) {
if (array.has_value()) {
if (array->size() < 1 || array->size() > 3) {
std::string msg = "size should be a list/array of size 1, 2, or 3.";
throw pybind11::value_error(msg);
}
for (int i = 0; i < 3; i++) {
des[i] = (i < array->size()) ? array->at(i) : 0;
}
}
};
"""
elif t == 'value':
code += """\n
auto set_value = [](auto&& des, auto&& val) {
python/mujoco/constants.cc
+1
View File
@@ -70,6 +70,7 @@ PYBIND11_MODULE(_constants, pymodule) {
X(mjNSENS);
X(mjNSOLVER);
X(mjNISLAND);
X(mjNPOLY);
// from mjvisualize.h
X(mjNGROUP);
python/mujoco/indexer_xmacro.h
+5 -1
View File
@@ -15,7 +15,7 @@
#ifndef MUJOCO_PYTHON_INDEXER_XMACRO_H_
#define MUJOCO_PYTHON_INDEXER_XMACRO_H_
#include <mujoco/mjxmacro.h>
#include <mujoco/mjxmacro.h> // IWYU pragma: keep
#define MJMODEL_ACTUATOR \
X( int, actuator_, trntype, nu, 1 ) \
@@ -132,6 +132,7 @@
X( mjtNum, jnt_, pos, njnt, 3 ) \
X( mjtNum, jnt_, axis, njnt, 3 ) \
X( mjtNum, jnt_, stiffness, njnt, 1 ) \
X( mjtNum, jnt_, stiffnesspoly, njnt, mjNPOLY ) \
X( mjtNum, jnt_, range, njnt, 2 ) \
X( mjtNum, jnt_, margin, njnt, 1 ) \
X( mjtNum, jnt_, user, njnt, MJ_M(nuser_jnt) ) \
@@ -145,6 +146,7 @@
X( mjtNum, dof_, frictionloss, nv, 1 ) \
X( mjtNum, dof_, armature, nv, 1 ) \
X( mjtNum, dof_, damping, nv, 1 ) \
X( mjtNum, dof_, dampingpoly, nv, mjNPOLY ) \
X( mjtNum, dof_, invweight0, nv, 1 ) \
X( mjtNum, dof_, M0, nv, 1 )
@@ -253,7 +255,9 @@
X( mjtNum, tendon, _range, ntendon, 2 ) \
X( mjtNum, tendon, _margin, ntendon, 1 ) \
X( mjtNum, tendon, _stiffness, ntendon, 1 ) \
X( mjtNum, tendon, _stiffnesspoly, ntendon, mjNPOLY ) \
X( mjtNum, tendon, _damping, ntendon, 1 ) \
X( mjtNum, tendon, _dampingpoly, ntendon, mjNPOLY ) \
X( mjtNum, tendon, _frictionloss, ntendon, 1 ) \
X( mjtNum, tendon, _lengthspring, ntendon, 1 ) \
X( mjtNum, tendon, _length0, ntendon, 1 ) \
python/mujoco/introspect/structs.py
+56 -12
View File
@@ -1732,9 +1732,17 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='stiffness coefficient',
doc='linear stiffness coefficient',
array_extent=('njnt',),
),
StructFieldDecl(
name='jnt_stiffnesspoly',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='high-order stiffness coefficients',
array_extent=('njnt', 'mjNPOLY'),
),
StructFieldDecl(
name='jnt_range',
type=PointerType(
@@ -1852,9 +1860,17 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='damping coefficient',
doc='linear damping coefficient',
array_extent=('nv',),
),
StructFieldDecl(
name='dof_dampingpoly',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='high-order damping coefficients',
array_extent=('nv', 'mjNPOLY'),
),
StructFieldDecl(
name='dof_invweight0',
type=PointerType(
@@ -4036,17 +4052,33 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='stiffness coefficient',
doc='linear stiffness coefficient',
array_extent=('ntendon',),
),
StructFieldDecl(
name='tendon_stiffnesspoly',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='high-order stiffness coefficients',
array_extent=('ntendon', 'mjNPOLY'),
),
StructFieldDecl(
name='tendon_damping',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='damping coefficient',
doc='linear damping coefficient',
array_extent=('ntendon',),
),
StructFieldDecl(
name='tendon_dampingpoly',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='high-order damping coefficients',
array_extent=('ntendon', 'mjNPOLY'),
),
StructFieldDecl(
name='tendon_armature',
type=PointerType(
@@ -7301,8 +7333,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
StructFieldDecl(
name='stiffness',
type=ValueType(name='double'),
doc='stiffness coefficient',
type=ArrayType(
inner_type=ValueType(name='double'),
extents=(3,),
),
doc='stiffness coefficients',
),
StructFieldDecl(
name='springref',
@@ -7371,8 +7406,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
StructFieldDecl(
name='damping',
type=ValueType(name='double'),
doc='damping coefficient',
type=ArrayType(
inner_type=ValueType(name='double'),
extents=(3,),
),
doc='damping coefficients',
),
StructFieldDecl(
name='frictionloss',
@@ -8934,8 +8972,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
StructFieldDecl(
name='stiffness',
type=ValueType(name='double'),
doc='stiffness coefficient',
type=ArrayType(
inner_type=ValueType(name='double'),
extents=(3,),
),
doc='stiffness coefficients',
),
StructFieldDecl(
name='springlength',
@@ -8947,8 +8988,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
StructFieldDecl(
name='damping',
type=ValueType(name='double'),
doc='damping coefficient',
type=ArrayType(
inner_type=ValueType(name='double'),
extents=(3,),
),
doc='damping coefficients',
),
StructFieldDecl(
name='frictionloss',
python/mujoco/specs_test.py
+3 -3
View File
@@ -169,7 +169,7 @@ class SpecsTest(absltest.TestCase):
# Add tendon.
tendon = spec.add_tendon(stiffness=2, springlength=[0.1, 0.2])
self.assertEqual(tendon.stiffness, 2)
np.testing.assert_array_equal(tendon.stiffness, [2, 0, 0])
np.testing.assert_array_equal(tendon.springlength, [0.1, 0.2])
# Add actuator.
@@ -478,14 +478,14 @@ class SpecsTest(absltest.TestCase):
body.add_geom(size=[])
self.assertEqual(
str(cm.exception),
'size should be a list/array of size 1, 2, or 3.',
'size should be a list/array of size 1 to 3.',
)
with self.assertRaises(ValueError) as cm:
body.add_geom(size=[1, 2, 3, 4])
self.assertEqual(
str(cm.exception),
'size should be a list/array of size 1, 2, or 3.',
'size should be a list/array of size 1 to 3.',
)
def test_load_xml(self):