diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py
index b54655c0..2975a32f 100644
--- a/mjx/mujoco/mjx/_src/scan.py
+++ b/mjx/mujoco/mjx/_src/scan.py
@@ -127,11 +127,11 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
def _check_input(m: Model, args: Any, in_types: str) -> None:
"""Checks that scan input has the right shape."""
size = {'b': m.nbody, 'j': m.njnt, 'q': m.nq, 'v': m.nv, 'u': m.nu, 'a': m.na}
- for idx, (a, typ) in enumerate(zip(args, in_types)):
- if len(a) != size[typ]:
+ for idx, (arg, typ) in enumerate(zip(args, in_types)):
+ if len(arg) != size[typ]:
raise IndexError(
(
- f'f argument "{idx}" with type "{typ}" has length "{len(a)}"'
+ f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
f' which does not match the in_types[{idx}] expected length of '
f'"{size[typ]}".'
)
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
index 1193a1f8..2b421f2f 100644
--- a/mjx/mujoco/mjx/_src/smooth.py
+++ b/mjx/mujoco/mjx/_src/smooth.py
@@ -462,5 +462,6 @@ def transmission(m: Model, d: Data) -> Data:
group_by='u',
)
moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val)
- d = d.replace(actuator_length=length.squeeze(), actuator_moment=moment)
+ length = length.reshape((m.nu,))
+ d = d.replace(actuator_length=length, actuator_moment=moment)
return d
diff --git a/mjx/mujoco/mjx/test_data/single_pendulum.xml b/mjx/mujoco/mjx/test_data/single_pendulum.xml
index 1877b500..98f5a0c6 100644
--- a/mjx/mujoco/mjx/test_data/single_pendulum.xml
+++ b/mjx/mujoco/mjx/test_data/single_pendulum.xml
@@ -2,8 +2,12 @@
-
+
+
+
+
+