Segmentation camera observations

[pthangeda]

Would be great to have per-pixel segmentation (geom IDs) as a camera data type alongside RGB and depth. This makes goal-conditioned vision tasks much easier.

I've implemented it and tested this on a 'multi-cube pick specified one' task and it works well!

mjlab_many.mp4

I can create PR here once my mujoco_warp segmentation PR is merged: google-deepmind/mujoco_warp#1236

[kevinzakka]

Hi @pthangeda, this is super useful! Would be great to merge once it lands in mjwarp :)

[pthangeda]

@kevinzakka sorry for the delay, was going back on forth on some things. Here is the PR: #862