google-deepmind / google-deepmind/mujoco
`model.bind()` with a slice of length 1 unexpectedly squeezes the batch dimension
- Dominant language
- C++
- Stars
- 15.2k
- Forks
- 1.8k
- Avg merge
- 10d 16h
- Merged PRs (30d)
- 25
Description
### Intro
Hi!
### My setup
```console
python -c "import mujoco; print(mujoco.__version__)"
3.5.0
```
### What's happening? What did you expect?
When using the `.bind()` method in the Python bindings (either for standard MuJoCo or MJX) on a list/slice of objects of length one, the binding unexpectedly squeezes the batch dimension.
For example, `model.bind(geoms[0:2]).size` returns an array of shape `(2, 3)` (which is correct), but `model.bind(geoms[0:1]).size` returns an array of shape `(3,)`, which is surprising to me. I expected the latter to return an array of shape `(1, 3)` to preserve the batch dimension, keeping it consistent with slicing of length > 1. This is a major potential source of error unless the user checks for the input length (with `len(geoms)`) or explicitly unsqueezes the result (with `output.reshape(-1, 3)`).
Note that this happens with both model and data fields.
### Steps for reproduction
1. Load a model with at least two bodies/geoms.
2. Bind a subset using a slice of length > 1 and check the shape (preserves batch dim).
3. Bind a subset using a slice of length == 1 and check the shape.
4. See that the dimension drops from 2D (1, N) to 1D (N,).
### Minimal model for reproduction
Minimal XML
```xml
```
### Code required for reproduction
```python
import mujoco
from mujoco import mjx
import numpy as np
xml = """
"""
s = mujoco.MjSpec.from_string(xml)
m = s.compile()
mx = mjx.put_model(m)
geoms = s.geoms[0:2]
print("--- MuJoCo (m.bind) ---")
print("m.bind(geoms[0:2]).size.shape:", np.array(m.bind(geoms[0:2]).size).shape) # Expected (2, 3)
# Actual (2, 3)
print("m.bind(geoms[0:1]).size.shape:", np.array(m.bind(geoms[0:1]).size).shape) # Expected (1, 3)
# Actual (3,)
print("\n--- MJX (mx.bind) ---")
print("mx.bind(geoms[0:2]).size.shape:", mx.bind(geoms[0:2]).size.shape) # Expected (2, 3)
# Actual (2, 3)
print("mx.bind(geoms[0:1]).size.shape:", mx.bind(geoms[0:1]).size.shape) # Expected (1, 3)
# Actual (3,)
```
https://github.com/google-deepmind/mujoco/compare/main...hartikainen:mujoco:bind-squeeze also includes a simple test case for this.
### 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.