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

Add mesh_normalnum to MJX API.

Browse Source 
PiperOrigin-RevId: 880957335
Change-Id: Ib3c37afa593c888ec109f520947958a5d47b0be2
This commit is contained in:
Alessio Quaglino
2026-03-09 11:33:41 -07:00
committed by Copybara-Service
parent 3a99f6110e
commit 84f09f3fa6
2 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/forward.py
+6
View File
@@ -213,6 +213,7 @@ def _forward_shim(
mesh_graphadr: wp.array(dtype=int),
mesh_normal: wp.array(dtype=wp.vec3),
mesh_normaladr: wp.array(dtype=int),
mesh_normalnum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
@@ -646,6 +647,7 @@ def _forward_shim(
_m.mesh_graphadr = mesh_graphadr
_m.mesh_normal = mesh_normal
_m.mesh_normaladr = mesh_normaladr
_m.mesh_normalnum = mesh_normalnum
_m.mesh_polyadr = mesh_polyadr
_m.mesh_polymap = mesh_polymap
_m.mesh_polymapadr = mesh_polymapadr
@@ -1447,6 +1449,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.mesh_graphadr,
m.mesh_normal,
m.mesh_normaladr,
m.mesh_normalnum,
m._impl.mesh_polyadr,
m._impl.mesh_polymap,
m._impl.mesh_polymapadr,
@@ -1989,6 +1992,7 @@ def _step_shim(
mesh_graphadr: wp.array(dtype=int),
mesh_normal: wp.array(dtype=wp.vec3),
mesh_normaladr: wp.array(dtype=int),
mesh_normalnum: wp.array(dtype=int),
mesh_polyadr: wp.array(dtype=int),
mesh_polymap: wp.array(dtype=int),
mesh_polymapadr: wp.array(dtype=int),
@@ -2424,6 +2428,7 @@ def _step_shim(
_m.mesh_graphadr = mesh_graphadr
_m.mesh_normal = mesh_normal
_m.mesh_normaladr = mesh_normaladr
_m.mesh_normalnum = mesh_normalnum
_m.mesh_polyadr = mesh_polyadr
_m.mesh_polymap = mesh_polymap
_m.mesh_polymapadr = mesh_polymapadr
@@ -3239,6 +3244,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.mesh_graphadr,
m.mesh_normal,
m.mesh_normaladr,
m.mesh_normalnum,
m._impl.mesh_polyadr,
m._impl.mesh_polymap,
m._impl.mesh_polymapadr,
mjx/mujoco/mjx/warp/types.py
+2
View File
@@ -669,6 +669,7 @@ _NDIM = {
'mesh_graphadr': 1,
'mesh_normal': 2,
'mesh_normaladr': 1,
'mesh_normalnum': 1,
'mesh_polyadr': 1,
'mesh_polymap': 1,
'mesh_polymapadr': 1,
@@ -1198,6 +1199,7 @@ _BATCH_DIM = {
'mesh_graphadr': False,
'mesh_normal': False,
'mesh_normaladr': False,
'mesh_normalnum': False,
'mesh_polyadr': False,
'mesh_polymap': False,
'mesh_polymapadr': False,