google-deepmind / google-deepmind/mujoco
[MJX] jax_enable_x64 not work in mujoco/mjx/viewer.py with put_data
- Dominant language
- C++
- Stars
- 15.2k
- Forks
- 1.8k
- Avg merge
- 10d 16h
- Merged PRs (30d)
- 25
Description
### Intro
Hi!
I am a student and use MuJoCo for my research on robot control. I want to use MJX and madrona to do RL.
### My setup
MuJoco version: 3.3.0
OS: Ubuntu 22.04
GPU: NVIDIA GeForce RTX 3090
Driver: 565.77
Cuda version: 12.5
### What's happening? What did you expect?
When debug `nan` [issue](https://github.com/google-deepmind/mujoco/issues/2541) occurs in MJX with jax.jit, `jax.config.update("jax_enable_x64", True)` is a way to show if it's a numerical precision issue.
But in `mujoco/mjx/viewer.py`, if set `jax.config.update("jax_enable_x64", True)` will cause errors like:
```shell
TypeError: Argument types differ from the types for which this computation was compiled. The mismatches are:
Argument 'd.contact.geom1' compiled with int32[496] and called with int64[496]
Argument 'd.contact.geom2' compiled with int32[496] and called with int64[496]
Argument 'd.contact.geom' compiled with int32[496,2] and called with int64[496,2]
```
This error is caused by `put_data`, now mujoco mjx only support make_data with `jax.config.update("jax_enable_x64", True)` in [io.py](https://github.com/google-deepmind/mujoco/blob/main/mjx/mujoco/mjx/_src/io.py). So we need to change `mujoco/mjx/viewer.py` with `make_data`.
And mujoco=3.3.0 not support tendon to int64 (just try [robotiq_2f85](https://github.com/google-deepmind/mujoco_menagerie/tree/main/robotiq_2f85)). This bug has fixed in latest main branch with this [issue](https://github.com/google-deepmind/mujoco/issues/2461) and [commit 4e0a4f4](https://github.com/Zezo-Ai/mujoco/commit/4e0a4f4d39eb16adaec47c5642942780521779ea)), very good job.
### Steps for reproduction
1. Edit `mujoco/mjx/viewer.py` and add `jax.config.update("jax_enable_x64", True)`
2. python viewer.py --mjcf xxx.xml
### Minimal model for reproduction
[report_issue.zip](https://github.com/user-attachments/files/19501845/report_issue.zip)
### Code required for reproduction
### Confirmations
- [x] I searched the [latest documentation](https://mujoco.readthedocs.io/en/latest/overview.html) thoroughly before posting.
- [x] I searched previous [Issues](https://github.com/google-deepmind/mujoco/issues) and [Discussions](https://github.com/google-deepmind/mujoco/discussions), I am certain this has not been raised before.
Contributor guide
Assessment
This issue has not been assessed yet.