MJX V3.1.6 Python 3.8 Compatibility

[willthibault]

Hi,

I believe that there may be an issue with MJX V3.1.6 when using Python 3.8. When importing mjx the following error occurs:

Python 3.8.10 (default, Nov 14 2022, 12:59:47) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from mujoco import mjx
/usr/lib/python3/dist-packages/requests/__init__.py:89: RequestsDependencyWarning: urllib3 (2.2.1) or chardet (3.0.4) doesn't match a supported version!
  warnings.warn("urllib3 ({}) or chardet ({}) doesn't match a supported "
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.8/dist-packages/mujoco/mjx/__init__.py", line 18, in <module>
    from mujoco.mjx._src.collision_driver import collision
  File "/usr/local/lib/python3.8/dist-packages/mujoco/mjx/_src/collision_driver.py", line 48, in <module>
    from mujoco.mjx._src.collision_convex import box_box
  File "/usr/local/lib/python3.8/dist-packages/mujoco/mjx/_src/collision_convex.py", line 1053, in <module>
    collider_fn: Callable[[_GeomInfo, _GeomInfo], Collision],
TypeError: 'ABCMeta' object is not subscriptable

Adding the import from typing import Callable to collision_convex.py fixes the issue. However, from collections.abc import Callable is already imported. Their functionality is slightly different, so I wanted to confirm this fix.

This issue is mentioned here, but I created a new issue for visibility.

Thanks!