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
This commit is contained in:
Saran Tunyasuvunakool
committed by
Copybara-Service
Copybara-Service
parent
b3f79e561d
commit
db1d51a1f5
+22
-16
@@ -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),
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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" && \
|
||||
|
||||
Reference in New Issue
Block a user