Fix circular buffer push broadcast shape for non-2D data

docs/source/changelog.rst

@@ -68,6 +68,11 @@ Fixed
   platform library not loaded`` on headless Linux hosts that don't pre-set
   ``MUJOCO_GL``. The default is now applied in ``mjlab/__init__.py`` (Linux
   only) so it takes effect before mujoco's GL backend selection runs.
+- Fixed ``CircularBuffer.push`` crashing with a broadcast shape mismatch
+  when the buffered data has more than two dimensions. The first-push
+  backfill previously hard-coded ``[None, :, None]`` and ``[None, :, :]``
+  views, which only matched 2-D tensors; it now computes the broadcast
+  shape dynamically from the data's ``ndim``.
 
 Version 1.4.0 (May 26, 2026)
 ----------------------------

src/mjlab/utils/buffers/circular_buffer.py

@@ -209,9 +209,8 @@ def append(self, data: torch.Tensor) -> None:
 
     # Backfill entire history with first frame for newly initialized batches.
     is_first_push = self._num_pushes == 0
-    torch.where(
-      is_first_push[None, :, None], data[None, :, :], self._buffer, out=self._buffer
-    )
+    cond = is_first_push.view(1, self._batch_size, *([1] * (data.ndim - 1)))
+    torch.where(cond, data.unsqueeze(0), self._buffer, out=self._buffer)
 
     self._num_pushes += 1
 

tests/test_dc_actuator.py

@@ -282,3 +282,57 @@ def test_dc_motor_warns_when_effort_limit_exceeds_saturation():
       velocity_limit=30.0,
       effort_limit=25.0,  # > saturation_effort.
     )
+
+
+def test_dc_motor_delay_position_mode(device, robot_xml):
+  """A 2-step lag should make the effective position target 2 steps behind."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+  effort_limit = 20.0  # Set to saturation so the DC curve itself is the limit.
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=effort_limit,
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+      delay_min_lag=2,
+      delay_max_lag=4,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Zero position and velocity.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[0.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Small targets so torque demand stays below saturation_effort and is
+  # therefore distinguishable across steps.
+  targets = [
+    torch.tensor([[0.05]], device=device),  # demand = 5.0
+    torch.tensor([[0.10]], device=device),  # demand = 10.0
+    torch.tensor([[0.15]], device=device),  # demand = 15.0
+  ]
+
+  zero_effort = torch.zeros(1, 1, device=device)
+  for target in targets:
+    entity.set_joint_position_target(target)
+    entity.set_joint_velocity_target(zero_effort)
+    entity.set_joint_effort_target(zero_effort)
+    entity.write_data_to_sim()
+
+  # With lag=2 and three writes:
+  #   t=0: append target0 → return target0 (backfill)
+  #   t=1: append target1 → return target0 (clamped to available history)
+  #   t=2: append target2 → return target0 (lag=2, T-2 = 0)
+  # So the effective target should still be targets[0].
+  ctrl = sim.data.ctrl[0]
+  expected = kp * targets[0][0, 0]  # 5.0
+  assert torch.allclose(ctrl, torch.tensor([expected], device=device), atol=1e-4)