search

google-deepmind / mujoco

Multi-Joint dynamics with Contact. A general purpose physics simulator.
https://mujoco.org
Apache License 2.0
7.85k stars 787 forks source link

[MJX] Incorrect accelerations when setting initial joint positions via mjx.Data #1117

Closed driesmarzougui closed 11 months ago

driesmarzougui commented 11 months ago

Hi

I'm currently trying to get a simple environment working in MJX. The environment looks like this: screenshot

The red sphere represents a target to which the robot has to locomote.

Upon environment resetting, I would like to relocate the red sphere on the XY plane randomly. I currently do so by overwriting the mjx.Data.qpos values of two sliding joints I added to the sphere's body, followed by an mjx.forward call. Doing so, however, currently causes the sphere to receive high and incorrect acceleration values. I have prepared a minimal code example to reproduce this issue: mjx_target_position_issue.zip

What's strange to me is that I do the same thing for the robot as well; This however does not induce the same issue. Furthermore, if I remove the robot from the environment, this problem no longer persists.

saran-t commented 11 months ago

Looking at your XML, the relevant section is:

<body name="target/">
    <joint name="target/target_x_slide/" class="/" type="slide" axis="1 0 0" limited="false" damping="1"/>
    <joint name="target/target_y_slide/" class="/" type="slide" axis="0 1 0" limited="false" damping="1"/>
    <site name="target/target_site" class="target/" type="sphere" rgba="0.71764705882352942 0.33725490196078434 0.34901960784313724 1" size="0.20000000000000001"/>
</body>

This body has no geom in it and therefore has no mass or inertia! In any case, in classic MuJoCo you wouldn't add joints to this, but rather just change site_xpos in mjModel directly. Not sure if MJX would be very happy about that, maybe @btaba can chime in.

In the meantime, add some large mass to the body using <inertial>.

driesmarzougui commented 11 months ago

@saran-t thank you for the reply. I had previously tried to replace the site with a heavy geom, but that does not solve the issue (and thus neither does the addition of an <inertial> element).

The reason why I added two joints was to more closely mimic the brax reacher example. Previously I indeed got it working by changing the model, but that requires to both change the mjx.Model for computation and the original mjModel for visualisation (which I didn't really like; Alternatively, only the mjModel needs to be adapted if we reput the model on the GPU, but I do not like that transfer to happen on every environment reset). Nevertheless, I will revert back to updating the mjx.Model. It could be handy however to have a function, similar to mjx.device_get_into, that can transfer a mjx.Model off the device back into a mjModel.

What's strange to me is that I do the same thing for the robot as well; This however does not induce the same issue. Furthermore, if I remove the robot from the environment, this problem no longer persists.

Do you have any intuition with respect to this behaviour though?

saran-t commented 11 months ago

Geom won't work unless you've also excluded contacts with the ground. Try again with the inertial, leave the site alone.

driesmarzougui commented 11 months ago

Tried both cases but the issue persists.

btaba commented 11 months ago

Thanks @driesmarzougui for the bug report with the code example!

I'm able to reproduce the issue with this minimal example:

xml = "<xml provided above with added <geom size="0.1"/> at the target>"

model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)
renderer = mujoco.Renderer(model)

jit_step = jax.jit(mjx.step)

def get_image(state) -> np.ndarray:
  d = mujoco.MjData(model)
  mjx.device_get_into(d, state)
  mujoco.mj_forward(model, d)
  renderer.update_scene(d)
  return renderer.render()

# grab a trajectory
images = []
for i in range(100):
  mjx_data = jit_step(mjx_model, mjx_data)
  images.append(get_image(mjx_data))

  if i == 50:
    # move the sphere
    target_joint_id = model.joint("target/target_x_slide/").id
    qpos = mjx_data.qpos
    qpos = qpos.at[target_joint_id: target_joint_id + 2].set(
        jp.array([1.0, 1.0]))
    mjx_data = mjx_data.replace(qpos=qpos)
    mjx_data = jax.jit(mjx.forward)(mjx_model, mjx_data)

media.show_video(images)

Behavior: after the ball is displaced to (1, 1) at step 50, it slides back to (0, 0) Expected behavior: the ball should remain at (1, 1) after step 50

If you add a joint range to the slide joints, the issue goes away, so I expect this is an issue with joint ranges that I think @erikfrey has a fix for

driesmarzougui commented 11 months ago

@btaba thanks! In any case, should behaviour such as random target positions be implemented through the Model or through the Data? I would prefer to keep things similar to classic MuJoCo and do it via the Model. But in that case, it would be nice to have a function similar to mjx.device_get_into that works on Models (for visualisation purposes).