Enable Warp execution on CPU devices in MJX.

This change modifies the Warp JAX FFI to register callbacks for both CUDA and CPU platforms, allowing Warp kernels to be launched on CPU devices. MJX's io.py is updated to allow Warp to use CPU devices if no CUDA GPU is available. Tests are adjusted to no longer skip when CUDA GPUs are absent. A minor fix for sorting precision in test_util.py is also included. Kristian Hartikainen - https://github.com/google-deepmind/mujoco/pull/2948 Github issue - https://github.com/google-deepmind/mujoco/issues/2947 PiperOrigin-RevId: 906018844 Change-Id: I9c026b4a4e2a0276d5d3af3954fd027cf75b6d7a
2026-04-26 14:18:49 -07:00
parent 8f32f6e72e
commit 521d152ec5
6 changed files with 191 additions and 85 deletions
@@ -96,16 +96,15 @@ def _resolve_device(
    return cpu_0

  if impl == types.Impl.WARP:
-    # WARP implementation requires a CUDA GPU.
-    cuda_gpus = [d for d in jax.devices('cuda')]
-    if not cuda_gpus:
-      raise AssertionError(
-          'No CUDA GPU devices found in'
-          f' jax.devices("cuda")={jax.devices("cuda")}.'
-      )
+    # WARP implementation requires a CUDA GPU or CPU.
+    if has_cuda_gpu_device():
+      cuda_gpus = [d for d in jax.devices('cuda')]
+      if cuda_gpus:
+        logging.debug('Picking default device: %s', cuda_gpus[0])
+        return cuda_gpus[0]

-    logging.debug('Picking default device: %s', cuda_gpus[0])
-    return cuda_gpus[0]
+    logging.debug('Picking default device for Warp: CPU')
+    return jax.devices('cpu')[0]

  raise ValueError(f'Unsupported implementation: {impl}')

@@ -121,9 +120,12 @@ def _check_impl_device_compatibility(
  impl = types.Impl(impl)

  if impl == types.Impl.WARP:
-    if not _is_cuda_gpu_device(device):
+    is_cuda_device = _is_cuda_gpu_device(device)
+    is_cpu_device = device.platform == 'cpu'
+    if not (is_cuda_device or is_cpu_device):
      raise AssertionError(
-          f'Warp implementation requires a CUDA GPU device, got {device}.'
+          'Warp implementation requires a CUDA GPU or CPU device, got '
+          f'{device}.'
      )
    _check_warp_installed()

@@ -425,8 +427,6 @@ def _put_model_jax(
  return _strip_weak_type(model)


-
-
 def _put_model_warp(
    m: mujoco.MjModel,
    graph_mode: mjxw.types.GraphMode,
@@ -719,8 +719,6 @@ def _make_data_jax(
  return d


-
-
 def _get_nested_attr(obj: Any, attr_name: str, split: str) -> Any:
  """Returns the nested attribute from an object."""
  for part in attr_name.split(split):
@@ -1084,8 +1082,6 @@ def _put_data_jax(
  return _strip_weak_type(data)


-
-
 # TODO(josechenf): Iterate on the keepalive implementation to make it easier to
 # use before OSS.
 def _put_data_cpp(
@@ -934,8 +934,8 @@ _DEVICE_TEST_CASES = [
    ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
    ('tpu', 'jax', ('tpu', Impl.JAX)),
    # WARP backend specified.
-    ('cpu', 'warp', ('cpu', 'error')),
-    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('cpu', 'warp', ('cpu', Impl.WARP)),
+    ('gpu-notnvidia', 'warp', ('gpu', 'error')),
    ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
    ('tpu', 'warp', ('tpu', 'error')),
    # CPP backend specified.
@@ -962,10 +962,10 @@ _DEFAULT_DEVICE_TEST_CASES = [
    ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
    ('tpu', 'jax', ('tpu', Impl.JAX)),
    # WARP backend impl specified.
-    ('cpu', 'warp', ('cpu', 'error')),
-    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('cpu', 'warp', ('cpu', Impl.WARP)),
+    ('gpu-notnvidia', 'warp', ('cpu', Impl.WARP)),
    ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
-    ('tpu', 'warp', ('tpu', 'error')),
+    ('tpu', 'warp', ('cpu', Impl.WARP)),
    # CPP backend impl specified, CPU should always be available.
    ('cpu', 'cpp', ('cpu', Impl.CPP)),
    ('gpu-notnvidia', 'cpp', ('cpu', Impl.CPP)),
@@ -1140,15 +1140,6 @@ class ResolveImplAndDeviceTest(parameterized.TestCase):
    self.mock_jax_backends.side_effect = backends_side_effect

    expected_device, expected_impl = expected
-    if (
-        expected_impl == 'error'
-        and default_device_str != 'gpu-nvidia'
-        and impl_str == 'warp'
-    ):
-      with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
-        mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
-      return
-
    if expected_impl == 'error':
      with self.assertRaises(AssertionError):
        mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
@@ -210,11 +210,19 @@ class FfiKernel:
        self.input_output_aliases = input_output_aliases

        # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
-        self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
-        ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
-        ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
-        jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
+                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
+        )
+
+        self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
+        ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
+        ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
+        jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
+
+        self.callback_func_host = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="Host"))
+        ffi_ccall_address_host = ctypes.cast(self.callback_func_host, ctypes.c_void_p)
+        ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
+        jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")

    def __call__(self, *args, output_dims=None, launch_dims=None, vmap_method=None):
        num_inputs = len(args)
@@ -241,18 +249,19 @@ class FfiKernel:
                # check dtype
                if input_value.dtype != input_arg.jax_scalar_type:
                    raise TypeError(
-                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
+                            f"Invalid data type for array argument '{input_arg.name}',"
+                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                    )
                # check ndim
                if input_value.ndim != input_arg.jax_ndim:
                    raise TypeError(
-                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                    )
                # check inner dims
                for d in range(input_arg.dtype_ndim):
                    if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                        raise TypeError(
-                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                        )
            else:
                # make sure scalar is not a traced variable, should be static
@@ -328,7 +337,7 @@ class FfiKernel:

        return call(*args, launch_id=launch_id)

-    def ffi_callback(self, call_frame):
+    def ffi_callback(self, call_frame, platform="CUDA"):
        try:
            # On the first call, XLA runtime will query the API version and traits
            # metadata using the |extension| field. Let us respond to that query
@@ -340,10 +349,11 @@ class FfiKernel:
                    metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
                    metadata_ext.contents.metadata.contents.api_version.major_version = 0
                    metadata_ext.contents.metadata.contents.api_version.minor_version = 1
-                    # Turn on CUDA graphs for this handler.
-                    metadata_ext.contents.metadata.contents.traits = (
-                        XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
-                    )
+                    # Turn on CUDA graphs for this handler if on CUDA platform.
+                    if platform == "CUDA":
+                        metadata_ext.contents.metadata.contents.traits = (
+                            XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
+                        )
                    return None

            # Lock is required to prevent race conditions when callback is invoked
@@ -423,29 +433,43 @@ class FfiKernel:
                kernel_params[0] = ctypes.addressof(launch_bounds)

                # get device and stream
-                device = wp.get_cuda_device(get_device_ordinal_from_callframe(call_frame.contents))
-                stream = get_stream_from_callframe(call_frame.contents)
+                if platform == "CUDA":
+                    device = wp.get_cuda_device(get_device_ordinal_from_callframe(call_frame.contents))
+                    stream = get_stream_from_callframe(call_frame.contents)
+                else:
+                    device = wp.get_device("cpu")
+                    stream = None

                # get kernel hooks
                hooks = self.kernel.module.get_kernel_hooks(self.kernel, device)
                assert hooks.forward, "Failed to find kernel entry point"

                # launch the kernel
-                wp._src.context.runtime.core.wp_cuda_launch_kernel(
-                    device.context,
-                    hooks.forward,
-                    launch_bounds.size,
-                    0,
-                    256,
-                    hooks.forward_smem_bytes,
-                    kernel_params,
-                    stream,
-                )
+                if device.is_cuda:
+                    wp._src.context.runtime.core.wp_cuda_launch_kernel(
+                        device.context,
+                        hooks.forward,
+                        launch_bounds.size,
+                        0,
+                        256,
+                        hooks.forward_smem_bytes,
+                        kernel_params,
+                        stream,
+                    )
+                else:
+                    wp._src.context.runtime.core.wp_cpu_launch_kernel(
+                        device.context,
+                        hooks.forward,
+                        launch_bounds.size,
+                        kernel_params,
+                    )

        except Exception as e:
            print(traceback.format_exc())
            return create_ffi_error(
-                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
+                call_frame.contents.api,
+                XLA_FFI_Error_Code.UNKNOWN,
+                f"FFI callback error: {type(e).__name__}: {e}",
            )


@@ -594,10 +618,16 @@ class FfiCallable:

        # register the callback
        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
-        self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
-        ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
-        ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
-        jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")
+
+        self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
+        ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
+        ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
+        jax.ffi.register_ffi_target(self.name, ffi_capsule_cuda, platform="CUDA")
+
+        self.callback_func_host = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="Host"))
+        ffi_ccall_address_host = ctypes.cast(self.callback_func_host, ctypes.c_void_p)
+        ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
+        jax.ffi.register_ffi_target(self.name, ffi_capsule_host, platform="Host")

    def __call__(self, *args, output_dims=None, vmap_method=None):
        num_inputs = len(args)
@@ -688,8 +718,7 @@ class FfiCallable:
                except Exception:
                    # ignore unsupported devices like TPUs
                    pass
-                # we only support CUDA devices for now
-                if dev.is_cuda:
+                if dev.is_cuda or dev.is_cpu:
                    module.load(dev)

        # save call data to be retrieved by callback
@@ -698,7 +727,7 @@ class FfiCallable:
        self.call_id += 1
        return call(*args, call_id=call_id)

-    def ffi_callback(self, call_frame):
+    def ffi_callback(self, call_frame, platform="CUDA"):
        try:
            # On the first call, XLA runtime will query the API version and traits
            # metadata using the |extension| field. Let us respond to that query
@@ -710,8 +739,8 @@ class FfiCallable:
                    metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
                    metadata_ext.contents.metadata.contents.api_version.major_version = 0
                    metadata_ext.contents.metadata.contents.api_version.minor_version = 1
-                    # Turn on CUDA graphs for this handler.
-                    if self.graph_mode is GraphMode.JAX:
+                    # Turn on CUDA graphs for this handler if on CUDA platform.
+                    if self.graph_mode is GraphMode.JAX and platform == "CUDA":
                        metadata_ext.contents.metadata.contents.traits = (
                            XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
                        )
@@ -738,6 +767,35 @@ class FfiCallable:
                assert num_inputs == self.num_inputs
                assert num_outputs == self.num_outputs

+                if platform == "Host":
+                    device = wp.get_device("cpu")
+                    # reconstruct the argument list
+                    arg_list = []
+
+                    # input and in-out args
+                    for i, arg in enumerate(self.input_args):
+                        if arg.is_array:
+                            buffer = inputs[i].contents
+                            shape = collapse_batch_dims(buffer.dims[: buffer.rank - arg.dtype_ndim], arg.type.ndim)
+                            arr = wp.array(ptr=buffer.data, dtype=arg.type.dtype, shape=shape, device=device)
+                            arg_list.append(arr)
+                        else:
+                            # scalar argument, get stashed value
+                            value = call_desc.static_inputs[arg.name]
+                            arg_list.append(value)
+
+                    # pure output args (skip in-out FFI buffers)
+                    for i, arg in enumerate(self.output_args):
+                        buffer = outputs[i + self.num_in_out].contents
+                        shape = collapse_batch_dims(buffer.dims[: buffer.rank - arg.dtype_ndim], arg.type.ndim)
+                        arr = wp.array(ptr=buffer.data, dtype=arg.type.dtype, shape=shape, device=device)
+                        arg_list.append(arr)
+
+                    # call the Python function with reconstructed arguments
+                    with wp.ScopedDevice(device):
+                        self.func(*arg_list)
+                    return
+
                cuda_stream = get_stream_from_callframe(call_frame.contents)
                device_ordinal = get_device_ordinal_from_callframe(call_frame.contents)

@@ -870,8 +928,8 @@ class FfiCallable:
                    arg_list.append(arr)

                # call the Python function with reconstructed arguments
-                with wp.ScopedStream(stream, sync_enter=False):
-                    if stream.is_capturing:
+                with wp.ScopedStream(stream, sync_enter=False) if stream else wp.ScopedDevice(device):
+                    if stream and stream.is_capturing:
                        # capturing with JAX
                        with wp.ScopedCapture(external=True) as capture:
                            self.func(*arg_list)
@@ -879,7 +937,7 @@ class FfiCallable:
                        # keep a reference to the capture object to prevent required modules getting unloaded
                        call_desc.capture = capture

-                    elif self.graph_mode == GraphMode.WARP:
+                    elif self.graph_mode == GraphMode.WARP and device.is_cuda:
                        # capturing with WARP
                        with wp.ScopedCapture() as capture:
                            self.func(*arg_list)
@@ -892,7 +950,7 @@ class FfiCallable:
                        if self._graph_cache_max is not None and len(self.captures) > self._graph_cache_max:
                            self.captures.popitem(last=False)

-                    elif self.graph_mode == GraphMode.WARP_STAGED_EX:
+                    elif self.graph_mode == GraphMode.WARP_STAGED_EX and device.is_cuda:
                        # capturing with WARP using staging buffers and memcopies done outside of the graph
                        wp_memcpy_batch = wp._src.context.runtime.core.wp_memcpy_batch

@@ -935,7 +993,7 @@ class FfiCallable:
                        # TODO: we should have a way of freeing this
                        call_desc.capture = capture

-                    elif self.graph_mode == GraphMode.WARP_STAGED:
+                    elif self.graph_mode == GraphMode.WARP_STAGED and device.is_cuda:
                        # capturing with WARP using staging buffers and memcopies done inside of the graph
                        wp_cuda_graph_insert_memcpy_batch = (
                            wp._src.context.runtime.core.wp_cuda_graph_insert_memcpy_batch
@@ -1013,7 +1071,7 @@ class FfiCallable:
                        call_desc.capture = capture

                    else:
-                        # not capturing
+                        # not capturing or on CPU
                        self.func(*arg_list)

        except Exception as e:
@@ -1621,7 +1679,7 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr

    # TODO check that the name is not already registered

-    def ffi_callback(call_frame):
+    def ffi_callback(call_frame, platform="CUDA"):
        try:
            extension = call_frame.contents.extension_start
            # On the first call, XLA runtime will query the API version and traits
@@ -1633,7 +1691,7 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
                    metadata_ext = ctypes.cast(extension, ctypes.POINTER(XLA_FFI_Metadata_Extension))
                    metadata_ext.contents.metadata.contents.api_version.major_version = 0
                    metadata_ext.contents.metadata.contents.api_version.minor_version = 1
-                    if graph_compatible:
+                    if graph_compatible and platform == "CUDA":
                        # Turn on CUDA graphs for this handler.
                        metadata_ext.contents.metadata.contents.traits = (
                            XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
@@ -1666,12 +1724,17 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
        return None

    FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
-    callback_func = FFI_CCALLFUNC(ffi_callback)
+    callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: ffi_callback(call_frame, platform="CUDA"))
+    callback_func_host = FFI_CCALLFUNC(lambda call_frame: ffi_callback(call_frame, platform="Host"))
    with _FFI_REGISTRY_LOCK:
-        _FFI_CALLBACK_REGISTRY[name] = callback_func
-    ffi_ccall_address = ctypes.cast(callback_func, ctypes.c_void_p)
-    ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
-    jax.ffi.register_ffi_target(name, ffi_capsule, platform="CUDA")
+        _FFI_CALLBACK_REGISTRY[f"{name}_cuda"] = callback_func_cuda
+        _FFI_CALLBACK_REGISTRY[f"{name}_host"] = callback_func_host
+    ffi_ccall_address_cuda = ctypes.cast(callback_func_cuda, ctypes.c_void_p)
+    ffi_capsule_cuda = jax.ffi.pycapsule(ffi_ccall_address_cuda.value)
+    jax.ffi.register_ffi_target(name, ffi_capsule_cuda, platform="CUDA")
+    ffi_ccall_address_host = ctypes.cast(callback_func_host, ctypes.c_void_p)
+    ffi_capsule_host = jax.ffi.pycapsule(ffi_ccall_address_host.value)
+    jax.ffi.register_ffi_target(name, ffi_capsule_host, platform="Host")


 ###############################################################################
@@ -31,6 +31,7 @@ from mujoco.mjx.warp import test_util as tu
 from mujoco.mjx.warp import warp as wp  # pylint: disable=g-importing-member
 import numpy as np

+
 try:
  from mujoco.mjx.warp import forward  # pylint: disable=g-import-not-at-top
 except ImportError:
@@ -300,6 +301,56 @@ class StepTest(parameterized.TestCase):
      tu.assert_attr_eq(dx, d, 'mocap_quat')
      tu.assert_attr_eq(dx, d, 'sensordata')

+  @parameterized.parameters(
+      'humanoid/humanoid.xml',
+      'pendula.xml',
+  )
+  def test_step_cpu(self, xml: str):
+    """Tests step on the CPU device."""
+    if not _FORCE_TEST:
+      if not mjxw.WARP_INSTALLED:
+        self.skipTest('Warp not installed.')
+
+    batch_size = 1
+    m = test_util.load_test_file(xml)
+    m.opt.iterations = 10
+    m.opt.ls_iterations = 10
+
+    cpu_device = jax.devices('cpu')[0]
+    mx = mjx.put_model(m, impl='warp', device=cpu_device)
+
+    d = mujoco.MjData(m)
+    worldids = jp.arange(batch_size)
+    dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
+    dx_batch = jax.device_put(dx_batch, cpu_device)
+    dx_batch_orig = dx_batch
+
+    for _ in range(10):
+      dx_batch = jax.vmap(forward.step, in_axes=(None, 0))(
+          mx, dx_batch
+      )
+
+    for i in range(batch_size):
+      dx = dx_batch[i]
+      dx_orig = dx_batch_orig[i]
+
+      d.qpos[:] = dx_orig.qpos
+      d.qvel[:] = dx_orig.qvel
+      d.ctrl[:] = dx_orig.ctrl
+      d.mocap_pos[:] = dx_orig.mocap_pos
+      d.mocap_quat[:] = dx_orig.mocap_quat
+      d.time = dx_orig.time
+      mujoco.mj_step(m, d, 10)
+
+      tu.assert_attr_eq(dx, d, 'qpos')
+      tu.assert_attr_eq(dx, d, 'qvel')
+      tu.assert_attr_eq(dx, d, 'time')
+      tu.assert_attr_eq(dx, d, 'ctrl')
+      tu.assert_attr_eq(dx, d, 'act')
+      tu.assert_attr_eq(dx, d, 'mocap_pos')
+      tu.assert_attr_eq(dx, d, 'mocap_quat')
+      tu.assert_attr_eq(dx, d, 'sensordata')
+
  def test_step_leading_dim_mismatch(self):
    if not _FORCE_TEST:
      if not mjxw.WARP_INSTALLED:
@@ -133,10 +133,11 @@ class SmoothTest(parameterized.TestCase):

  def test_kinematics_vmap(self):
    """Tests kinematics with batched data."""
-    if not mjxw.WARP_INSTALLED:
-      self.skipTest('Warp not installed.')
-    if not io.has_cuda_gpu_device():
-      self.skipTest('No CUDA GPU device available.')
+    if not _FORCE_TEST:
+      if not mjxw.WARP_INSTALLED:
+        self.skipTest('Warp not installed.')
+      if not io.has_cuda_gpu_device():
+        self.skipTest('No CUDA GPU device available.')

    m = tu.load_test_file('pendula.xml')

@@ -153,7 +153,9 @@ def _mjx_efc(dx, worldid: int):
  efc_pos = select(dx._impl.efc__pos)[:nefc]
  efc_type = select(dx._impl.efc__type)[:nefc]
  efc_d = select(dx._impl.efc__D)[:nefc]
-  keys_sorted = np.lexsort((-efc_pos, efc_type, efc_d))
+  keys_sorted = np.lexsort(
+      (-np.round(efc_pos, 12), efc_type, np.round(efc_d, 12))
+  )
  keys = keys[keys_sorted]

  nefc = len(keys)
@@ -180,7 +182,9 @@ def _mj_efc(d):
  else:
    efc_j = d.efc_J.reshape((-1, d.qvel.shape[0]))

-  keys = np.lexsort((-d.efc_pos, d.efc_type, d.efc_D))
+  keys = np.lexsort(
+      (-np.round(d.efc_pos, 12), d.efc_type, np.round(d.efc_D, 12))
+  )
  type_ = d.efc_type[keys]
  pos = d.efc_pos[keys]
  efc_j = efc_j[keys]