From 84f09f3fa6ea58d7c768ee68671829e4575c7a1f Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 9 Mar 2026 11:33:41 -0700 Subject: [PATCH] Add mesh_normalnum to MJX API. PiperOrigin-RevId: 880957335 Change-Id: Ib3c37afa593c888ec109f520947958a5d47b0be2 --- mjx/mujoco/mjx/warp/forward.py | 6 ++++++ mjx/mujoco/mjx/warp/types.py | 2 ++ 2 files changed, 8 insertions(+) diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 11f7dd9f..6496f460 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -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, diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index d86108f6..56ee25f8 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -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,