google-deepmind / google-deepmind/mujoco

[MuJoCo Warp] Selective reset_data corrupts packed contact membership across worlds

Open
#3,572 3 comments 0 reactions 1 assignee Claimed by @thowell View on GitHub
bug WARP
Dominant language
C++
Stars
15.2k
Forks
1.8k
Avg merge
10d 16h
Merged PRs (30d)
25

Description

### Intro

Hi! I am testing MuJoCo Warp's Python APIs for batched simulation and selective world resets. I found a contact-bookkeeping discrepancy affecting worlds that were not selected for reset. The affected component is `google-deepmind/mujoco_warp`.

### My setup

- API: Python, using MuJoCo Warp's exported `put_model`, `make_data`, `forward`, `reset_data`, and `get_data_into` functions.
- MuJoCo Warp: 3.13.0, tested `main` commit `d8a241ef1643295ebb85bfdfdf9171360bfb2c0e` on 2026-09-11.
- MuJoCo: 3.12.1.
- NVIDIA Warp: 1.15.0.
- Python: 3.12.13.
- OS / architecture: Linux 7.0.0-28-generic, x86_64, glibc 2.39.
- GPU: NVIDIA GeForce RTX 5060 Ti; reproduction executed on CUDA.
- CUDA Toolkit used by Warp: 12.9.
- NVIDIA driver: 595.84; `nvidia-smi` CUDA compatibility: 13.2.

### What's happening? What did you expect?

Calling `reset_data` with a per-world mask corrupts the globally packed contact prefix when more than one world currently has contacts.

Expected: reset worlds have no active contacts, while contacts belonging to unselected worlds remain active. With one contact in each of two worlds, resetting only world 0 should leave `(nacon, worldids) == (1, [1])`; resetting only world 1 should leave `(1, [0])`.

Actual:

- `[True, False]` sets the global `nacon` to zero and hides the unselected world-1 contact.
- `[False, True]` leaves `nacon` unchanged and relabels the cleared world-1 row as an additional world-0 contact.

The attached `poc/repro.py`, also included in the code field below, produced:

```text
before packed=(2, [0, 1]) host_ncon=[1, 1]
reset_world_0 packed=(0, []) world_1_host_ncon=0
reset_world_1 packed=(2, [0, 0]) world_0_host_ncon=2
BUG_REPRODUCED=True
```

`get_data_into` exposes the discrepancy through another public API: it reports no contact for the untouched world in the first case and two contacts for world 0 in the second. This discrepancy exists immediately after selective reset; a later `forward` rebuilds contacts. I am not claiming that the incorrect rows persist through a subsequent collision pass.

Contacts for all worlds share `contact[0:d.nacon[0]]` and use `contact.worldid` for membership. The selective reset kernel clears selected rows in place and changes their `worldid` to zero without compacting the prefix. Separately, the per-world reset kernel clears the single global `d.nacon[0]` only from thread `worldid == 0`, after checking the mask.

Source at the tested revision:
https://github.com/google-deepmind/mujoco_warp/blob/d8a241ef1643295ebb85bfdfdf9171360bfb2c0e/mujoco_warp/_src/io.py#L2487-L2609

Related work is in the MuJoCo Warp repository: [issue #716](https://github.com/google-deepmind/mujoco_warp/issues/716) requested per-world reset, [PR #720](https://github.com/google-deepmind/mujoco_warp/pull/720) implemented it, and [PR #586](https://github.com/google-deepmind/mujoco_warp/pull/586) introduced the original all-world reset. The closest existing `test_reset_data_world` checks only `qvel` and constructs no contacts. A bounded search found no existing report for this packed-contact failure.

### Steps for reproduction

1. Extract the attached ZIP into its own directory. It contains `poc/repro.py`, `poc/README.md`, and `poc/observed.txt`.
2. Use a CUDA-capable environment with the versions above and a clean MuJoCo Warp checkout at the tested revision:

```bash
git clone https://github.com/google-deepmind/mujoco_warp.git
cd mujoco_warp
git checkout d8a241ef1643295ebb85bfdfdf9171360bfb2c0e
uv run --frozen python /absolute/path/to/extracted/poc/repro.py
```

3. The script creates two worlds, each with one plane/sphere contact. It resets world 0 only, checks the packed prefix and host contact count, rebuilds contacts, and then repeats with world 1 only.
4. Compare the output with the expected prefixes `(1, [1])` and `(1, [0])`. On the tested revision, the output above ends with `BUG_REPRODUCED=True`.

Reproduction package: [selective-reset-packed-contact-corruption-poc.zip](https://github.com/user-attachments/files/32137029/selective-reset-packed-contact-corruption-poc.zip)

### Minimal model for reproduction

This is the complete MJCF embedded in the script. It requires no external assets.

```xml








```

### Code required for reproduction

This is the attached `poc/repro.py` verbatim. It embeds the model above and runs as a standalone script within the MuJoCo Warp checkout's environment.

```python
#!/usr/bin/env python3
"""Reproduce packed-contact corruption after per-world reset_data calls."""

import mujoco
import mujoco_warp as mjw
import warp as wp

XML = """








"""

def packed_contacts(data):
"""Return the active packed count and world IDs."""
count = int(data.nacon.numpy()[0])
return count, data.contact.worldid.numpy()[:count].tolist()

def host_contact_count(model, data, world_id):
"""Exercise the public readback path for one world."""
result = mujoco.MjData(model)
mjw.get_data_into(result, model, data, world_id=world_id)
return result.ncon

mjm = mujoco.MjModel.from_xml_string(XML)
model = mjw.put_model(mjm)
data = mjw.make_data(mjm, nworld=2, nconmax=20, njmax=40)

mjw.forward(model, data)
print(
"before",
f"packed={packed_contacts(data)}",
f"host_ncon={[host_contact_count(mjm, data, 0), host_contact_count(mjm, data, 1)]}",
)

# Reset world 0 only. reset_nworld runs one thread per world and clears the
# single global nacon only from selected world 0, hiding world 1's untouched row.
mjw.reset_data(model, data, wp.array([True, False], dtype=bool))
world0_count, world0_ids = packed_contacts(data)
world1_host_count = host_contact_count(mjm, data, 1)
print(
"reset_world_0",
f"packed={(world0_count, world0_ids)}",
f"world_1_host_ncon={world1_host_count}",
)

# Rebuild contacts, then reset world 1 only. The global count is left at two,
# while reset_contact zeroes world 1's row and relabels it as world 0.
mjw.forward(model, data)
mjw.reset_data(model, data, wp.array([False, True], dtype=bool))
world1_count, world1_ids = packed_contacts(data)
world0_host_count = host_contact_count(mjm, data, 0)
print(
"reset_world_1",
f"packed={(world1_count, world1_ids)}",
f"world_0_host_ncon={world0_host_count}",
)

bug_reproduced = (
world0_count == 0
and world1_host_count == 0
and world1_count == 2
and world1_ids == [0, 0]
and world0_host_count == 2
)
print(f"BUG_REPRODUCED={bug_reproduced}")
if not bug_reproduced:
raise SystemExit(
"current source no longer has the expected selective-reset behavior"
)
```

### 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.