google-deepmind / google-deepmind/mujoco

[MJX] jax_enable_x64 not work in mujoco/mjx/viewer.py with put_data

Open
#2,565 2 comments 0 reactions 0 assignees View on GitHub
bug MJX
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.