Fix broadcasting bug

Author: oleflb

tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/joint_actuator_info.py

@@ -10,9 +10,9 @@ class JointActuatorInfo:
 
     def __init__(self, name: str, model: MjModel) -> None:
         self.name = name
-        self.qpos_addr = model.joint(name).qposadr
-        self.qvel_addr = model.joint(name).dofadr
-        self.qacc_addr = model.joint(name).dofadr
+        self.qpos_addr = model.joint(name).qposadr.item()
+        self.qvel_addr = model.joint(name).dofadr.item()
+        self.qacc_addr = model.joint(name).dofadr.item()
         self.qfrc_actuator_addr = model.actuator(name).id
 
 

tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/low_command.py

@@ -1,19 +1,29 @@
 import numpy as np
 from mujoco import MjData, MjModel
 from mujoco_rust_server.booster_types import LowCommand
+from numpy.typing import NDArray
 
 from mujoco_simulator.joint_actuator_info import JointActuatorInfo
 
 
 class RobotPositionControl:
+    model: MjModel
+    qpos_indices: NDArray
+    qvel_indices: NDArray
+    actuator_indices: NDArray
+
     def __init__(
         self, model: MjModel, actuator_info: list[JointActuatorInfo]
     ) -> None:
         self.model = model
-        self.qpos_indices = np.array([info.qpos_addr for info in actuator_info])
-        self.qvel_indices = np.array([info.qvel_addr for info in actuator_info])
+        self.qpos_indices = np.array(
+            [info.qpos_addr for info in actuator_info], dtype=np.uint32
+        )
+        self.qvel_indices = np.array(
+            [info.qvel_addr for info in actuator_info], dtype=np.uint32
+        )
         self.actuator_indices = np.array(
-            [info.qfrc_actuator_addr for info in actuator_info]
+            [info.qfrc_actuator_addr for info in actuator_info], dtype=np.uint32
         )
 
     def apply_control(
@@ -33,13 +43,17 @@ def apply_control(
         current_velocity = data.qvel[self.qvel_indices]
 
         # TODO(oleflb): booster supposedly clips position to joint limits first
-        control_torque = np.clip(
-            kp * (q - current_position) + kd * (dq - current_velocity) + tau,
-            self.model.actuator_ctrlrange[:, 0],
-            self.model.actuator_ctrlrange[:, 1],
+        desired = (
+            kp * (q - current_position) + kd * (dq - current_velocity) + tau
         )
+
+        ctrl_min = self.model.actuator_ctrlrange[self.actuator_indices, 0]
+        ctrl_max = self.model.actuator_ctrlrange[self.actuator_indices, 1]
+        control_torque = np.clip(desired, ctrl_min, ctrl_max)
+
+        # Ensure existing control values are used as floats when smoothing
+        current_ctrl = data.ctrl[self.actuator_indices]
         smoothed_control_torque = (
-            weight * control_torque
-            + (1 - weight) * data.ctrl[self.actuator_indices]
+            weight * control_torque + (1 - weight) * current_ctrl
         )
         data.ctrl[self.actuator_indices] = smoothed_control_torque