google_barkour_vb/scene_hfield_mjx.xml - robot sinks into the floor

Which model is the issue affecting?

google_barkour_vb/scene_hfield_mjx.xml

What is the issue?

Issue 1)

The robot sinks into and falls below the floor when I simulate the scene using the (non-jax) basic simulation loop from https://github.com/google-deepmind/mujoco/blob/d029e4ff991851c12fdb2a49555bd65d1dbdf705/mjx/tutorial.ipynb

https://github.com/google-deepmind/mujoco_menagerie/assets/174741646/c505364f-5090-4aff-b3f1-4a1b5a187e9a

The robot bounces back and stabilizes on top of the height field if I run the simulation using jax:

https://github.com/google-deepmind/mujoco_menagerie/assets/174741646/f00ef2cb-48e7-44ad-8842-cec7751f64ca

The starting height of the robot seems to be too low here. In both cases it spawns well into the height field.

If I set the freejoint z coordinate to e.g. 2.0 (mj_data.qpos[2] = 2) then the robot bounces back in the non-jax case, so setting the initial height higher seems like a good solution here.

Issue 2)

If I simulate using jax then the resulting video rendered with mujoco.Renderer is colored incorrectly, as you can see in the latter video.

Is there any additional context you can provide (e.g., a spec sheet or a URDF to show a value mismatch)?

The code I used for testing:

import argparse

import cv2
import mujoco

parser = argparse.ArgumentParser()
parser.add_argument("--jax", action="store_true")
parser.add_argument("--elevate", type=float)
args = parser.parse_args()

width, height = 640, 480

mj_model = mujoco.MjModel.from_xml_path("google_barkour_vb/scene_hfield_mjx.xml")
mj_data = mujoco.MjData(mj_model)
renderer = mujoco.Renderer(mj_model, height=height, width=width)

duration = 3.8  # (seconds)
framerate = 60  # (Hz)

frames = []
mujoco.mj_resetData(mj_model, mj_data)

if args.elevate is not None:
    mj_data.qpos[2] = args.elevate

if not args.jax:
    while mj_data.time < duration:
        mujoco.mj_step(mj_model, mj_data)
        if len(frames) < mj_data.time * framerate:
            renderer.update_scene(mj_data)
            pixels = renderer.render()
            frames.append(pixels)
else:
    import jax
    from mujoco import mjx

    jit_step = jax.jit(mjx.step)
    mjx_model = mjx.put_model(mj_model)
    mjx_data = mjx.put_data(mj_model, mj_data)

    while mjx_data.time < duration:
        mjx_data = jit_step(mjx_model, mjx_data)
        mjx.get_data_into(mj_data, mj_model, mjx_data)
        if len(frames) < mjx_data.time * framerate:
            renderer.update_scene(mj_data)
            pixels = renderer.render()
            frames.append(pixels)

renderer.close()

device_suffix = "mjx" if args.jax else "mj"
elevate_suffix = f"elevate{'X' if args.elevate is None else args.elevate}"
save_suffix = [device_suffix, elevate_suffix]
save_suffix = "_".join(save_suffix)

codec = cv2.VideoWriter_fourcc(*"mp4v")
video_writer = cv2.VideoWriter(f"video_{save_suffix}.mp4", codec, 60, (width, height))
for frame in frames:
    video_writer.write(frame[:, :, [2, 1, 0]])
video_writer.release()

P.s. had to do a h264 conversion to make the videos work on github: ffmpeg -i input.mp4 -vcodec libx264 output.mp4

google-deepmind / mujoco_menagerie

google_barkour_vb/scene_hfield_mjx.xml - robot sinks into the floor #70