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

Add more fields to mjx.Option and mjx.Statistic for MuJoCo compatibility.

Browse Source 
Also add metadata fields to the Python bindings of mjModel and mjOption,
and fix dtype of a number of fields.

PiperOrigin-RevId: 654138876
Change-Id: If2bbe41bf3dfdb350d60b5289afc65fffb397ce9
This commit is contained in:
Saran Tunyasuvunakool
2024-07-19 15:11:57 -07:00
committed by Copybara-Service
parent b3f79e561d
commit db1d51a1f5
3 changed files with 129 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/structs.cc
+50
View File
@@ -1383,6 +1383,30 @@ PYBIND11_MODULE(_structs, m) {
MJOPTION_VECTORS
#undef X
mjOption.def_property_readonly_static("_float_fields", [](py::object) {
std::vector<std::string> field_names;
#define X(type, var) field_names.push_back(#var);
MJOPTION_FLOATS
#undef X
return py::tuple(py::cast(field_names));
});
mjOption.def_property_readonly_static("_int_fields", [](py::object) {
std::vector<std::string> field_names;
#define X(type, var) field_names.push_back(#var);
MJOPTION_INTS
#undef X
return py::tuple(py::cast(field_names));
});
mjOption.def_property_readonly_static("_floatarray_fields", [](py::object) {
std::vector<std::string> field_names;
#define X(var, sz) field_names.push_back(#var);
MJOPTION_VECTORS
#undef X
return py::tuple(py::cast(field_names));
});
// ==================== MJVISUAL =============================================
py::class_<MjVisualWrapper> mjVisual(m, "MjVisual");
mjVisual.def("__copy__", [](const MjVisualWrapper& other) {
@@ -1628,6 +1652,32 @@ This is useful for example when the MJB is not available as a file on disk.)"));
MJMODEL_INTS
#undef X
mjModel.def_property_readonly("_sizes", [](const MjModelWrapper& m) {
int nint = 0;
#define X(var) ++nint;
MJMODEL_INTS
#undef X
py::array_t<std::int64_t> sizes(nint);
{
int i = 0;
auto data = sizes.mutable_unchecked();
#define X(var) data[i++] = m.get()->var;
MJMODEL_INTS
#undef X
}
py::detail::array_proxy(sizes.ptr())->flags &=
~py::detail::npy_api::NPY_ARRAY_WRITEABLE_;
return sizes;
});
mjModel.def_property_readonly_static("_size_fields", [](py::object) {
std::vector<std::string> fields;
#define X(var) fields.push_back(#var);
MJMODEL_INTS
#undef X
return py::tuple(py::cast(fields));
});
#define X(dtype, var, dim0, dim1) \
if constexpr (std::string_view(#var) != "text_data" && \
std::string_view(#var) != "names" && \