[MJX] visualize parallel environments in one viewer

[panwliu]

First of all, thanks for the great work! MJX could make a great impact in the robotics community.

I'm doing massive parallel simulation using MJX. During the development phase, I'd like to visualize the parallel environments. For example, I spawned 4096 humanoids, and I'd like to visualize them all in a viewer as the image (a screenshot of Isaac Gym). I saw that mjx/mujoco/mjx/viewer.py is used to visualize one environment. How to visualize the parallel environments in one viewer?

Thanks

[rs545837]

Yes, I am also very curious to know if there's a way to do it @yuvaltassa.

[btaba]

Hi folks,

For the MJX part, you can simulate N environments and write them back to MjData via get_data which accepts a batched mjx.Data.

Now for visualizing N environments, the simplest way I can think of is to create an MjModel/MjData with N characters (procedurally generate the XML or use something like dm_env composer), and then feed the qpos from the above step back into your MjData. Others should comment on this replication bit though. I know that there has been some active discussion on better ways to repeat bodies.

[panwliu]

Hi @btaba , thanks for the response.

Since mujoco doesn't allow repeated bodies in the XML, it actually makes procedurally generating N characters not trivial. Also needs careful handling in feeding the batched mjx.Data to the right bodies.

Another possible way would be a "batched" version of mjModel. Would you think this as an option?

[yuvaltassa]

I think what you want is visual only, no? Like, your humanoids don't need to touch each other at all, right?

This is something that @vikashplus (and, as it turns out, #1395) have asked for.

What you actually want, IIUC, is something like mjv_updateScene, except you want "mjv_addToScene", i.e. some way to add together the geoms from multiple mjData instances into one visual scene. Correct?

[vikashplus]

+1 mjv_addToScene will be amazing. I have been trying to unsuccessfully hack this in multiple ways.

[panwliu]

some way to add together the geoms from multiple mjData instances into one visual scene

@yuvaltassa, exactly!

[yuvaltassa]

Okay, this is MUCH easier than your initial proposal. I'll try to do something like this in the next few days, will update here.

[AlexS28]

@yuvaltassa was wondering if there was any update on this feature. I think it would be really great to have, and not only good to get a nice overview for how the robots are behaving, but also nice for making videos for research papers etc.

[alexhansson]

@yuvaltassa Any update on this? Would definitely be huge if this is somehow possible!