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

Make from_file() and from_string() classmethod functions.

Browse Source 
This prevents the user from modifying an mjSpec and then calling the functions to load it from string/file.

PiperOrigin-RevId: 679311284
Change-Id: I4f1ae3cc8f15bece245f624303470d538deec410
This commit is contained in:
Alessio Quaglino
2024-09-26 15:15:01 -07:00
committed by Copybara-Service
parent 8bbf9c7b4f
commit 6a415da267
3 changed files with 106 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+92 -73
View File
@@ -70,6 +70,23 @@ using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
struct MjSpec {
MjSpec() : ptr(mj_makeSpec()) {}
MjSpec(raw::MjSpec* ptr) : ptr(ptr) {}
// copy constructor and assignment
MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {}
MjSpec& operator=(const MjSpec& other) {
ptr = mj_copySpec(other.ptr);
return *this;
}
// move constructor and move assignment
MjSpec(MjSpec&& other) : ptr(other.ptr) { other.ptr = nullptr; }
MjSpec& operator=(MjSpec&& other) {
ptr = other.ptr;
other.ptr = nullptr;
return *this;
}
~MjSpec() { mj_deleteSpec(ptr); }
raw::MjSpec* ptr;
};
@@ -147,6 +164,81 @@ PYBIND11_MODULE(_specs, m) {
// ============================= MJSPEC =====================================
mjSpec.def(py::init<>());
mjSpec.def_static(
"from_file",
[](std::string& filename,
std::optional<std::unordered_map<std::string, py::bytes>>& assets)
-> MjSpec {
const auto converted_assets = _impl::ConvertAssetsDict(assets);
raw::MjSpec* spec;
{
py::gil_scoped_release no_gil;
char error[1024];
spec = LoadSpecFileImpl(
filename, converted_assets,
[&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error));
});
if (!spec) {
throw py::value_error(error);
}
}
return MjSpec(spec);
},
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", py::return_value_policy::move);
mjSpec.def_static(
"from_string",
[](std::string& xml,
std::optional<std::unordered_map<std::string, py::bytes>>& assets)
-> MjSpec {
auto converted_assets = _impl::ConvertAssetsDict(assets);
raw::MjSpec* spec;
{
py::gil_scoped_release no_gil;
std::string model_filename = "model_.xml";
if (assets.has_value()) {
while (assets->find(model_filename) != assets->end()) {
model_filename =
model_filename.substr(0, model_filename.size() - 4) + "_.xml";
}
}
converted_assets.emplace_back(
model_filename.c_str(), xml.c_str(), xml.length());
char error[1024];
spec = LoadSpecFileImpl(
model_filename, converted_assets,
[&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error));
});
if (!spec) {
throw py::value_error(error);
}
}
return MjSpec(spec);
},
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", py::return_value_policy::move);
mjSpec.def("recompile", [mjmodel_mjdata_from_spec_ptr](
const MjSpec& self, py::object m, py::object d) {
return mjmodel_mjdata_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr),
@@ -225,79 +317,6 @@ PYBIND11_MODULE(_specs, m) {
}
return std::string(buf.get());
});
mjSpec.def(
"from_file",
[](MjSpec& self, std::string& filename,
std::optional<std::unordered_map<std::string, py::bytes>>& assets)
-> void {
const auto converted_assets = _impl::ConvertAssetsDict(assets);
{
py::gil_scoped_release no_gil;
char error[1024];
mj_deleteSpec(self.ptr);
self.ptr = LoadSpecFileImpl(
filename, converted_assets,
[&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error));
});
if (!self.ptr) {
throw py::value_error(error);
}
}
},
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<std::unordered_map<std::string, py::bytes>>& assets)
-> void {
auto converted_assets = _impl::ConvertAssetsDict(assets);
{
py::gil_scoped_release no_gil;
std::string model_filename = "model_.xml";
if (assets.has_value()) {
while (assets->find(model_filename) != assets->end()) {
model_filename =
model_filename.substr(0, model_filename.size() - 4) + "_.xml";
}
}
converted_assets.emplace_back(
model_filename.c_str(), xml.c_str(), xml.length());
char error[1024];
mj_deleteSpec(self.ptr);
self.ptr = LoadSpecFileImpl(
model_filename, converted_assets,
[&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error));
});
if (!self.ptr) {
throw py::value_error(error);
}
}
},
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
+10 -18
View File
@@ -417,8 +417,7 @@ class SpecsTest(absltest.TestCase):
state_type = mujoco.mjtState.mjSTATE_INTEGRATION
# Load from file.
spec1 = mujoco.MjSpec()
spec1.from_file(filename)
spec1 = mujoco.MjSpec.from_file(filename)
model1 = spec1.compile()
data1 = mujoco.MjData(model1)
mujoco.mj_step(model1, data1)
@@ -427,9 +426,8 @@ class SpecsTest(absltest.TestCase):
mujoco.mj_getState(model1, data1, state1, state_type)
# Load from string.
spec2 = mujoco.MjSpec()
with open(filename, 'r') as file:
spec2.from_string(file.read().rstrip())
spec2 = mujoco.MjSpec.from_string(file.read().rstrip())
model2 = spec2.compile()
data2 = mujoco.MjData(model2)
mujoco.mj_step(model2, data2)
@@ -621,8 +619,7 @@ class SpecsTest(absltest.TestCase):
</worldbody>
</mujoco>
"""
spec = mujoco.MjSpec()
spec.from_string(textwrap.dedent("""
spec = mujoco.MjSpec.from_string(textwrap.dedent("""
<mujoco model="MuJoCo Model">
<include file="included.xml"/>
</mujoco>
@@ -633,8 +630,7 @@ class SpecsTest(absltest.TestCase):
def test_delete(self):
filename = '../../test/testdata/model.xml'
spec = mujoco.MjSpec()
spec.from_file(filename)
spec = mujoco.MjSpec.from_file(filename)
model = spec.compile()
self.assertIsNotNone(model)
@@ -664,8 +660,7 @@ class SpecsTest(absltest.TestCase):
</mujoco>
"""
spec = mujoco.MjSpec()
spec.from_string(xml)
spec = mujoco.MjSpec.from_string(xml)
self.assertIsNotNone(spec.worldbody)
body = spec.worldbody.add_body()
@@ -696,8 +691,7 @@ class SpecsTest(absltest.TestCase):
</mujoco>
"""
spec = mujoco.MjSpec()
spec.from_string(main_xml)
spec = mujoco.MjSpec.from_string(main_xml)
model = spec.compile()
data = mujoco.MjData(model)
@@ -710,8 +704,7 @@ class SpecsTest(absltest.TestCase):
spec.recompile(model, data)
def test_delete_unused_plugin(self):
spec = mujoco.MjSpec()
spec.from_string(textwrap.dedent("""
spec = mujoco.MjSpec.from_string("""
<mujoco model="MuJoCo Model">
<extension>
<plugin plugin="mujoco.pid">
@@ -727,7 +720,7 @@ class SpecsTest(absltest.TestCase):
</body>
</worldbody>
</mujoco>
"""))
""")
plugin = spec.plugins[0]
self.assertIsNotNone(plugin)
plugin.delete()
@@ -737,8 +730,7 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(model.nplugin, 0)
def test_access_option_stat_visual(self):
spec = mujoco.MjSpec()
spec.from_string(textwrap.dedent("""
spec = mujoco.MjSpec.from_string("""
<mujoco model="MuJoCo Model">
<option timestep="0.001"/>
<statistic meansize="0.05"/>
@@ -746,7 +738,7 @@ class SpecsTest(absltest.TestCase):
<quality shadowsize="4096"/>
</visual>
</mujoco>
"""))
""")
self.assertEqual(spec.option.timestep, 0.001)
self.assertEqual(spec.stat.meansize, 0.05)
self.assertEqual(spec.visual.quality.shadowsize, 4096)