MJX collision between Ellipsoid / Cylinder and Hfield

[youwyu]

The feature, motivation and pitch

As shown in MJX 3.2.7, ELLIPSOID and CYLINDER are implemented but only collide with other primitives. If I use mesh for wheeled robots, the GPU memory consumption will be incredible for Rough environments with hfield. If I only use the empty world with flat plane, the GPU memory consumption is also not reasonable.

Alternatives

No response

Additional context

Details for my settings. I use Mujoco-Playground and parallel 8192 wheeled robots. I use train_jax_ppo.py for training.

1. Flat Plane
   The GPU memory will take around 20GB, much bigger than other examples (e.g., GO1) that take no more than 5GB.
2. Rough Hfield
   W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3021] Can't reduce memory use below 17.42GiB (18704526550 bytes) by rematerialization; only reduced to 168.35GiB (180768929689 bytes), down from 169.22GiB (181694524993 bytes) originally

[hello-robot-shehab]

Hi @yuvaltassa, I didn't want to open a duplicate/related issue; do you know if supporting cylinder to mesh geom collisions in MJX is planned to be supported?

We're trying to use this mobile manipulator from the Mujoco Menagerie https://github.com/google-deepmind/mujoco_menagerie/tree/main/hello_robot_stretch_3 with MJX, but the main issue with using this MJCF are the cylinders used for wheel collision. Like @youwyu says, using a mesh cylinders for wheels is no-go.

Thank you for your thoughts on what we could do here!

Side-question: I don't know much about how these geom colliders are built, but I see sphere <> mesh collisions implemented, is there a way to re-purpose that sphere implementation for cylinders? Thanks!