[MJX] Model mesh_vert no longer jax array

[areiner222]

I noticed that the mesh vertices (previously geom_convex_verts) of the mjx.Model are no longer jax arrays. I think this precludes me from having different mesh vertices over my vmapped batch dimension which is what I've done historically - use shape parameters to modify a set of convex meshes and recompute inertial properties, etc in jax.

[btaba]

Hi @areiner222 , this got refactored in a large change where condim was added, but I plan on adding these fields back the the mjx.Model at mjx.put_model time. Apologies for the inconvenience in the interim.

You can pin to a prior version of MJ/MJX (likely 3.1.4), would that work for the time being?

cc @erikfrey

[areiner222]

Thanks for the update @btaba.

Yes will pin a prior version for now - 3.1.3 works just fine and will try 3.1.4.
I was looking to upgrade in order to be able to try higher condim to improve stability for "mesh feet/hands" x floor plane collisions. Can hold off on that until the the fields are added back.

[areiner222]

I see that you map to the convex mesh inside the collider function and compute face normals, edge directions, unique edges there.
If the mesh vertices are jax arrays, is there a plan for where you would perform those operations in the pipeline? If not mistaken, I believe to use trimesh, map to face/facets, computing the edge directions will all be challenging to do in jax with static shape constraints?

Please let me know if I'm not thinking about it correctly!

[btaba]

Hi @areiner222 , we wouldn't perform the mesh "processing" in jax, as you noted jax puts a very hard limitation on implementing any of those processing functions. We'd precompute and store results in mjx.Model as jax.Array