[MJX] does MJX have its own renderer or does it rely on MuJoCo's renderer

[Kallinteris-Andreas]

Hi,

I'm a maintainer of Gymnasium & Gymnasium-Robotics, and I'm trying to use MuJoCo-MJX for "prototyping MJX-based RL environments in Gymnasium".

Here is the model tested: Gymnasium/HalfCheetah (though it should be not relevant for this question)

The MJX tutorial in the section "Define a render utility function", you create a copy of the mjx_data into d (type(d)==mujoco.mj_data) and using it with base MuJoCo's render function to render the rgb_array.

def get_image(state: State, camera: str) -> np.ndarray:
  """Renders the environment state."""
  d = mujoco.MjData(env.model)
  # write the mjx.Data into an mjData object
  mjx.device_get_into(d, state.pipeline_state)
  mujoco.mj_forward(env.model, d)
  # use the mjData object to update the renderer
  renderer.update_scene(d, camera=camera)
  return renderer.render()

Didn't brax have its own render utilities?
Does MJX plan to add its own render functions, or it is the plan to keep using MuJoCo's render functions.

Thanks!

[btaba]

Hi @Kallinteris-Andreas , brax had its owner render utilities but they are not compatible with MJX. We plan to keep using MuJoCo render functions in MJX.

If MJX is ported back to brax (needs a lot of plumbing), we'll likely keep the brax HTML rendering there (but not in the MuJoCo repository.

[rs545837]

Instead of saving or showing the video with the help of mediap-py library is there also a way to render the model with the help of mujoco viewer while the RL training is going on.

[btaba]

Hey @rs545837 , here is how the MuJoCo viewer is used with MJX. You could try loading a checkpoint and running a policy using that sample code

For checkpointing in Brax see: google/brax#299
For loading a model in Brax see: google/brax#403

[Kallinteris-Andreas]

in gymnasium you can do gymnasium.make("Ant-v5", render_mode="human")

[rs545837]

Hey @btaba, I am a bit confused when trying to do the same. Can you please elaborate the whole process if you don't mind? I understand the part where you said about saving and loading a checkpoint but then how to use viewer to show the state.

Hey @rs545837 , here is how the MuJoCo viewer is used with MJX. You could try loading a checkpoint and running a policy using that sample code

For checkpointing in Brax see: google/brax#299 For loading a model in Brax see: google/brax#403

Also isn't saving and loading is already there in the tutorial.py for MJX, should I just use policy_params_fn instead of params_fn in the same code while getting make_inference_fn and params.

[lucattyy]

Hi,

I'm a maintainer of Gymnasium & Gymnasium-Robotics, and I'm trying to use MuJoCo-MJX for "prototyping MJX-based RL environments in Gymnasium".

Here is the model tested: Gymnasium/HalfCheetah (though it should be not relevant for this question)

The MJX tutorial in the section "Define a render utility function", you create a copy of the mjx_data into d (type(d)==mujoco.mj_data) and using it with base MuJoCo's render function to render the rgb_array.

def get_image(state: State, camera: str) -> np.ndarray:
  """Renders the environment state."""
  d = mujoco.MjData(env.model)
  # write the mjx.Data into an mjData object
  mjx.device_get_into(d, state.pipeline_state)
  mujoco.mj_forward(env.model, d)
  # use the mjData object to update the renderer
  renderer.update_scene(d, camera=camera)
  return renderer.render()

Didn't brax have its own render utilities? Does MJX plan to add its own render functions, or it is the plan to keep using MuJoCo's render functions.

Thanks!

Excuse me. Have you completed the development? where can I see the program if it is done.
Thanks!

[Kallinteris-Andreas]

@lucattyy Farama-Foundation/Gymnasium#833

[Andrew-Luo1]

If you're doing RL using depth-images, MJX rendering via mjx.ray can be useful; see my response here. This MJX implementation of a depth-camera can render up to millions of FPS on a very simple setup - not sure how the ray tracing scales with geom count.