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

Add assets to MjSpec.compile()

Browse Source 
PiperOrigin-RevId: 653965318
Change-Id: Ib9ec1c70a7241f7ee547b799a08cc58e7d4bace7
This commit is contained in:
Alessio Quaglino
2024-07-19 04:33:35 -07:00
committed by Copybara-Service
parent 85532b74e1
commit 70ac76bb4b
5 changed files with 130 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+93 -18
View File
@@ -16,6 +16,7 @@
#include <cstddef> // IWYU pragma: keep
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <string_view> // IWYU pragma: keep
#include <vector> // IWYU pragma: keep
@@ -159,6 +160,30 @@ PYBIND11_MODULE(_specs, m) {
mjSpec.def("compile", [mjmodel_from_spec_ptr](MjSpec& self) {
return mjmodel_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr));
});
mjSpec.def(
"compile",
[mjmodel_from_spec_ptr](MjSpec& self, py::dict& assets) -> py::object {
mjVFS vfs;
mj_defaultVFS(&vfs);
for (auto item : assets) {
std::string buffer = py::cast<std::string>(item.second);
mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
buffer.c_str(), buffer.size());
};
auto model =
mjmodel_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr),
reinterpret_cast<uintptr_t>(&vfs));
mj_deleteVFS(&vfs);
return model;
}, R"mydelimiter(
Compiles the spec and returns the compiled model.
Parameters
----------
assets : dict, optional
A dictionary of assets to be used by the spec. The keys are asset names
and the values are asset contents.
)mydelimiter");
mjSpec.def(
"copy_back",
[](MjSpec& self, raw::MjModel& model) {
@@ -177,24 +202,74 @@ PYBIND11_MODULE(_specs, m) {
}
return std::string(buf.get());
});
mjSpec.def("from_file", [](MjSpec& self, std::string& filename) -> void {
std::array<char, 1024> err;
err[0] = '\0';
mj_deleteSpec(self.ptr);
self.ptr = mj_parseXML(filename.c_str(), 0, err.data(), err.size());
if (!self.ptr) {
throw FatalError(std::string(err.data()));
}
});
mjSpec.def("from_string", [](MjSpec& self, std::string& xml) -> void {
std::array<char, 1024> err;
err[0] = '\0';
mj_deleteSpec(self.ptr);
self.ptr = mj_parseXMLString(xml.c_str(), 0, err.data(), err.size());
if (!self.ptr) {
throw FatalError(std::string(err.data()));
}
});
mjSpec.def(
"from_file",
[](MjSpec& self, std::string& filename,
std::optional<py::dict>& assets) -> void {
mjVFS vfs;
mj_defaultVFS(&vfs);
if (assets.has_value()) {
for (auto item : assets.value()) {
std::string buffer = py::cast<std::string>(item.second);
mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
buffer.c_str(), buffer.size());
};
}
std::array<char, 1024> err;
err[0] = '\0';
mj_deleteSpec(self.ptr);
self.ptr = mj_parseXML(filename.c_str(), &vfs, err.data(), err.size());
mj_deleteVFS(&vfs);
if (!self.ptr) {
throw FatalError(std::string(err.data()));
}
},
py::arg("filename"),
py::arg("assets") = py::none(), R"mydelimiter(
Creates a spec from an XML file.
Parameters
----------
filename : str
Path to the XML file.
assets : dict, optional
A dictionary of assets to be used by the spec. The keys are asset names
and the values are asset contents.
)mydelimiter");
mjSpec.def(
"from_string",
[](MjSpec& self, std::string& xml,
std::optional<py::dict>& assets) -> void {
mjVFS vfs;
mj_defaultVFS(&vfs);
if (assets.has_value()) {
for (auto item : assets.value()) {
std::string buffer = py::cast<std::string>(item.second);
mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
buffer.c_str(), buffer.size());
};
}
std::array<char, 1024> err;
err[0] = '\0';
mj_deleteSpec(self.ptr);
self.ptr = mj_parseXMLString(xml.c_str(), &vfs, err.data(), err.size());
mj_deleteVFS(&vfs);
if (!self.ptr) {
throw FatalError(std::string(err.data()));
}
},
py::arg("xml"),
py::arg("assets") = py::none(), R"mydelimiter(
Creates a spec from an XML string.
Parameters
----------
xml : str
XML string.
assets : dict, optional
A dictionary of assets to be used by the spec. The keys are asset names
and the values are asset contents.
)mydelimiter");
mjSpec.def(
"add_default",
[](MjSpec* spec, std::string& classname,
python/mujoco/specs_test.py
+20
View File
@@ -249,6 +249,26 @@ class SpecsTest(absltest.TestCase):
geom = spec.worldbody.next_geom(geom)
i += 1
def test_assets(self):
cube = """
v -1 -1 1
v 1 -1 1
v -1 1 1
v 1 1 1
v -1 1 -1
v 1 1 -1
v -1 -1 -1
v 1 -1 -1"""
spec = mujoco.MjSpec()
mesh = spec.add_mesh()
mesh.name = 'cube'
mesh.file = 'cube.obj'
geom = spec.worldbody.add_geom()
geom.type = mujoco.mjtGeom.mjGEOM_MESH
geom.meshname = 'cube'
model = spec.compile({'cube.obj': cube})
self.assertEqual(model.nmeshvert, 8)
if __name__ == '__main__':
absltest.main()
python/mujoco/structs.cc
+11 -3
View File
@@ -423,8 +423,9 @@ MjModelWrapper MjModelWrapper::LoadXML(
return MjModelWrapper(model);
}
MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec) {
auto m = mj_compile(spec, nullptr);
MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec,
const mjVFS* vfs) {
auto m = mj_compile(spec, vfs);
if (!m || mjs_isWarning(spec)) {
throw py::value_error(mjs_getError(spec));
}
@@ -1572,7 +1573,14 @@ R"(Loads an MjModel from an XML string and an optional assets dictionary.)"));
mjModel.def_static(
"_from_spec_ptr", [](uintptr_t addr) {
return MjModelWrapper::CompileSpec(
reinterpret_cast<raw::MjSpec*>(addr));
reinterpret_cast<raw::MjSpec*>(addr),
nullptr);
});
mjModel.def_static(
"_from_spec_ptr", [](uintptr_t addr, uintptr_t vfs) {
return MjModelWrapper::CompileSpec(
reinterpret_cast<raw::MjSpec*>(addr),
reinterpret_cast<mjVFS*>(vfs));
});
mjModel.def_static(
"from_xml_path", &MjModelWrapper::LoadXMLFile,
python/mujoco/structs.h
+1 -1
View File
@@ -489,7 +489,7 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
static MjWrapper CompileSpec(raw::MjSpec* spec);
static MjWrapper CompileSpec(raw::MjSpec* spec, const mjVFS* vfs);
static constexpr char kFromRawPointer[] =
"__MUJOCO_STRUCTS_MJMODELWRAPPER_LOOKUP";