MJX does not keep equality constraints

[EnricoMingo]

Hi,

I'm trying to use MuJoCo (MJX) to simulate a humanoid bipedal robot with several closed linkages. I am using MJX inside BRAX and when I visualize some seconds of trajectory I clearly see that the equality constraints that I put inside the model are not kept. These constraints normally work using "normal" MUJOCO.

I am using the default settings for MjModel:

mj_model.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON 
mj_model.opt.iterations = 100
mj_model.opt.ls_iterations = 50

notice that I have 32 equality constraints in my model (for most of them solref = "0.005 1"). Thank you for your help!

[yuvaltassa]

Can you please attach a minimal model that reproduces this issue?

[EnricoMingo]

Thank you for the prompt reply. I am following the tutorial here and I changed the model from the humanoid to my robot. I also have commented out the following lines of code:

mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG mj_model.opt.iterations = 6 mj_model.opt.ls_iterations = 6

I check the output from the "# grab a trajectory" part of the tutorial. I thought it was related with the different settings of the solver so I tried to put the default values as well (the ones in the previous post), but I am still having this issue. The model of the robot is quite complex, I will try to prepare a smaller version with less constraints and see if the error remains.

[sck-robot]

I seem to be having a similar issue. I created a closed-loop four-link model, which works properly under mujoco. I used connect to implement this closed loop. But under mjx, the closed loop linkage doesn't seem to be working properly. I have provided a code to reproduce this problem. Does anyone have any ideas about this? Thanks a lot.

https://github.com/google-deepmind/mujoco/assets/60764118/25d1a3ce-c87c-4e60-b3ef-957d9ffed96e

https://github.com/google-deepmind/mujoco/assets/60764118/48bddfc5-6faa-43bf-a29e-f5ac220ca8b7

close_loop_test.TXT

[EnricoMingo]

@sck-robot thank you for your example. I was planning to do something similar next week but you anticipated me :)

[EnricoMingo]

🙄

[heleiduan]

echoing a similar closed-loop violations seen in an existing model from mujoco_menagerie. an earlier thread #1129 includes the example code that we tried as well.

[erikfrey]

Hi all - confirming that I can repro this on my end using the model provided by @sck-robot. Will report back soon with a fix.

[erikfrey]

OK, this is fixed as of 67fa7c1dc21b535b9b5f906830f5ebf81dbf4d7e which will be available in MuJoCo 3.0.2. In the meantime, I can explain the bug if you would like to fix it in your user code and still use 3.0.1:

MJX is incorrectly setting d.ne to 0 in mjx.device_put and is then ignoring equality constraints during the step. For the time being, you can fix this by using mjx.make_data to initialize mjx.Data instead of mjx.device_put. So in your example @sck-robot make the following change:

79c79
<     mjx_data = mjx.forward(mjx_model, mjx_data)
---
>     mjx_data = mjx.make_data(mjx_model)

In MuJoCo 3.0.2 we will offer some new functions for placing data on device: mjx.put_data and mjx.put_model that work better in these scenarios.

[EnricoMingo]

Thank you @erikfrey for being so fast in solving the issue! Is this something I can also temporarily fix in my example? I am asking because I have pip installed mujoco 3.0.1. Thank you again.

[sck-robot]

Thank you @erikfrey for solving the issue so instantly! Now I can create the closed loop correctly under MJX 3.0.1 with your advice. I'm really looking forward to MJX's new features. Thank you again for your efforts.

[EnricoMingo]

I was able to fix my script as well. It is based on the mujoco_mjx environment for Brax example. It was just needed to apply the suggestion of @erikfrey inside the pipeline_init function.

[heleiduan]

@erikfrey Thank you for taking care of this! The closed loop in robot model is able to work properly with compiled version of commit https://github.com/google-deepmind/mujoco/commit/67fa7c1dc21b535b9b5f906830f5ebf81dbf4d7e.