chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add naccdmax to mjx io.

Browse Source 
PiperOrigin-RevId: 879780258
Change-Id: I24f9d0b3205577322bebadc3d03796c8dcadde48
This commit is contained in:
Baruch Tabanpour
2026-03-06 13:02:16 -08:00
committed by Copybara-Service
parent 6a0ebac4b1
commit 6bc6c11c49
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+8 -2
View File
@@ -872,6 +872,7 @@ def _make_data_warp(
m: Union[types.Model, mujoco.MjModel], m: Union[types.Model, mujoco.MjModel],
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
naconmax: Optional[int] = None, naconmax: Optional[int] = None,
naccdmax: Optional[int] = None,
njmax: Optional[int] = None, njmax: Optional[int] = None,
) -> types.Data: ) -> types.Data:
"""Allocate and initialize Data for the Warp implementation.""" """Allocate and initialize Data for the Warp implementation."""
@@ -882,7 +883,7 @@ def _make_data_warp(
) )
with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable
dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, njmax=njmax) # pylint: disable=undefined-variable dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, naccdmax=naccdmax, njmax=njmax) # pylint: disable=undefined-variable
fields = _make_data_public_fields(m) fields = _make_data_public_fields(m)
for k in fields: for k in fields:
@@ -971,6 +972,7 @@ def make_data(
_full_compat: bool = False, # pylint: disable=invalid-name _full_compat: bool = False, # pylint: disable=invalid-name
nconmax: Optional[int] = None, nconmax: Optional[int] = None,
naconmax: Optional[int] = None, naconmax: Optional[int] = None,
naccdmax: Optional[int] = None,
njmax: Optional[int] = None, njmax: Optional[int] = None,
keepalive_refs: Optional[Dict[int, Any]] = None, keepalive_refs: Optional[Dict[int, Any]] = None,
) -> types.Data: ) -> types.Data:
@@ -989,6 +991,10 @@ def make_data(
Since the number of worlds is **not** pre-defined in JAX, we use the Since the number of worlds is **not** pre-defined in JAX, we use the
`naconmax` argument to set the upper bound for the number of contacts `naconmax` argument to set the upper bound for the number of contacts
across all worlds, rather than the `nconmax` argument from MuJoCo Warp. across all worlds, rather than the `nconmax` argument from MuJoCo Warp.
naccdmax: maximum number of contacts for GJK collision detection across all
worlds. Since the number of worlds is **not** pre-defined in JAX, we use the
`naccdmax` argument to set the upper bound for the number of contacts
across all worlds, rather than the `nccdmax` argument from MuJoCo Warp.
njmax: maximum number of constraints to allocate for warp across all worlds njmax: maximum number of constraints to allocate for warp across all worlds
keepalive_refs: optional dict to store references to underlying MuJoCo keepalive_refs: optional dict to store references to underlying MuJoCo
objects, preventing them from being garbage collected. Required for CPP objects, preventing them from being garbage collected. Required for CPP
@@ -1025,7 +1031,7 @@ def make_data(
elif impl == types.Impl.WARP: elif impl == types.Impl.WARP:
_check_warp_installed() _check_warp_installed()
naconmax = nconmax if naconmax is None else naconmax naconmax = nconmax if naconmax is None else naconmax
return _make_data_warp(m, device, naconmax, njmax) return _make_data_warp(m, device, naconmax, naccdmax, njmax)
raise NotImplementedError( raise NotImplementedError(
f'make_data for implementation "{impl}" not implemented yet.' f'make_data for implementation "{impl}" not implemented yet.'