Internal change.

PiperOrigin-RevId: 704273388
Change-Id: I6de2aa5bf80353dae4f2faf408c4b4c1c2e14f15

mjx/mujoco/mjx/_src/io.py

@@ -37,8 +37,11 @@ def _strip_weak_type(tree):
   return jax.tree_util.tree_map(f, tree)
 
 
-def _make_option(o: mujoco.MjOption) -> types.Option:
+def _make_option(
+    o: mujoco.MjOption, _full_compat: bool = False
+) -> types.Option:
   """Returns mjx.Option given mujoco.MjOption."""
+  if not _full_compat:
     if o.integrator not in set(types.IntegratorType):
       raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
 
@@ -57,6 +60,7 @@ def _make_option(o: mujoco.MjOption) -> types.Option:
 
   has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
   implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST
+  if not _full_compat:
     if implicitfast and has_fluid_params:
       raise NotImplementedError('implicitfast not implemented for fluid drag.')
 
@@ -182,11 +186,12 @@ def put_model(
   fields['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0
   fields['geom_rbound_hfield'] = fields['geom_rbound']
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
-  fields['opt'] = _make_option(m.opt)
+  fields['opt'] = _make_option(m.opt, _full_compat=_full_compat)
   fields['stat'] = _make_statistic(m.stat)
 
   # Pre-compile meshes for MJX collisions.
   fields['mesh_convex'] = [None] * m.nmesh
+  if not _full_compat:
     for i in mesh_geomid:
       dataid = m.geom_dataid[i]
       if fields['mesh_convex'][dataid] is None:
@@ -496,8 +501,9 @@ def get_data_into(
         value = np.ones(m.nv)
 
       if isinstance(value, np.ndarray) and value.shape:
-        if restricted_to in ('mujoco', 'mjx') and value.shape == (0,):
+        if restricted_to in ('mujoco', 'mjx'):
           continue  # don't copy fields that are mujoco-only or MJX-only
+        else:
           getattr(result_i, field.name)[:] = value
       else:
         setattr(result_i, field.name, value)