Allow for different MJX backend implementations.

PiperOrigin-RevId: 755935704
Change-Id: Ic135cd00137c2857c73c683ed9fdc5ac4418715d

mjx/mujoco/mjx/_src/forward.py

@@ -33,12 +33,14 @@ from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import BiasType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import DynType
 from mujoco.mjx._src.types import GainType
 from mujoco.mjx._src.types import IntegratorType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -81,10 +83,13 @@ def fwd_position(m: Model, d: Data) -> Data:
 @named_scope
 def fwd_velocity(m: Model, d: Data) -> Data:
   """Velocity-dependent computations."""
-  d = d.replace(
-      actuator_velocity=d.actuator_moment @ d.qvel,
-      ten_velocity=d.ten_J @ d.qvel,
-  )
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('fwd_velocity requires JAX backend implementation.')
+
+  d = d.tree_replace({
+      '_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
+      '_impl.ten_velocity': d._impl.ten_J @ d.qvel,
+  })
   d = smooth.com_vel(m, d)
   d = passive.passive(m, d)
   d = smooth.rne(m, d)
@@ -173,8 +178,8 @@ def fwd_actuation(m: Model, d: Data) -> Data:
       m.actuator_gainprm,
       m.actuator_biastype,
       m.actuator_biasprm,
-      d.actuator_length,
-      d.actuator_velocity,
+      d._impl.actuator_length,
+      d._impl.actuator_velocity,
       ctrl_act,
       jp.array(m.actuator_lengthrange),
       jp.array(m.actuator_acc0),
@@ -215,7 +220,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
   )
   force = jp.clip(force, forcerange[:, 0], forcerange[:, 1])
 
-  qfrc_actuator = d.actuator_moment.T @ force
+  qfrc_actuator = d._impl.actuator_moment.T @ force
 
   if m.ngravcomp:
     # actuator-level gravity compensation, skip if added as passive force
@@ -331,13 +336,17 @@ def _advance(
 @named_scope
 def euler(m: Model, d: Data) -> Data:
   """Euler integrator, semi-implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('euler requires JAX backend implementation.')
+
   # integrate damping implicitly
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     if support.is_sparse(m):
-      dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
     else:
-      dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
+    dh = d.tree_replace({'_impl.qM': qM})
     dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
@@ -392,13 +401,15 @@ def rungekutta4(m: Model, d: Data) -> Data:
 @named_scope
 def implicit(m: Model, d: Data) -> Data:
   """Integrates fully implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('implicit requires JAX backend implementation.')
 
   qderiv = derivative.deriv_smooth_vel(m, d)
 
   qacc = d.qacc
   if qderiv is not None:
     # TODO(robotics-simulation): use smooth.factor_m / solve_m here:
-    qm = support.full_m(m, d) if support.is_sparse(m) else d.qM
+    qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
     qm -= m.opt.timestep * qderiv
     qh, _ = jax.scipy.linalg.cho_factor(qm)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
@@ -410,6 +421,9 @@ def implicit(m: Model, d: Data) -> Data:
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('forward requires JAX backend implementation.')
+
   d = fwd_position(m, d)
   d = sensor.sensor_pos(m, d)
   d = fwd_velocity(m, d)
@@ -418,7 +432,7 @@ def forward(m: Model, d: Data) -> Data:
   d = fwd_acceleration(m, d)
   d = sensor.sensor_acc(m, d)
 
-  if d.efc_J.size == 0:
+  if d._impl.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)
     return d