From db1d51a1f5a6cd3740ad79cd6db5b502a2cd6029 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Fri, 19 Jul 2024 15:11:57 -0700 Subject: [PATCH] Add more fields to mjx.Option and mjx.Statistic for MuJoCo compatibility. 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 --- mjx/mujoco/mjx/_src/io.py | 38 ++++++++++------- mjx/mujoco/mjx/_src/types.py | 82 +++++++++++++++++++++++++----------- python/mujoco/structs.cc | 50 ++++++++++++++++++++++ 3 files changed, 129 insertions(+), 41 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 1cdb015f..373e9b0b 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -60,7 +60,13 @@ def _make_option(o: mujoco.MjOption) -> types.Option: def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" - return types.Statistic(meaninertia=s.meaninertia) + return types.Statistic( + meaninertia=s.meaninertia, + meanmass=s.meanmass, + meansize=s.meansize, + extent=s.extent, + center=s.center, + ) def put_model(m: mujoco.MjModel, device=None) -> types.Model: @@ -193,19 +199,19 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: cinert=jp.zeros((m.nbody, 10), dtype=float), flexvert_xpos=jp.zeros((m.nflexvert, 3), dtype=float), flexelem_aabb=jp.zeros((m.nflexelem, 6), dtype=float), - flexedge_J_rownnz=jp.zeros((m.nflexedge,), dtype=int), - flexedge_J_rowadr=jp.zeros((m.nflexedge,), dtype=int), - flexedge_J_colind=jp.zeros((m.nflexedge, m.nv), dtype=int), + flexedge_J_rownnz=jp.zeros((m.nflexedge,), dtype=jp.int32), + flexedge_J_rowadr=jp.zeros((m.nflexedge,), dtype=jp.int32), + flexedge_J_colind=jp.zeros((m.nflexedge, m.nv), dtype=jp.int32), flexedge_J=jp.zeros((m.nflexedge, m.nv), dtype=float), flexedge_length=jp.zeros((m.nflexedge,), dtype=float), - ten_wrapadr=jp.zeros((m.ntendon,), dtype=int), - ten_wrapnum=jp.zeros((m.ntendon,), dtype=int), - ten_J_rownnz=jp.zeros((m.ntendon,), dtype=int), - ten_J_rowadr=jp.zeros((m.ntendon,), dtype=int), - ten_J_colind=jp.zeros((m.ntendon, m.nv), dtype=int), + ten_wrapadr=jp.zeros((m.ntendon,), dtype=jp.int32), + ten_wrapnum=jp.zeros((m.ntendon,), dtype=jp.int32), + ten_J_rownnz=jp.zeros((m.ntendon,), dtype=jp.int32), + ten_J_rowadr=jp.zeros((m.ntendon,), dtype=jp.int32), + ten_J_colind=jp.zeros((m.ntendon, m.nv), dtype=jp.int32), ten_J=jp.zeros((m.ntendon, m.nv), dtype=float), ten_length=jp.zeros((m.ntendon,), dtype=float), - wrap_obj=jp.zeros((m.nwrap, 2), dtype=int), + wrap_obj=jp.zeros((m.nwrap, 2), dtype=jp.int32), wrap_xpos=jp.zeros((m.nwrap, 6), dtype=float), actuator_length=jp.zeros((m.nu,), dtype=float), actuator_moment=jp.zeros((m.nu, m.nv), dtype=float), @@ -242,12 +248,12 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: subtree_angmom=jp.zeros((m.nbody, 3), dtype=float), qH=jp.zeros((m.nM,), dtype=float), qHDiagInv=jp.zeros((m.nv,), dtype=float), - D_rownnz=jp.zeros((m.nv,), dtype=int), - D_rowadr=jp.zeros((m.nv,), dtype=int), - D_colind=jp.zeros((m.nD,), dtype=int), - B_rownnz=jp.zeros((m.nbody,), dtype=int), - B_rowadr=jp.zeros((m.nbody,), dtype=int), - B_colind=jp.zeros((m.nB,), dtype=int), + D_rownnz=jp.zeros((m.nv,), dtype=jp.int32), + D_rowadr=jp.zeros((m.nv,), dtype=jp.int32), + D_colind=jp.zeros((m.nD,), dtype=jp.int32), + B_rownnz=jp.zeros((m.nbody,), dtype=jp.int32), + B_rowadr=jp.zeros((m.nbody,), dtype=jp.int32), + B_colind=jp.zeros((m.nB,), dtype=jp.int32), qDeriv=jp.zeros((m.nD,), dtype=float), qLU=jp.zeros((m.nD,), dtype=float), actuator_force=jp.zeros((m.nu,), dtype=float), diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index f20f1cbb..d9e3375c 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -271,48 +271,71 @@ class Option(PyTreeNode): """Physics options. Attributes: - timestep: timestep - impratio: ratio of friction-to-normal contact impedance - tolerance: main solver tolerance - ls_tolerance: CG/Newton linesearch tolerance - gravity: gravitational acceleration (3,) - wind: wind (for lift, drag and viscosity) - density: density of medium - viscosity: viscosity of medium - has_fluid_params: automatically set by mjx if wind/density/viscosity are + timestep: timestep + apirate: update rate for remote API (Hz) (not used) + impratio: ratio of friction-to-normal contact impedance + tolerance: main solver tolerance + ls_tolerance: CG/Newton linesearch tolerance + noslip_tolerance: noslip solver tolerance (not used) + mpr_tolerance: MPR solver tolerance (not used) + gravity: gravitational acceleration (3,) + wind: wind (for lift, drag and viscosity) + magnetic: global magnetic flux (not used) + density: density of medium + viscosity: viscosity of medium + o_margin: contact solver override: margin (not used) + o_solref: contact solver override: solref (not used) + o_solimp: contact solver override: solimp (not used) + o_friction[5]: contact solver override: friction (not used) + has_fluid_params: automatically set by mjx if wind/density/viscosity are nonzero. Not used by mj - integrator: integration mode - cone: type of friction cone - jacobian: matrix layout for mass matrices (dense or sparse) - (note that this is different from MuJoCo, where jacobian - specifies whether efc_J and its accompanying matrices - are dense or sparse. - solver: solver algorithm - iterations: number of main solver iterations - ls_iterations: maximum number of CG/Newton linesearch iterations - disableflags: bit flags for disabling standard features + integrator: integration mode + cone: type of friction cone + jacobian: matrix layout for mass matrices (dense or sparse) + (note that this is different from MuJoCo, where jacobian + specifies whether efc_J and its accompanying matrices + are dense or sparse. + solver: solver algorithm + iterations: number of main solver iterations + ls_iterations: maximum number of CG/Newton linesearch iterations + noslip_iterations: maximum number of noslip solver iterations (not used) + mpr_iterations: maximum number of MPR solver iterations (not used) + disableflags: bit flags for disabling standard features + enableflags: bit flags for enabling optional features (not used) + disableactuator: bit flags for disabling actuators by group id (not used) + sdf_initpoints: number of starting points for gradient descent (not used) + sdf_iterations: max number of iterations for gradient descent (not used) """ timestep: jax.Array - # unsupported: apirate + apirate: jax.Array impratio: jax.Array tolerance: jax.Array ls_tolerance: jax.Array - # unsupported: noslip_tolerance, mpr_tolerance + noslip_tolerance: jax.Array + mpr_tolerance: jax.Array gravity: jax.Array wind: jax.Array + magnetic: jax.Array density: jax.Array viscosity: jax.Array + o_margin: jax.Array + o_solref: jax.Array + o_solimp: jax.Array + o_friction: jax.Array has_fluid_params: bool - # unsupported: magnetic, o_margin, o_solref, o_solimp integrator: IntegratorType cone: ConeType jacobian: JacobianType solver: SolverType iterations: int ls_iterations: int - # unsupported: noslip_iterations, mpr_iterations + noslip_iterations: int + mpr_iterations: int disableflags: DisableBit - # unsupported: enableflags + enableflags: int + disableactuator: int + sdf_initpoints: int + sdf_iterations: int class Statistic(PyTreeNode): @@ -320,9 +343,16 @@ class Statistic(PyTreeNode): Attributes: meaninertia: mean diagonal inertia + meanmass: mean body mass (not used) + meansize: mean body size (not used) + extent: spatial extent (not used) + center: center of model (not used) """ meaninertia: jax.Array - # unsupported: meanmass, meansize, extent, center + meanmass: jax.Array + meansize: jax.Array + extent: jax.Array + center: jax.Array class Model(PyTreeNode): @@ -907,6 +937,7 @@ class Model(PyTreeNode): actuator_cranklength: np.ndarray actuator_acc0: np.ndarray actuator_lengthrange: np.ndarray + actuator_plugin: np.ndarray sensor_type: np.ndarray sensor_datatype: np.ndarray sensor_needstage: np.ndarray @@ -939,6 +970,7 @@ class Model(PyTreeNode): name_tupleadr: np.ndarray name_keyadr: np.ndarray names: bytes + _sizes: jax.Array class Contact(PyTreeNode): diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index fd8443d9..373593c7 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -1383,6 +1383,30 @@ PYBIND11_MODULE(_structs, m) { MJOPTION_VECTORS #undef X + mjOption.def_property_readonly_static("_float_fields", [](py::object) { + std::vector 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 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 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_ 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 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 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" && \