Pack an arm's state once, for every arm that states the same things

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

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))

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)