[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

[Servo97]

Hi, very sorry for the late reply.
I was finally able to add mjx.get_data to my _get_obs function and this is the error I get:

ConcretizationTypeError                   Traceback (most recent call last)
Cell In[67], line 2
      1 # initialize the state
----> 2 state = jit_reset(jax.random.PRNGKey(0))
      3 rollout = [state.pipeline_state]
      5 # grab a trajectory

    [... skipping hidden 12 frame]

Cell In[65], line 33, in DeltaArrayMJX.reset(self, rng)
     31 qvel = jp.zeros(self.sys.nv)
     32 data = self.pipeline_init(qpos, qvel)
---> 33 obs = self._get_obs(data)
     34 x = jax.random.uniform(rng_x, minval=0.009, maxval=0.19)
     35 y = jax.random.uniform(rng_y, minval=0.007, maxval=0.23)

Cell In[65], line 56, in DeltaArrayMJX._get_obs(self, data)
     55 def _get_obs(self, data: mjx.Data) -> jp.ndarray:
---> 56     mj_data = mjx.get_data(self.sys.mj_model, data)
     57     self.renderer.update_scene(mj_data, camera=self.cam)
     58     img = self.renderer.render()

File [~/anaconda3/envs/py310/lib/python3.10/site-packages/mujoco/mjx/_src/io.py:250](https://file+.vscode-resource.vscode-cdn.net/home/sarvesh/Sarvesh/Delta_Array/dexterous-manipulation-delta-arrays/Delta_Array_MJX/~/anaconda3/envs/py310/lib/python3.10/site-packages/mujoco/mjx/_src/io.py:250), in get_data(m, d)
    247 else:
    248   result = mujoco.MjData(m)
--> 250 get_data_into(result, m, d)
    252 return result

File [~/anaconda3/envs/py310/lib/python3.10/site-packages/mujoco/mjx/_src/io.py:292](https://file+.vscode-resource.vscode-cdn.net/home/sarvesh/Sarvesh/Delta_Array/dexterous-manipulation-delta-arrays/Delta_Array_MJX/~/anaconda3/envs/py310/lib/python3.10/site-packages/mujoco/mjx/_src/io.py:292), in get_data_into(result, m, d)
    290 efc_active = (d_i.efc_J != 0).any(axis=1)
    291 efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
--> 292 nefc, nc = int(efc_active.sum()), int((efc_active & efc_con).sum())
    293 result_i.nnzJ = nefc * m.nv
    294 if ncon != result_i.ncon or nefc != result_i.nefc:

    [... skipping hidden 1 frame]

File [~/anaconda3/envs/py310/lib/python3.10/site-packages/jax/_src/core.py:1516](https://file+.vscode-resource.vscode-cdn.net/home/sarvesh/Sarvesh/Delta_Array/dexterous-manipulation-delta-arrays/Delta_Array_MJX/~/anaconda3/envs/py310/lib/python3.10/site-packages/jax/_src/core.py:1516), in concretization_function_error..error(self, arg)
   1515 def error(self, arg):
-> 1516   raise ConcretizationTypeError(arg, fname_context)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.
The error occurred while tracing the function reset at [/tmp/ipykernel_2471467/3418647397.py:26](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:26) for jit. This value became a tracer due to JAX operations on these lines:

  operation a:f32[135] = copy b
    from line [/tmp/ipykernel_2471467/3418647397.py:32](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:32) (reset)

  operation a:f32[1,0,3] = pjit[
  name=_take
  jaxpr={ lambda ; b:f32[129,3] c:i32[1,0]. let
      _:i32[1,0] = pjit[
        name=remainder
        jaxpr={ lambda ; d:i32[1,0] e:i32[]. let
            f:bool[] = eq e 0
            g:i32[] = pjit[
              name=_where
              jaxpr={ lambda ; h:bool[] i:i32[] j:i32[]. let
                  k:i32[] = select_n h j i
                in (k,) }
            ] f 1 e
            l:i32[1,0] = rem d g
            m:bool[1,0] = ne l 0
            n:bool[1,0] = lt l 0
            o:bool[] = lt g 0
            p:bool[1,0] = ne n o
            q:bool[1,0] = and p m
            r:i32[1,0] = add l g
            s:i32[1,0] = select_n q l r
          in (s,) }
      ] c 129
      t:f32[1,0,3] = broadcast_in_dim[broadcast_dimensions=() shape=(1, 0, 3)] 0.0
    in (t,) }
] u v
    from line [/tmp/ipykernel_2471467/3418647397.py:32](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:32) (reset)

  operation a:f32[1,0,3] = pjit[
  name=_take
  jaxpr={ lambda ; b:f32[129,3] c:i32[1,0]. let
      _:i32[1,0] = pjit[
        name=remainder
        jaxpr={ lambda ; d:i32[1,0] e:i32[]. let
            f:bool[] = eq e 0
            g:i32[] = pjit[
              name=_where
              jaxpr={ lambda ; h:bool[] i:i32[] j:i32[]. let
                  k:i32[] = select_n h j i
                in (k,) }
            ] f 1 e
            l:i32[1,0] = rem d g
            m:bool[1,0] = ne l 0
            n:bool[1,0] = lt l 0
            o:bool[] = lt g 0
            p:bool[1,0] = ne n o
            q:bool[1,0] = and p m
            r:i32[1,0] = add l g
            s:i32[1,0] = select_n q l r
          in (s,) }
      ] c 129
      t:f32[1,0,3] = broadcast_in_dim[broadcast_dimensions=() shape=(1, 0, 3)] 0.0
    in (t,) }
] u v
    from line [/tmp/ipykernel_2471467/3418647397.py:32](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:32) (reset)

  operation a:i32[1,0] = convert_element_type[new_dtype=int32 weak_type=False] b
    from line [/tmp/ipykernel_2471467/3418647397.py:32](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:32) (reset)

  operation a:f32[1,0] = pjit[
  name=_take
  jaxpr={ lambda ; b:f32[135] c:i32[1,0]. let
      _:i32[1,0] = pjit[
        name=remainder
        jaxpr={ lambda ; d:i32[1,0] e:i32[]. let
            f:bool[] = eq e 0
            g:i32[] = pjit[
              name=_where
              jaxpr={ lambda ; h:bool[] i:i32[] j:i32[]. let
                  k:i32[] = select_n h j i
                in (k,) }
            ] f 1 e
            l:i32[1,0] = rem d g
            m:bool[1,0] = ne l 0
            n:bool[1,0] = lt l 0
            o:bool[] = lt g 0
            p:bool[1,0] = ne n o
            q:bool[1,0] = and p m
            r:i32[1,0] = add l g
            s:i32[1,0] = select_n q l r
          in (s,) }
      ] c 135
      t:f32[1,0] = broadcast_in_dim[broadcast_dimensions=() shape=(1, 0)] 0.0
    in (t,) }
] u v
    from line [/tmp/ipykernel_2471467/3418647397.py:32](https://file+.vscode-resource.vscode-cdn.net/tmp/ipykernel_2471467/3418647397.py:32) (reset)

(Additional originating lines are not shown.)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

I believe there are posts that are doing similar stuff e.g. this one.

But IIUC, most of them are using the rendering pipeline for generating a video of the scene. I want to use the rendered images, preprocess them, and use them as observations for my RL pipeline.

Do you have any suggestions for fixing this?

Thanks!

[btaba]

Hi @Servo97 , get_data is not jittable. We are actively working on rendering within a jitted function. For now, you can tinker with the ray casting in mjx.ray

See #1202

[Servo97]

Ahhh I see! Thanks for the prompt response! I will check out ray casting!

[markusheimerl]

You can always save all mjx_data states across all time steps of the simulation in an array, return that from your jitted function and then have another loop at the end of your script that "replays" these states through setting the data of each time step with renderer.update_scene(mj_data) followed by frames.append(renderer.render()) and finally a media.write_video('simulation_mjx.mp4', frames, fps=framerate)

[btaba]

YMMV, but there is an experimental branch here where MJX is hooked up to madrona for rendering. Here's a demo on the Aloha robot: https://github.com/shacklettbp/madrona_mjx/blob/aloha/scripts/mjx_aloha_demo.py