google-deepmind / google-deepmind/mujoco

`model.bind()` with a slice of length 1 unexpectedly squeezes the batch dimension

Open
#3,128 6 comments 0 reactions 1 assignee Claimed by @thowell View on GitHub
bug
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

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.