Change keyword argument asset to include in spec.from_file and spec.from_string.

PiperOrigin-RevId: 720611881
Change-Id: I92f5e72cece8d3223bc552ea4f2ef742c47cb8c9
This commit is contained in:
Alessio Quaglino
2025-01-28 09:54:48 -08:00
committed by Copybara-Service
parent f4096bcab9
commit 5b330ca19c
2 changed files with 24 additions and 33 deletions
+23 -32
View File
@@ -71,13 +71,7 @@ using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
struct MjSpec { struct MjSpec {
MjSpec() : ptr(mj_makeSpec()) {} MjSpec() : ptr(mj_makeSpec()) {}
MjSpec(raw::MjSpec* ptr, MjSpec(raw::MjSpec* ptr) : ptr(ptr) {}
const std::unordered_map<std::string, py::bytes>& assets_ = {})
: ptr(ptr) {
for (const auto& asset : assets_) {
assets[asset.first.c_str()] = asset.second;
}
}
// copy constructor and assignment // copy constructor and assignment
MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) { MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {
@@ -269,15 +263,15 @@ PYBIND11_MODULE(_specs, m) {
mjSpec.def_static( mjSpec.def_static(
"from_file", "from_file",
[](std::string& filename, [](std::string& filename,
std::optional<std::unordered_map<std::string, py::bytes>>& assets) std::optional<std::unordered_map<std::string, py::bytes>>& include)
-> MjSpec { -> MjSpec {
const auto converted_assets = _impl::ConvertAssetsDict(assets); const auto files = _impl::ConvertAssetsDict(include);
raw::MjSpec* spec; raw::MjSpec* spec;
{ {
py::gil_scoped_release no_gil; py::gil_scoped_release no_gil;
char error[1024]; char error[1024];
spec = LoadSpecFileImpl( spec = LoadSpecFileImpl(
filename, converted_assets, filename, files,
[&error](const char* filename, const mjVFS* vfs) { [&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)( return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error)); filename, vfs, error, sizeof(error));
@@ -286,43 +280,42 @@ PYBIND11_MODULE(_specs, m) {
throw py::value_error(error); throw py::value_error(error);
} }
} }
if (assets.has_value()) {
return MjSpec(spec, assets.value());
}
return MjSpec(spec); return MjSpec(spec);
}, },
py::arg("filename"), py::arg("assets") = py::none(), R"mydelimiter( py::arg("filename"), py::arg("include") = py::none(), R"mydelimiter(
Creates a spec from an XML file. Creates a spec from an XML file.
Parameters Parameters
---------- ----------
filename : str filename : str
Path to the XML file. Path to the XML file.
assets : dict, optional include : dict, optional
A dictionary of assets to be used by the spec. The keys are asset names A dictionary of xml files included by the model. The keys are file names
and the values are asset contents. and the values are file contents.
)mydelimiter", py::return_value_policy::move); )mydelimiter",
py::return_value_policy::move);
mjSpec.def_static( mjSpec.def_static(
"from_string", "from_string",
[](std::string& xml, [](std::string& xml,
std::optional<std::unordered_map<std::string, py::bytes>>& assets) std::optional<std::unordered_map<std::string, py::bytes>>& include)
-> MjSpec { -> MjSpec {
auto converted_assets = _impl::ConvertAssetsDict(assets); auto files = _impl::ConvertAssetsDict(include);
raw::MjSpec* spec; raw::MjSpec* spec;
{ {
py::gil_scoped_release no_gil; py::gil_scoped_release no_gil;
std::string model_filename = "model_.xml"; std::string model_filename = "model_.xml";
if (assets.has_value()) { if (include.has_value()) {
while (assets->find(model_filename) != assets->end()) { while (include->find(model_filename) !=
include->end()) {
model_filename = model_filename =
model_filename.substr(0, model_filename.size() - 4) + "_.xml"; model_filename.substr(0, model_filename.size() - 4) + "_.xml";
} }
} }
converted_assets.emplace_back( files.emplace_back(
model_filename.c_str(), xml.c_str(), xml.length()); model_filename.c_str(), xml.c_str(), xml.length());
char error[1024]; char error[1024];
spec = LoadSpecFileImpl( spec = LoadSpecFileImpl(
model_filename, converted_assets, model_filename, files,
[&error](const char* filename, const mjVFS* vfs) { [&error](const char* filename, const mjVFS* vfs) {
return InterceptMjErrors(mj_parseXML)( return InterceptMjErrors(mj_parseXML)(
filename, vfs, error, sizeof(error)); filename, vfs, error, sizeof(error));
@@ -331,22 +324,20 @@ PYBIND11_MODULE(_specs, m) {
throw py::value_error(error); throw py::value_error(error);
} }
} }
if (assets.has_value()) {
return MjSpec(spec, assets.value());
}
return MjSpec(spec); return MjSpec(spec);
}, },
py::arg("xml"), py::arg("assets") = py::none(), R"mydelimiter( py::arg("xml"), py::arg("include") = py::none(), R"mydelimiter(
Creates a spec from an XML string. Creates a spec from an XML string.
Parameters Parameters
---------- ----------
xml : str xml : str
XML string. XML string.
assets : dict, optional include : dict, optional
A dictionary of assets to be used by the spec. The keys are asset names A dictionary of xml files included by the model. The keys are file names
and the values are asset contents. and the values are file contents.
)mydelimiter", py::return_value_policy::move); )mydelimiter",
py::return_value_policy::move);
mjSpec.def("recompile", [mjmodel_mjdata_from_spec_ptr]( mjSpec.def("recompile", [mjmodel_mjdata_from_spec_ptr](
const MjSpec& self, py::object m, py::object d) { const MjSpec& self, py::object m, py::object d) {
return mjmodel_mjdata_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr), return mjmodel_mjdata_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr),
+1 -1
View File
@@ -765,7 +765,7 @@ class SpecsTest(absltest.TestCase):
<include file="included.xml"/> <include file="included.xml"/>
</mujoco> </mujoco>
"""), """),
{'included.xml': included_xml.encode('utf-8')}, include={'included.xml': included_xml.encode('utf-8')},
) )
self.assertEqual( self.assertEqual(
spec.worldbody.first_body().first_geom().type, mujoco.mjtGeom.mjGEOM_BOX spec.worldbody.first_body().first_geom().type, mujoco.mjtGeom.mjGEOM_BOX