From d49d4bf64cd83b9d0c8e5cddde4b0d3b614f25de Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Fri, 19 Jul 2024 15:22:40 -0700 Subject: [PATCH] Make unsupported feature checking optional in put_model. PiperOrigin-RevId: 654141976 Change-Id: I35b5f1e4c5b07e32075e847d4a757698a0b603c1 --- mjx/mujoco/mjx/_src/io.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 373e9b0b..b9657c4f 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -69,17 +69,19 @@ def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: ) -def put_model(m: mujoco.MjModel, device=None) -> types.Model: +def put_model( + m: mujoco.MjModel, device=None, _check_unsupported=True +) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" - if m.ntendon: + if _check_unsupported and m.ntendon: raise NotImplementedError('tendons are not supported') mesh_geomid = set() for g1, g2, ip in collision_driver.geom_pairs(m): t1, t2 = m.geom_type[[g1, g2]] # check collision function exists for type pair - if not collision_driver.has_collision_fn(t1, t2): + if _check_unsupported and not collision_driver.has_collision_fn(t1, t2): t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') # margin/gap not supported for meshes and height fields @@ -89,7 +91,7 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: margin = m.pair_margin[ip] else: margin = m.geom_margin[g1] + m.geom_margin[g2] - if margin.any(): + if _check_unsupported and margin.any(): t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') for t, g in [(t1, g1), (t2, g2)]: @@ -104,12 +106,12 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: (m.eq_type, types.EqType, mujoco.mjtEq), ): missing = set(enum_field) - set(enum_type) - if missing: + if _check_unsupported and missing: raise NotImplementedError( f'{[mj_type(m) for m in missing]} not supported' ) - if not np.allclose(m.dof_frictionloss, 0): + if _check_unsupported and not np.allclose(m.dof_frictionloss, 0): raise NotImplementedError('dof_frictionloss is not implemented.') mjx_only = {'mesh_convex', 'geom_rbound_hfield'}