Publish complete robot commands and owned state snapshots

i2rt/robots/motor_chain_robot.py

@@ -303,11 +303,13 @@ def _check_current_qpos_in_joint_limits(self, buffer_rad: float = 0.1) -> None:
 
     def get_robot_info(self) -> Dict[str, Any]:
         """Get the robot information, such as kp, kd, joint limits, gripper limits, etc."""
+        with self._command_lock:
+            kp, kd = self._kp.copy(), self._kd.copy()
         info: Dict[str, Any] = {
             "arm_type": self._arm_type,
             "gripper_type": self._gripper_type,
-            "kp": self._kp,
-            "kd": self._kd,
+            "kp": kp,
+            "kd": kd,
             "grav_comp_kd": self._grav_comp_kd,
             "coulomb_friction": self._coulomb_friction,
             "use_coulomb_friction": self.use_coulomb_friction,
@@ -393,7 +395,7 @@ def update(self) -> None:
                     max(self._gripper_limits),
                 )
                 self._last_gripper_command_qpos = joint_commands.pos[self._gripper_index]
-            self._update_joint_state(motor_torques, joint_commands)
+        self._update_joint_state(motor_torques, joint_commands)
 
     def _update_joint_state(
         self,
@@ -421,7 +423,8 @@ def _update_joint_state(
             kp=joint_commands.kp,
             kd=joint_commands.kd,
         )
-        self._joint_state = self._motor_state_to_joint_state(motor_state)
+        with self._state_lock:
+            self._joint_state = self._motor_state_to_joint_state(motor_state)
 
         with self._mcap_lock:
             if self._mcap_recorder is not None:
@@ -532,7 +535,8 @@ def num_dofs(self) -> int:
 
     def get_motor_torques(self) -> Optional[np.ndarray]:
         """Return the last computed motor torques (gravity comp + any command torques)."""
-        return self._last_motor_torques
+        with self._state_lock:
+            return self._last_motor_torques.copy() if self._last_motor_torques is not None else None
 
     def get_joint_pos(self) -> np.ndarray:
         """Get the current state of the leader robot, including the gripper in radian.
@@ -541,7 +545,7 @@ def get_joint_pos(self) -> np.ndarray:
             T: The current state of the leader robot.
         """
         with self._state_lock:
-            return self._joint_state.pos
+            return self._joint_state.pos.copy()
 
     def _clip_robot_joint_pos_command(self, pos: np.ndarray) -> np.ndarray:
         """Clip the robot joint pos command to the joint limits. Do not clip the gripper pos.
@@ -551,6 +555,7 @@ def _clip_robot_joint_pos_command(self, pos: np.ndarray) -> np.ndarray:
             np.ndarray: The clipped joint pos command.
         """
 
+        pos = np.array(pos, copy=True)
         if self._joint_limits is not None:
             if self._gripper_index is not None:
                 pos[: self._gripper_index] = np.clip(
@@ -572,8 +577,8 @@ def command_joint_pos(self, joint_pos: np.ndarray) -> None:
         with self._command_lock:
             self._commands = JointCommands.init_all_zero(len(self.motor_chain))
             self._commands.pos = self.remapper.to_robot_joint_pos_space(pos)
-            self._commands.kp = self._kp
-            self._commands.kd = self._kd
+            self._commands.kp = self._kp.copy()
+            self._commands.kd = self._kd.copy()
 
     def command_joint_state(self, joint_state: Dict[str, np.ndarray]) -> None:
         """Command the leader robot to a given state.
@@ -582,15 +587,15 @@ def command_joint_state(self, joint_state: Dict[str, np.ndarray]) -> None:
             joint_state (Dict[str, np.ndarray]): The state to command the leader robot to.
         """
         pos = self._clip_robot_joint_pos_command(joint_state["pos"])
-        vel = joint_state["vel"]
-        self._commands = JointCommands.init_all_zero(len(self.motor_chain))
-        kp = joint_state.get("kp", self._kp)
-        kd = joint_state.get("kd", self._kd)
+        vel = np.array(joint_state["vel"], copy=True)
         with self._command_lock:
-            self._commands.pos = self.remapper.to_robot_joint_pos_space(pos)
-            self._commands.vel = self.remapper.to_robot_joint_vel_space(vel)
-            self._commands.kp = kp
-            self._commands.kd = kd
+            self._commands = JointCommands(
+                torques=np.zeros(len(self.motor_chain)),
+                pos=self.remapper.to_robot_joint_pos_space(pos),
+                vel=self.remapper.to_robot_joint_vel_space(vel),
+                kp=np.array(joint_state.get("kp", self._kp), copy=True),
+                kd=np.array(joint_state.get("kd", self._kd), copy=True),
+            )
 
     def zero_torque_mode(self) -> None:
         logging.info(f"Entering zero_torque_mode for {self}")
@@ -612,22 +617,22 @@ def get_observations(self) -> Dict[str, np.ndarray]:
         with self._state_lock:
             if self._gripper_index is None:
                 result = {
-                    "joint_pos": self._joint_state.pos,
-                    "joint_vel": self._joint_state.vel,
-                    "joint_eff": self._joint_state.eff,
+                    "joint_pos": self._joint_state.pos.copy(),
+                    "joint_vel": self._joint_state.vel.copy(),
+                    "joint_eff": self._joint_state.eff.copy(),
                 }
             else:
                 result = {
-                    "joint_pos": self._joint_state.pos[: self._gripper_index],
-                    "joint_vel": self._joint_state.vel[: self._gripper_index],
-                    "joint_eff": self._joint_state.eff[: self._gripper_index],
+                    "joint_pos": self._joint_state.pos[: self._gripper_index].copy(),
+                    "joint_vel": self._joint_state.vel[: self._gripper_index].copy(),
+                    "joint_eff": self._joint_state.eff[: self._gripper_index].copy(),
                     "gripper_pos": np.array([self._joint_state.pos[self._gripper_index]]),
                     "gripper_vel": np.array([self._joint_state.vel[self._gripper_index]]),
                     "gripper_eff": np.array([self._joint_state.eff[self._gripper_index]]),
                 }
             if self.temp_record_flag:
-                result["temp_mos"] = self._joint_state.temp_mos
-                result["temp_rotor"] = self._joint_state.temp_rotor
+                result["temp_mos"] = self._joint_state.temp_mos.copy()
+                result["temp_rotor"] = self._joint_state.temp_rotor.copy()
             return result
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
@@ -637,7 +642,7 @@ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
     def move_joints(self, target_joint_positions: np.ndarray, time_interval_s: float = 2.0) -> None:
         """Move the robot to a given joint positions."""
         with self._state_lock:
-            current_pos = self._joint_state.pos
+            current_pos = self._joint_state.pos.copy()
         assert len(current_pos) == len(target_joint_positions)
         steps = 50  # 50 steps over time_interval_s
         for i in range(steps + 1):
@@ -659,8 +664,9 @@ def close(self) -> None:
 
     def update_kp_kd(self, kp: np.ndarray, kd: np.ndarray) -> None:
         assert kp.shape == self._kp.shape == kd.shape
-        self._kp = kp
-        self._kd = kd
+        with self._command_lock:
+            self._kp = kp.copy()
+            self._kd = kd.copy()
 
     def enter_gravity_comp_idle(self) -> None:
         """Reset active commands to gravity-comp idle.

i2rt/robots/sim_robot.py

@@ -71,7 +71,7 @@ def __init__(
             if jnt_lo != 0.0 or jnt_hi != 1.0:
                 self._gripper_qpos_range = (float(jnt_lo), float(jnt_hi))
 
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         self._qpos = np.zeros(n_dofs) if initial_qpos is None else np.array(initial_qpos, dtype=float)
         self._qvel = np.zeros(n_dofs)
 
@@ -271,9 +271,10 @@ def command_target_vel(self, joint_vel: np.ndarray) -> None:
             self._update_joint_state()
 
     def command_joint_state(self, joint_state: Dict[str, np.ndarray]) -> None:
-        self.command_joint_pos(joint_state["pos"])
-        if "vel" in joint_state:
-            self.command_target_vel(joint_state["vel"])
+        with self._lock:
+            self.command_joint_pos(joint_state["pos"])
+            if "vel" in joint_state:
+                self.command_target_vel(joint_state["vel"])
 
     def get_observations(self) -> Dict[str, np.ndarray]:
         with self._lock:

i2rt/robots/tests/test_robot_command_concurrency.py

@@ -0,0 +1,108 @@
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import pytest
+
+from i2rt.robots.get_robot import get_yam_robot
+from i2rt.robots.motor_chain_robot import JointCommands, JointStates, MotorChainRobot
+from i2rt.robots.utils import GripperType, JointMapper
+
+
+def _robot() -> MotorChainRobot:
+    robot = MotorChainRobot.__new__(MotorChainRobot)
+    robot.motor_chain = [None, None]
+    robot._command_lock = threading.Lock()
+    robot._state_lock = threading.Lock()
+    robot._commands = JointCommands.init_all_zero(2)
+    robot._kp, robot._kd = np.ones(2), np.ones(2)
+    robot._joint_limits = None
+    robot.remapper = JointMapper({}, 2)
+    robot._joint_state = JointStates([], np.ones(2), np.zeros(2), np.zeros(2), np.zeros(2), np.zeros(2), 0.0)
+    return robot
+
+
+def test_robot_commands_and_observations_own_their_arrays() -> None:
+    robot = _robot()
+    pos, vel, kp, kd = [np.ones(2) for _ in range(4)]
+    robot.update_kp_kd(kp, kd)
+    robot.command_joint_state({"pos": pos, "vel": vel})
+    for value in (pos, vel, kp, kd, robot.get_joint_pos()):
+        value[:] = 99.0
+    for value in (robot._commands.pos, robot._commands.vel, robot._commands.kp, robot._commands.kd):
+        np.testing.assert_array_equal(value, np.ones(2))
+    np.testing.assert_array_equal(robot._joint_state.pos, np.ones(2))
+
+
+def test_joint_state_command_is_not_replaced_before_taking_the_lock() -> None:
+    robot = _robot()
+    lock, attempted = threading.Lock(), threading.Event()
+
+    class ObservedLock:
+        def __enter__(self) -> None:
+            attempted.set()
+            lock.acquire()
+
+        def __exit__(self, *_args: object) -> None:
+            lock.release()
+
+    robot._command_lock = ObservedLock()
+    previous = robot._commands
+    with ThreadPoolExecutor(1) as pool:
+        with lock:
+            future = pool.submit(robot.command_joint_state, {"pos": np.ones(2), "vel": np.zeros(2)})
+            assert attempted.wait(1.0)
+            assert robot._commands is previous
+        future.result(timeout=1.0)
+    np.testing.assert_array_equal(robot._commands.pos, np.ones(2))
+
+
+def test_robot_readers_are_not_blocked_by_motor_io() -> None:
+    robot = _robot()
+    robot._compute_gravity_compensation = lambda _state: np.zeros(2)
+    robot.use_coulomb_friction = False
+    robot.gravity_comp_factor = np.ones(2)
+    robot._clip_motor_torque = np.inf
+    robot._gripper_index = None
+    entered, release = threading.Event(), threading.Event()
+
+    def motor_io(*_args: object) -> None:
+        entered.set()
+        assert release.wait(2.0)
+
+    robot._update_joint_state = motor_io
+    with ThreadPoolExecutor(2) as pool:
+        writer = pool.submit(robot.update)
+        try:
+            assert entered.wait(1.0)
+            np.testing.assert_array_equal(pool.submit(robot.get_joint_pos).result(timeout=0.2), np.ones(2))
+        finally:
+            release.set()
+        writer.result(timeout=1.0)
+
+
+def test_sim_joint_state_is_published_atomically(monkeypatch: pytest.MonkeyPatch) -> None:
+    robot = get_yam_robot(gripper_type=GripperType.NO_GRIPPER, sim=True)
+    entered, release = threading.Event(), threading.Event()
+    original = robot.command_joint_pos
+
+    def set_position(pos: np.ndarray) -> None:
+        original(pos)
+        entered.set()
+        assert release.wait(2.0)
+
+    monkeypatch.setattr(robot, "command_joint_pos", set_position)
+    command = {"pos": np.full(6, 0.2), "vel": np.full(6, 0.3)}
+    with ThreadPoolExecutor(2) as pool:
+        writer = pool.submit(robot.command_joint_state, command)
+        try:
+            assert entered.wait(1.0)
+            reader = pool.submit(robot.get_joint_state)
+            with pytest.raises(TimeoutError):
+                reader.result(timeout=0.1)
+        finally:
+            release.set()
+            robot.close()
+        writer.result(timeout=1.0)
+        for key, value in reader.result(timeout=1.0).items():
+            np.testing.assert_array_equal(value, command[key])