[MJX] Capturing Images of Scene in MJX

[Servo97]

Hi,

I'm a PhD student and I'm trying to use MuJoCo for a multi agent environment with a camera.

I can easily get a camera image using the traditional pipeline of mujoco (src)

glfw.init()
glfw.window_hint(glfw.VISIBLE, 0)
window = glfw.create_window(width, height, "Offscreen", None, None)
glfw.make_context_current(window)
glfw.swap_interval(1)
framebuffer_width, framebuffer_height = glfw.get_framebuffer_size(window)

opt = mujoco.MjvOption()
cam = mujoco.MjvCamera()
cam.lookat = np.array((0.13125, 0.1407285, 1.5))
cam.distance = 0.85
cam.azimuth = 0
cam.elevation = 90
scene = mujoco.MjvScene(model, maxgeom=10000)
pert = mujoco.MjvPerturb()
context = mujoco.MjrContext(model, mujoco.mjtFontScale.mjFONTSCALE_150.value)
viewport = mujoco.MjrRect(0, 0, framebuffer_width, framebuffer_height)
rgb_pixels = np.zeros((height, width, 3), dtype=np.uint8)
mujoco.mjv_updateScene(model, data, opt, pert, cam, mujoco.mjtCatBit.mjCAT_ALL.value, scene)

while True:
       % Simulation Steps...

        mujoco.mjr_render(viewport, scene, context)
        mujoco.mjr_readPixels(rgb_pixels, None, viewport, context)
        plt.imshow(np.flipud(rgb_pixels))
        plt.show()

I've been following the tutorial here, but I'm unsure as how to get similar "rgb_pixels" in the MJX pipeline.

Can someone please help me with this?

[btaba]

Hi @Servo97 , have you considered doing an mjx.get_data and using the same pipeline you have for MuJoCo? You can also take a look at the MJX colab tutorial which has some rendering examples