mjx.put_data() and mjx.put_model() takes lot of time to resolve

[AnxSnip]

Hi,

I'm a student and I'm trying to use MuJoCo to model forces applied on a cargo shipment.

I'm looking for some help with the loading of the data from mujoco to mjx.

Every is working fine with the mujoco model and the simulation works well (takes 1 to 10 seconds depending on the force). But I wanted it to be faster so I began using MJX. But I noticed that the conversion to jax array and the jitted step method takes way longer time than the classic mujoco version.

Am I doing something wrong ? or it is just that MJX cannot handle so many shapes?

Here's the code I use:
`
def run_jitted(xml):
duration = 30
jit_step = jax.jit(mjx.step)

mj_model = mujoco.MjModel.from_xml_string(xml)
mj_data = mujoco.MjData(mj_model)
mjx_model = mjx.put_model(mj_model)

mujoco.mj_resetData(mj_model, mj_data)
mjx_data = mjx.put_data(mj_model, mj_data)
while mjx_data.time < duration:
    mjx_data = jit_step(mjx_model, mjx_data)
return mj_model, mjx.get_data(mj_model, mjx_data)

`

Here is a model which explains my question:

https://pastebin.com/W0xQXiFL

[btaba]

Hi @AnxSnip , see here for more info about the performance characteristics of MJX: https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits

Seems like you are trying to simulate a bunch of boxes? Try to fiddle with some of these parameters (NB: max contacts, the solver iters):

mujoco/mjx/mujoco/mjx/test_data/shadow_hand/right_hand.xml

Lines 4 to 11 in e885a3d

	<option impratio="10" iterations="1" ls_iterations="4" timestep="0.001">
	<flag eulerdamp="disable"/>
	</option>
	<custom>
	<numeric data="15" name="max_contact_points"/>
	<numeric data="15" name="max_geom_pairs"/>
	</custom>

The timestep seems really high, but I notice you use RK4. Consider lowering the timestep and using Newton for MJX. Hope that helps.

[AnxSnip]

Hello @btaba, thank you for your quick response.
I am currently trying to do some supervised learning on the stability of a bunch of boxes in a container under acceleration force (as if they were in a truck)
My goal would be to be automatically load a setup from my database to mujoco and then get the state before and after the simulation to compute the stability of the model. The boxes will be under force for around 10 to 30s. But I need it to be able to do fast simulation.
The simulation with this xml takes 10s to complete, but I wanted it to be faster so I looked into MJX but even after 30min it is still not finishing.
I choose these parameters because it seemed to me that they were close to what I expected to happen in real life. But because it is the first time I use physics simulation, I don't really know the meaning of each parameter
I tried changing the values of the parameters you provided but It behave weirdly: some blocks are flying, or exploding without any apparent reason.
If you have any lead or reading that would help me, it will be really kind.
Thank you!

[btaba]

Hi @AnxSnip , sorry this slipped through. Could you re-paste a full example XML so we can help debug? Looks like the previous link expired

[AnxSnip]

Hello @btaba, no need to worry, I managed to stabilize my model and realized that handling collisions with multiple items doesn't work well with MJX. So, I think I'll stick with the classic version. However, I now have some questions about how to apply force to a body. I noticed that we can use site and actuator in the XML or xfrc_applied in the code, but there aren't many examples or much documentation on this. Can you point me to some resources?
This is my new model : https://pastebin.com/4xEFzr1r
I want to add movement to the container, ideally at the bottom toward any horizontal direction.
Thank you!