Euler/Advance (#19)

- Adds euler integrator and advance function
- Adds internal API for solve/factor m with explicit arguments

Author: adenzler-nvidia

mujoco/mjx/__init__.py

@@ -1,5 +1,21 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 """Public API for MJX."""
 
+from ._src.forward import euler
 from ._src.forward import forward
 from ._src.forward import fwd_acceleration
 from ._src.forward import fwd_position
@@ -12,9 +28,9 @@
 from ._src.smooth import com_vel
 from ._src.smooth import crb
 from ._src.smooth import factor_m
-from ._src.smooth import solve_m
 from ._src.smooth import kinematics
 from ._src.smooth import rne
+from ._src.smooth import solve_m
 from ._src.support import is_sparse
 from ._src.test_util import benchmark
 from ._src.types import *

mujoco/mjx/_src/forward.py

@@ -1,9 +1,210 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional
+
 import warp as wp
+
+from . import math
 from . import passive
 from . import smooth
 
+from .types import array2df
 from .types import Model
 from .types import Data
+from .types import MJ_MINVAL
+from .types import MJ_DSBL_EULERDAMP
+
+
+def _advance(
+  m: Model,
+  d: Data,
+  act_dot: wp.array,
+  qacc: wp.array,
+  qvel: Optional[wp.array] = None,
+) -> Data:
+  """Advance state and time given activation derivatives and acceleration."""
+
+  # TODO(team): can we assume static timesteps?
+
+  @wp.kernel
+  def next_activation(
+    m: Model,
+    d: Data,
+    act_dot_in: array2df,
+  ):
+    worldId, actid = wp.tid()
+
+    # get the high/low range for each actuator state
+    limited = m.actuator_actlimited[actid]
+    range_low = wp.select(limited, -wp.inf, m.actuator_actrange[actid, 0])
+    range_high = wp.select(limited, wp.inf, m.actuator_actrange[actid, 1])
+
+    # get the actual actuation - skip if -1 (means stateless actuator)
+    act_adr = m.actuator_actadr[actid]
+    if act_adr == -1:
+      return
+
+    acts = d.act[worldId]
+    acts_dot = act_dot_in[worldId]
+
+    act = acts[act_adr]
+    act_dot = acts_dot[act_adr]
+
+    # check dynType
+    dyn_type = m.actuator_dyntype[actid]
+    dyn_prm = m.actuator_dynprm[actid, 0]
+
+    # advance the actuation
+    if dyn_type == 3:  # wp.static(WarpDynType.FILTEREXACT):
+      tau = wp.select(dyn_prm < MJ_MINVAL, dyn_prm, MJ_MINVAL)
+      act = act + act_dot * tau * (1.0 - wp.exp(-m.opt.timestep / tau))
+    else:
+      act = act + act_dot * m.opt.timestep
+
+    # apply limits
+    wp.clamp(act, range_low, range_high)
+
+    acts[act_adr] = act
+
+  @wp.kernel
+  def advance_velocities(m: Model, d: Data, qacc: array2df):
+    worldId, tid = wp.tid()
+    d.qvel[worldId, tid] = d.qvel[worldId, tid] + qacc[worldId, tid] * m.opt.timestep
+
+  @wp.kernel
+  def integrate_joint_positions(m: Model, d: Data, qvel_in: array2df):
+    worldId, jntid = wp.tid()
+
+    jnt_type = m.jnt_type[jntid]
+    qpos_adr = m.jnt_qposadr[jntid]
+    dof_adr = m.jnt_dofadr[jntid]
+    qpos = d.qpos[worldId]
+    qvel = qvel_in[worldId]
+
+    if jnt_type == 0:  # free joint
+      qpos_pos = wp.vec3(qpos[qpos_adr], qpos[qpos_adr + 1], qpos[qpos_adr + 2])
+      qvel_lin = wp.vec3(qvel[dof_adr], qvel[dof_adr + 1], qvel[dof_adr + 2])
+
+      qpos_new = qpos_pos + m.opt.timestep * qvel_lin
+
+      qpos_quat = wp.quat(
+        qpos[qpos_adr + 3],
+        qpos[qpos_adr + 4],
+        qpos[qpos_adr + 5],
+        qpos[qpos_adr + 6],
+      )
+      qvel_ang = wp.vec3(qvel[dof_adr + 3], qvel[dof_adr + 4], qvel[dof_adr + 5])
+
+      qpos_quat_new = math.quat_integrate(qpos_quat, qvel_ang, m.opt.timestep)
+
+      qpos[qpos_adr] = qpos_new[0]
+      qpos[qpos_adr + 1] = qpos_new[1]
+      qpos[qpos_adr + 2] = qpos_new[2]
+      qpos[qpos_adr + 3] = qpos_quat_new[0]
+      qpos[qpos_adr + 4] = qpos_quat_new[1]
+      qpos[qpos_adr + 5] = qpos_quat_new[2]
+      qpos[qpos_adr + 6] = qpos_quat_new[3]
+
+    elif jnt_type == 1:  # ball joint
+      qpos_quat = wp.quat(
+        qpos[qpos_adr],
+        qpos[qpos_adr + 1],
+        qpos[qpos_adr + 2],
+        qpos[qpos_adr + 3],
+      )
+      qvel_ang = wp.vec3(qvel[dof_adr], qvel[dof_adr + 1], qvel[dof_adr + 2])
+
+      qpos_quat_new = math.quat_integrate(qpos_quat, qvel_ang, m.opt.timestep)
+
+      qpos[qpos_adr] = qpos_quat_new[0]
+      qpos[qpos_adr + 1] = qpos_quat_new[1]
+      qpos[qpos_adr + 2] = qpos_quat_new[2]
+      qpos[qpos_adr + 3] = qpos_quat_new[3]
+
+    else:  # if jnt_type in (JointType.HINGE, JointType.SLIDE):
+      qpos[qpos_adr] = qpos[qpos_adr] + m.opt.timestep * qvel[dof_adr]
+
+  # skip if no stateful actuators.
+  if m.na:
+    wp.launch(next_activation, dim=(d.nworld, m.nu), inputs=[m, d, act_dot])
+
+  wp.launch(advance_velocities, dim=(d.nworld, m.nv), inputs=[m, d, qacc])
+
+  # advance positions with qvel if given, d.qvel otherwise (semi-implicit)
+  if qvel is not None:
+    qvel_in = qvel
+  else:
+    qvel_in = d.qvel
+
+  wp.launch(integrate_joint_positions, dim=(d.nworld, m.njnt), inputs=[m, d, qvel_in])
+
+  d.time = d.time + m.opt.timestep
+  return d
+
+
+def euler(m: Model, d: Data) -> Data:
+  """Euler integrator, semi-implicit in velocity."""
+  # integrate damping implicitly
+
+  def add_damping_sum_qfrc(m: Model, d: Data, is_sparse: bool):
+    @wp.kernel
+    def add_damping_sum_qfrc_kernel_sparse(m: Model, d: Data):
+      worldId, tid = wp.tid()
+
+      dof_Madr = m.dof_Madr[tid]
+      d.qM_integration[worldId, 0, dof_Madr] += m.opt.timestep * m.dof_damping[dof_Madr]
+
+      d.qfrc_integration[worldId, tid] = (
+        d.qfrc_smooth[worldId, tid] + d.qfrc_constraint[worldId, tid]
+      )
+
+    @wp.kernel
+    def add_damping_sum_qfrc_kernel_dense(m: Model, d: Data):
+      worldid, i, j = wp.tid()
+
+      damping = wp.select(i == j, 0.0, m.opt.timestep * m.dof_damping[i])
+      d.qM_integration[worldid, i, j] = d.qM[worldid, i, j] + damping
+
+      if i == 0:
+        d.qfrc_integration[worldid, j] = (
+          d.qfrc_smooth[worldid, j] + d.qfrc_constraint[worldid, j]
+        )
+
+    if is_sparse:
+      wp.copy(d.qM_integration, d.qM)
+      wp.launch(add_damping_sum_qfrc_kernel_sparse, dim=(d.nworld, m.nv), inputs=[m, d])
+    else:
+      wp.launch(
+        add_damping_sum_qfrc_kernel_dense, dim=(d.nworld, m.nv, m.nv), inputs=[m, d]
+      )
+
+  if not m.opt.disableflags & MJ_DSBL_EULERDAMP:
+    add_damping_sum_qfrc(m, d, m.opt.is_sparse)
+    smooth.factor_i(m, d, d.qM_integration, d.qLD_integration, d.qLDiagInv_integration)
+    smooth.solve_LD(
+      m,
+      d,
+      d.qLD_integration,
+      d.qLDiagInv_integration,
+      d.qacc_integration,
+      d.qfrc_integration,
+    )
+    return _advance(m, d, d.act_dot, d.qacc_integration)
+
+  return _advance(m, d, d.act_dot, d.qacc)
 
 
 def fwd_position(m: Model, d: Data):

mujoco/mjx/_src/forward_test.py

@@ -17,11 +17,12 @@
 
 from absl.testing import absltest
 from etils import epath
-import mujoco
-from mujoco import mjx
 import numpy as np
 import warp as wp
 
+import mujoco
+from mujoco import mjx
+
 # tolerance for difference between MuJoCo and MJX smooth calculations - mostly
 # due to float precision
 _TOLERANCE = 5e-5
@@ -71,6 +72,58 @@ def test_fwd_acceleration(self):
     _assert_eq(d.qfrc_smooth.numpy()[0], mjd.qfrc_smooth, "qfrc_smooth")
     _assert_eq(d.qacc_smooth.numpy()[0], mjd.qacc_smooth, "qacc_smooth")
 
+  def test_eulerdamp(self):
+    path = epath.resource_path("mujoco.mjx") / "test_data/pendula.xml"
+    mjm = mujoco.MjModel.from_xml_path(path.as_posix())
+    self.assertTrue((mjm.dof_damping > 0).any())
+
+    mjd = mujoco.MjData(mjm)
+    mjd.qvel[:] = 1.0
+    mjd.qacc[:] = 1.0
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjx.put_model(mjm)
+    d = mjx.put_data(mjm, mjd)
+
+    mjx.euler(m, d)
+    mujoco.mj_Euler(mjm, mjd)
+
+    _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
+    _assert_eq(d.act.numpy()[0], mjd.act, "act")
+
+    # also test sparse
+    mjm.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    mjd = mujoco.MjData(mjm)
+    mjd.qvel[:] = 1.0
+    mjd.qacc[:] = 1.0
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjx.put_model(mjm)
+    d = mjx.put_data(mjm, mjd)
+
+    mjx.euler(m, d)
+    mujoco.mj_Euler(mjm, mjd)
+
+    _assert_eq(d.qpos.numpy()[0], mjd.qpos, "qpos")
+    _assert_eq(d.act.numpy()[0], mjd.act, "act")
+
+  def test_disable_eulerdamp(self):
+    path = epath.resource_path("mujoco.mjx") / "test_data/pendula.xml"
+    mjm = mujoco.MjModel.from_xml_path(path.as_posix())
+    mjm.opt.disableflags = mjm.opt.disableflags | mujoco.mjtDisableBit.mjDSBL_EULERDAMP
+
+    mjd = mujoco.MjData(mjm)
+    mujoco.mj_forward(mjm, mjd)
+    mjd.qvel[:] = 1.0
+    mjd.qacc[:] = 1.0
+
+    m = mjx.put_model(mjm)
+    d = mjx.put_data(mjm, mjd)
+
+    mjx.euler(m, d)
+
+    np.testing.assert_allclose(d.qvel.numpy()[0], 1 + mjm.opt.timestep)
+
 
 if __name__ == "__main__":
   wp.init()