[MJX] Any plans to soon support collisions for Height fields?

[AlexS28]

I am currently working on an RL controller for Biped locomotion for a journal submission. I am very happy with MJX so far, and able to make my robot walk and follow joystick commands with the controller. However, one of the most critical parts for robustness and making the robot actually handle outside environments is to include rough terrain. I don't think it's possible to create a robust robot locomotion outside the lab without training on rough terrain to some extend. I know how to make rough terrain using the heightfield parameter, but currently MJX does not support it? Are there any alternatives at all? My guess is that a mesh would have computational issues? Either way, I was wondering how soon we could expect the height field collisions to work for MJX in terms of development?

Any other suggestions also appreciated (currently I decided just to make a bunch of thin plates using box geometry). But for our paper we'll need to put the robot outside the lab, and it would need to handle all kinds of rough terrain, so I am not sure how good using thin box shapes can really represent that type of environment for testing.

[btaba]

Hi @AlexS28 , could you share a bit more about your embodiment, and what kind of geometries you are using for the feet? Height-fields are not being prioritized, but we are happy to accept contributions, and I may be able to give some pointers.

[breakds]

A height map with random pattern (e.g. Perlin noise) is usually used as a terrain to train locomotion of robot e.g. with RL so that the produced policy is robust. A pattern of 40 x 40 vertices will produce a height map of 1600 and makes MuJoCo simulation much slower based my past experience. Do you think there is a way MJX can potentially accelerate this?

Thanks!

[btaba]

Hi @breakds , as a first pass, we'll probably implement something pretty similar to what is used in MuJoCo, at least for sphere-hfield, and we can vmap over the hfield prisms. I'm not sure how fast this will be, but it'll definitely be slower than plane collisions :) I'm guessing with further fussing with geoms/conaffinity/contype, performance might be somewhat reasonable, but I'm not sure

[breakds]

I see. Thanks for the explanation @btaba ! In that case, maybe an alternative would be to approximate uneven plane with something like "stairs"?

[btaba]

Are you thinking to use boxes for now? Go for it! If you have a lot of boxes, consider adding these fields to your model

https://github.com/google-deepmind/mujoco/blob/a8db22f0d077aee86f771808fd24ea60a148c93f/mjx/mujoco/mjx/test_data/shadow_hand/right_hand.xml#L8-L11

max_geom_pairs will do some rough broadphase

[breakds]

Yes using boxes should also work.

Thanks for sharing the code snippet! Didn't know about those two options, TIL.

[driesmarzougui]

Are you thinking to use boxes for now? Go for it! If you have a lot of boxes, consider adding these fields to your model

https://github.com/google-deepmind/mujoco/blob/a8db22f0d077aee86f771808fd24ea60a148c93f/mjx/mujoco/mjx/test_data/shadow_hand/right_hand.xml#L8-L11

max_geom_pairs will do some rough broadphase

Hey @btaba! Are those fields documented somewhere?

[btaba]

@driesmarzougui it's on the TODO list to document these, will open as a separate issue