Help Needed with IndexError in Brax Environment on Kaggle

[AliSouliman1]

Hi everyone,

I'm encountering an IndexError when trying to run a MuJoCo/Brax simulation on Kaggle. The code was working previously, but I made some changes related to package versions and now I'm stuck.

I'm using the notebook from this tutorial: https://github.com/goncalog/ai-robotics/blob/main/tutorials/robot_tricks.ipynb

Initially, it was working (though I did get some warnings about jax.tree_map which I was trying to resolve). I decided to remove the specific version numbers for brax, mujoco, and mujoco_mjx in the installation commands. This is when the IndexError started.

Here's the full error traceback:

IndexError Traceback (most recent call last)
in <cell line: 3>()
1 # Instantiate the environment
2 env_name = "op3"
----> 3 env = envs.get_environment(env_name)
4
5 # Define the jit reset/step functions

/usr/local/lib/python3.10/dist-packages/brax/envs/init.py in get_environment(env_name, **kwargs)
61 env: an environment
62 """
---> 63 return _envsenv_name
64
65

in init(self, terminate_when_unhealthy, reset_noise_scale, exclude_current_positions_from_observation, **kwargs)
15 mj_model.opt.ls_iterations = 6
16
---> 17 sys = mjcf.load_model(mj_model)
18
19 physics_steps_per_control_step = 5

/usr/local/lib/python3.10/dist-packages/brax/io/mjcf.py in load_model(mj)
446 link.transform.rot[free_idx] = np.array([1.0, 0.0, 0.0, 0.0])
447
--> 448 mjx_model = mjx.put_model(mj)
449
450 sys = System( # pytype: disable=wrong-arg-types # jax-ndarray

/usr/local/lib/python3.10/dist-packages/mujoco/mjx/_src/io.py in put_model(m, device, _full_compat)
197 dataid = m.geom_dataid[i]
198 if fields['mesh_convex'][dataid] is None:
--> 199 fields['mesh_convex'][dataid] = mesh.convex(m, dataid) # pytype: disable=unsupported-operands
200 fields['mesh_convex'] = tuple(fields['mesh_convex'])
201

/usr/local/lib/python3.10/dist-packages/mujoco/mjx/_src/mesh.py in convex(m, data_id)
262 graph_idx = 0
263
--> 264 numvert, numface = graph[0], graph[1]
265 graph_idx += 2
266

IndexError: index 0 is out of bounds for axis 0 with size 0

The error seems to be related to the mesh files used by MuJoCo, I'm not sure how removing the version specifications would cause this. Any help would be greatly appreciated! I'm particularly wondering if there are any known compatibility issues with different versions of these libraries.