diff --git a/i2rt/robots/motor_chain_robot.py b/i2rt/robots/motor_chain_robot.py index 18da115..61e99d9 100644 --- a/i2rt/robots/motor_chain_robot.py +++ b/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. diff --git a/i2rt/robots/sim_robot.py b/i2rt/robots/sim_robot.py index 470e3f1..3e752cb 100644 --- a/i2rt/robots/sim_robot.py +++ b/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: diff --git a/i2rt/robots/tests/test_robot_command_concurrency.py b/i2rt/robots/tests/test_robot_command_concurrency.py new file mode 100644 index 0000000..454fe0f --- /dev/null +++ b/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])