Fix bug with single actuator in transmission.

PiperOrigin-RevId: 574637727
Change-Id: Ibfc01ff3bc80afe5da253fd9d6d76609b8e9198d

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)):
+  for idx, (arg, typ) in enumerate(zip(args, in_types)):
-    if len(a) != size[typ]:
+    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]}".'
           )

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

mjx/mujoco/mjx/test_data/single_pendulum.xml

@@ -2,8 +2,12 @@
   <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
   <worldbody>
     <body>
-      <joint axis="1 0 0" type="hinge"/>
+      <joint name="slider" axis="1 0 0" type="hinge"/>
       <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
     </body>
   </worldbody>
+  <!-- Tests that a single actuator doesn't get mangled in a physics step. -->
+  <actuator>
+    <motor name="slide" joint="slider" gear="10" ctrllimited="true" ctrlrange="-1 1"/>
+  </actuator>
 </mujoco>