diff --git a/positronic/drivers/roboarm/state.py b/positronic/drivers/roboarm/state.py new file mode 100644 index 000000000..ad8b4e8b5 --- /dev/null +++ b/positronic/drivers/roboarm/state.py @@ -0,0 +1,57 @@ +"""The arm state a driver ships to the rest of a run.""" + +from typing import Any + +import numpy as np + +import pimm +from positronic import geom + +from . import RobotStatus, State + + +class PackedState(State, pimm.shared_memory.NumpySMAdapter): + """An arm's joints, their velocities, the end effector pose and the arm's status, in one float32 array. + + Shared memory carries a fixed-size payload, so the state is packed rather than shipped as fields. Every + arm this drives states the same four things and differs only in how many joints it has. + """ + + def __init__(self, n_joints: int): + self._q_slice = slice(0, n_joints) + self._dq_slice = slice(n_joints, 2 * n_joints) + self._ee_slice = slice(2 * n_joints, 2 * n_joints + 7) + self._status_index = 2 * n_joints + 7 + super().__init__(shape=(self._status_index + 1,), dtype=np.dtype(np.float32)) + + @property + def n_joints(self) -> int: + """How many joints the layout carries. Read off the layout, so the two cannot disagree.""" + return self._q_slice.stop - self._q_slice.start + + def instantiation_params(self) -> tuple[Any, ...]: + return (self.n_joints,) + + @property + def q(self) -> np.ndarray: + return self.array[self._q_slice].copy() + + @property + def dq(self) -> np.ndarray: + return self.array[self._dq_slice].copy() + + @property + def ee_pose(self) -> geom.Transform3D: + pose = self.array[self._ee_slice].copy() + return geom.Transform3D(pose[:3], geom.Rotation.from_quat(pose[3:7])) + + @property + def status(self) -> RobotStatus: + return RobotStatus(int(self.array[self._status_index])) + + def encode(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, status: RobotStatus) -> None: + self.array[self._q_slice] = q + self.array[self._dq_slice] = dq + self.array[self._ee_slice.start : self._ee_slice.start + 3] = ee_pose.translation + self.array[self._ee_slice.start + 3 : self._ee_slice.stop] = ee_pose.rotation.as_quat + self.array[self._status_index] = status.value diff --git a/positronic/drivers/roboarm/tests/test_state.py b/positronic/drivers/roboarm/tests/test_state.py new file mode 100644 index 000000000..0d3afbc6c --- /dev/null +++ b/positronic/drivers/roboarm/tests/test_state.py @@ -0,0 +1,48 @@ +"""What ``PackedState`` carries, and what it must carry across a process boundary.""" + +import numpy as np +import pytest + +from positronic import geom +from positronic.drivers.roboarm import RobotStatus +from positronic.drivers.roboarm.state import PackedState + + +@pytest.mark.parametrize('n_joints', [5, 6, 7]) +def test_what_goes_in_comes_back_out(n_joints): + state = PackedState(n_joints) + q = np.arange(n_joints, dtype=np.float64) * 0.1 + dq = np.arange(n_joints, dtype=np.float64) * -0.01 + pose = geom.Transform3D(np.array([0.3, -0.2, 0.5]), geom.Rotation.from_rotvec(np.array([0.1, 0.2, 0.3]))) + + state.encode(q, dq, pose, RobotStatus.BUSY) + + np.testing.assert_allclose(state.q, q, atol=1e-6) + np.testing.assert_allclose(state.dq, dq, atol=1e-6) + np.testing.assert_allclose(state.ee_pose.translation, pose.translation, atol=1e-6) + np.testing.assert_allclose(state.ee_pose.rotation.as_quat, pose.rotation.as_quat, atol=1e-6) + assert state.status is RobotStatus.BUSY + + +@pytest.mark.parametrize('n_joints', [5, 6, 7]) +def test_the_state_says_how_to_build_it_again(n_joints): + """Shared memory rebuilds the payload on the far side from ``instantiation_params``, so a state that + does not name its joint count comes back the wrong size and reads another arm's numbers.""" + state = PackedState(n_joints) + + rebuilt = PackedState(*state.instantiation_params()) + + assert rebuilt.n_joints == n_joints + assert rebuilt.array.shape == state.array.shape + + +def test_a_reading_is_a_copy_of_what_the_buffer_holds(): + """The buffer is written again every tick, so a reader that kept a view would watch its own reading + change under it.""" + state = PackedState(6) + state.encode(np.zeros(6), np.zeros(6), geom.Transform3D(), RobotStatus.AVAILABLE) + q = state.q + + state.encode(np.ones(6), np.zeros(6), geom.Transform3D(), RobotStatus.AVAILABLE) + + np.testing.assert_allclose(q, np.zeros(6)) diff --git a/positronic/drivers/roboarm/yam.py b/positronic/drivers/roboarm/yam.py index 6611068ee..c5cb5708e 100644 --- a/positronic/drivers/roboarm/yam.py +++ b/positronic/drivers/roboarm/yam.py @@ -28,9 +28,10 @@ from positronic.drivers.utils import DriverRun, MoveAbandoned, MoveStatus, log_failure from positronic.utils import package_assets_path -from . import RobotStatus, State, command +from . import RobotStatus, command from .ik import qpos_from_site_pose from .models import DEFAULT_FRAME +from .state import PackedState # i2rt lives in the `yam` extra, which the type-check environment does not install. with vendor_import('i2rt', 'YAM support', hint='Re-run with the yam extra:\n uv run --locked --extra yam ...\n'): @@ -64,44 +65,6 @@ def _connect(channel: str, sim: bool): return get_yam_robot(channel, gripper_type=GripperType.LINEAR_4310, zero_gravity_mode=False, sim=sim) -class YamState(State, pimm.shared_memory.NumpySMAdapter): - Q_OFFSET = 0 - DQ_OFFSET = Q_OFFSET + 6 - EE_POSE_OFFSET = DQ_OFFSET + 6 - STATUS_OFFSET = EE_POSE_OFFSET + 7 - TOTAL = STATUS_OFFSET + 1 - - def __init__(self): - super().__init__(shape=(YamState.TOTAL,), dtype=np.dtype(np.float32)) - - def instantiation_params(self) -> tuple[Any, ...]: - return () - - @property - def q(self) -> np.ndarray: - return self.array[YamState.Q_OFFSET : YamState.Q_OFFSET + 6].copy() - - @property - def dq(self) -> np.ndarray: - return self.array[YamState.DQ_OFFSET : YamState.DQ_OFFSET + 6].copy() - - @property - def ee_pose(self) -> geom.Transform3D: - pose = self.array[YamState.EE_POSE_OFFSET : YamState.EE_POSE_OFFSET + 7].copy() - return geom.Transform3D(pose[:3], geom.Rotation.from_quat(pose[3:7])) - - @property - def status(self) -> RobotStatus: - return RobotStatus(int(self.array[YamState.STATUS_OFFSET])) - - def encode(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, status: RobotStatus): - self.array[YamState.Q_OFFSET : YamState.Q_OFFSET + 6] = q - self.array[YamState.DQ_OFFSET : YamState.DQ_OFFSET + 6] = dq - self.array[YamState.EE_POSE_OFFSET : YamState.EE_POSE_OFFSET + 3] = ee_pose.translation - self.array[YamState.EE_POSE_OFFSET + 3 : YamState.EE_POSE_OFFSET + 7] = ee_pose.rotation.as_quat - self.array[YamState.STATUS_OFFSET] = status.value - - class _Kinematics: """FK/IK on the vendored YAM MJCF at ``DEFAULT_FRAME``, in the arm-base frame. @@ -169,7 +132,7 @@ def __init__( vendor: Any, sync_move: pimm.calls.ControlSystemHandler[command.CommandType, None], async_move: pimm.SignalReceiver[command.CommandType], - out: pimm.SignalEmitter[YamState], + out: pimm.SignalEmitter[PackedState], grip_out: pimm.SignalEmitter[float], base_pose: geom.Transform3D, should_stop: pimm.SignalReceiver, @@ -179,7 +142,7 @@ def __init__( self.vendor = vendor self.out = out self.grip_out = grip_out - self.state = YamState() + self.state = PackedState(len(_JOINT_NAMES)) self._base_pose = base_pose self._kin = _Kinematics() @@ -350,7 +313,7 @@ def __init__( self.commands = pimm.ControlSystemReceiver[command.CommandType](self) self.sync_move = pimm.calls.ControlSystemHandler[command.CommandType, None](self) self.target_grip = pimm.ControlSystemReceiver[float](self) - self.state = pimm.ControlSystemEmitter[YamState](self) + self.state = pimm.ControlSystemEmitter[PackedState](self) self.grip = pimm.ControlSystemEmitter[float](self) self.robot_meta = pimm.ControlSystemEmitter[dict[str, Any]](self)