Fix bug with single actuator in transmission.

PiperOrigin-RevId: 574637727 Change-Id: Ibfc01ff3bc80afe5da253fd9d6d76609b8e9198d
2023-10-18 16:17:22 -07:00
parent b989a85364
commit 44ba099fd6
3 changed files with 10 additions and 5 deletions
@@ -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]}".'
          )
@@ -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
@@ -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>