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

Add mjs_makeFlex to the MuJoCo C and Python APIs.

Browse Source 
- Add mjs_makeFlex C API that wraps mjCFlexcomp::Make(), creating
  a flex with auto-generated bodies, joints, and equality constraints
  from parameters (type, dim, dof, count, cellcount, spacing, scale,
  radius, mass, inertiabox, equality, rigid, flatskin, elastic2d,
  pos, quat, origin, file, vfs).
- Add Python binding body.make_flex() in specs.cc with keyword args.
- Refactor flex_spec_provision.py to use make_flex() instead of XML
  string templating + separate spec + attach. Eliminates textwrap
  dependency.
- Merge user_flexcomp into user_model BUILD target to resolve
  circular dependency.

PiperOrigin-RevId: 922725304
Change-Id: Ice13787c45f3173407fec659f9590a6453bc65dd
This commit is contained in:
Alessio Quaglino
2026-05-28 05:58:47 -07:00
committed by Copybara-Service
parent b9fb817af8
commit 4c381635e1
11 changed files with 506 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/introspect/functions.py
+138
View File
@@ -10366,6 +10366,144 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Add flex.',
)),
('mjs_makeFlex',
FunctionDecl(
name='mjs_makeFlex',
return_type=PointerType(
inner_type=ValueType(name='mjsFlex'),
),
parameters=(
FunctionParameterDecl(
name='body',
type=PointerType(
inner_type=ValueType(name='mjsBody'),
),
),
FunctionParameterDecl(
name='name',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='type',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
nullable=True,
),
FunctionParameterDecl(
name='dim',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='dof',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
nullable=True,
),
FunctionParameterDecl(
name='count',
type=ArrayType(
inner_type=ValueType(name='int', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='cellcount',
type=ArrayType(
inner_type=ValueType(name='int', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='spacing',
type=ArrayType(
inner_type=ValueType(name='double', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='scale',
type=ArrayType(
inner_type=ValueType(name='double', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='radius',
type=ValueType(name='double'),
),
FunctionParameterDecl(
name='mass',
type=ValueType(name='double'),
),
FunctionParameterDecl(
name='inertiabox',
type=ValueType(name='double'),
),
FunctionParameterDecl(
name='equality',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='rigid',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='flatskin',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='elastic2d',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='pos',
type=ArrayType(
inner_type=ValueType(name='double', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='quat',
type=ArrayType(
inner_type=ValueType(name='double', is_const=True),
extents=(4,),
),
nullable=True,
),
FunctionParameterDecl(
name='origin',
type=ArrayType(
inner_type=ValueType(name='double', is_const=True),
extents=(3,),
),
nullable=True,
),
FunctionParameterDecl(
name='file',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
nullable=True,
),
FunctionParameterDecl(
name='vfs',
type=PointerType(
inner_type=ValueType(name='mjVFS', is_const=True),
),
nullable=True,
),
),
doc='Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.', # pylint: disable=line-too-long
)),
('mjs_addPair',
FunctionDecl(
name='mjs_addPair',
python/mujoco/specs.cc
+124
View File
@@ -842,6 +842,130 @@ PYBIND11_MODULE(_specs, m) {
return FindAllImpl(self, objtype, true);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"make_flex",
[](raw::MjsBody& self,
const std::string& name,
std::optional<std::string> type,
int dim,
std::optional<std::string> dof,
std::optional<std::vector<int>> count,
std::optional<std::vector<int>> cellcount,
std::optional<std::vector<double>> spacing,
std::optional<std::vector<double>> scale,
double radius,
double mass,
double inertiabox,
int equality,
int rigid,
int flatskin,
int elastic2d,
std::optional<std::vector<double>> pos,
std::optional<std::vector<double>> quat,
std::optional<std::vector<double>> origin,
std::optional<std::string> file,
MjVfs* vfs) -> raw::MjsFlex* {
const char* type_str = type.has_value() ? type->c_str() : nullptr;
const char* dof_str = dof.has_value() ? dof->c_str() : nullptr;
const char* file_str = file.has_value() ? file->c_str() : nullptr;
const mjVFS* vfs_ptr = vfs ? vfs->get() : nullptr;
int count_arr[3] = {10, 10, 10};
if (count.has_value()) {
if (count->size() != 3) {
throw pybind11::value_error("count must have 3 elements");
}
for (int i = 0; i < 3; i++) count_arr[i] = (*count)[i];
}
const int* count_ptr = count.has_value() ? count_arr : nullptr;
int cellcount_arr[3] = {-1, -1, -1};
if (cellcount.has_value()) {
if (cellcount->size() != 3) {
throw pybind11::value_error("cellcount must have 3 elements");
}
for (int i = 0; i < 3; i++) cellcount_arr[i] = (*cellcount)[i];
}
const int* cellcount_ptr =
cellcount.has_value() ? cellcount_arr : nullptr;
double spacing_arr[3] = {0.02, 0.02, 0.02};
if (spacing.has_value()) {
if (spacing->size() != 3) {
throw pybind11::value_error("spacing must have 3 elements");
}
for (int i = 0; i < 3; i++) spacing_arr[i] = (*spacing)[i];
}
const double* spacing_ptr = spacing.has_value() ? spacing_arr : nullptr;
double scale_arr[3] = {1, 1, 1};
if (scale.has_value()) {
if (scale->size() != 3) {
throw pybind11::value_error("scale must have 3 elements");
}
for (int i = 0; i < 3; i++) scale_arr[i] = (*scale)[i];
}
const double* scale_ptr = scale.has_value() ? scale_arr : nullptr;
double pos_arr[3] = {0, 0, 0};
if (pos.has_value()) {
if (pos->size() != 3) {
throw pybind11::value_error("pos must have 3 elements");
}
for (int i = 0; i < 3; i++) pos_arr[i] = (*pos)[i];
}
const double* pos_ptr = pos.has_value() ? pos_arr : nullptr;
double quat_arr[4] = {1, 0, 0, 0};
if (quat.has_value()) {
if (quat->size() != 4) {
throw pybind11::value_error("quat must have 4 elements");
}
for (int i = 0; i < 4; i++) quat_arr[i] = (*quat)[i];
}
const double* quat_ptr = quat.has_value() ? quat_arr : nullptr;
double origin_arr[3] = {0, 0, 0};
if (origin.has_value()) {
if (origin->size() != 3) {
throw pybind11::value_error("origin must have 3 elements");
}
for (int i = 0; i < 3; i++) origin_arr[i] = (*origin)[i];
}
const double* origin_ptr = origin.has_value() ? origin_arr : nullptr;
auto out = mjs_makeFlex(
&self, name.c_str(), type_str, dim, dof_str,
count_ptr, cellcount_ptr, spacing_ptr, scale_ptr,
radius, mass, inertiabox, equality, rigid, flatskin, elastic2d,
pos_ptr, quat_ptr, origin_ptr, file_str, vfs_ptr);
if (!out) {
raw::MjSpec* spec = mjs_getSpec(self.element);
throw pybind11::value_error(mjs_getError(spec));
}
return out;
},
py::arg("name"),
py::arg("type") = py::none(),
py::arg("dim") = 3,
py::arg("dof") = py::none(),
py::arg("count") = py::none(),
py::arg("cellcount") = py::none(),
py::arg("spacing") = py::none(),
py::arg("scale") = py::none(),
py::arg("radius") = 0.0,
py::arg("mass") = 1.0,
py::arg("inertiabox") = 0.005,
py::arg("equality") = 0,
py::arg("rigid") = 0,
py::arg("flatskin") = 0,
py::arg("elastic2d") = 0,
py::arg("pos") = py::none(),
py::arg("quat") = py::none(),
py::arg("origin") = py::none(),
py::arg("file") = py::none(),
py::arg("vfs") = py::none(),
py::return_value_policy::reference_internal);
mjsBody.def(
"find_child",
[](raw::MjsBody& self, std::string& name) -> raw::MjsBody* {
python/mujoco/specs_test.py
+70
View File
@@ -2032,5 +2032,75 @@ class SpecsTest(absltest.TestCase):
with self.assertRaises(mujoco.FatalError):
spec.encode(filename, model)
def test_make_flex_grid(self):
# Create a spec with a flexcomp grid.
spec = mujoco.MjSpec()
body = spec.worldbody.add_body(name='flex_body')
flex = body.make_flex(
name='test_flex',
type='grid',
dim=3,
count=[4, 4, 4],
spacing=[0.05, 0.05, 0.05],
mass=0.5,
equality=1,
)
self.assertIsNotNone(flex)
model = spec.compile()
self.assertIsNotNone(model)
self.assertGreater(model.nflex, 0)
# Verify elastic2d is forwarded to Make() and produces shell-mode
# strain constraints (equality=3 + elastic2d=2 triggers shell path).
spec2 = mujoco.MjSpec()
body2 = spec2.worldbody.add_body(name='shell_body')
flex2 = body2.make_flex(
name='shell_flex',
type='grid',
dim=3,
count=[3, 3, 3],
spacing=[0.1, 0.1, 0.1],
dof='trilinear',
cellcount=[2, 2, 1],
mass=0.5,
equality=3, # strain
elastic2d=2, # bend
)
flex2.young = 1e3
flex2.thickness = 0.01
flex2.selfcollide = mujoco.mjtFlexSelf.mjFLEXSELF_NONE
self.assertIsNotNone(flex2)
model2 = spec2.compile()
self.assertIsNotNone(model2)
# Shell mode creates face-based constraints; verify they exist.
self.assertGreater(model2.neq, 0)
def test_make_flex_defaults(self):
# Create a spec with minimal flexcomp args (defaults).
spec = mujoco.MjSpec()
body = spec.worldbody.add_body(name='flex_body')
flex = body.make_flex(name='default_flex', equality=1)
self.assertIsNotNone(flex)
model = spec.compile()
self.assertIsNotNone(model)
def test_make_flex_with_pos_quat(self):
# Create a spec with flexcomp that has a pose.
spec = mujoco.MjSpec()
body = spec.worldbody.add_body(name='flex_body')
flex = body.make_flex(
name='posed_flex',
type='grid',
dim=2,
count=[3, 3, 1],
spacing=[0.1, 0.1, 0.1],
pos=[1.0, 2.0, 3.0],
quat=[1.0, 0.0, 0.0, 0.0],
equality=1,
)
self.assertIsNotNone(flex)
model = spec.compile()
self.assertIsNotNone(model)
if __name__ == '__main__':
absltest.main()