From eb3614d2fc4390a5774bf25cf1e2329c3055aafd Mon Sep 17 00:00:00 2001 From: DarksaCY Date: Tue, 8 Sep 2026 10:21:35 +0300 Subject: [PATCH 1/3] Pack an arm's state once, for every arm that states the same things `YamState` is `PackedState(n_joints)` in `state.py`: the joints, their velocities, the end effector pose and the status, in one float32 array. Five of the six state classes in the repository are that layout with a different joint count, and only the Franka's differs -- it carries a wrench. The offsets are the instance's, computed from the joint count, and `instantiation_params` names that count: shared memory rebuilds the payload on the far side from it, and a state that did not name it would come back the wrong size. The YAM reads it; `python -m positronic.drivers.roboarm.yam --fake` passes. Kinova, SO-101 and the MuJoCo sim can follow. --- positronic/drivers/roboarm/state.py | 53 +++++++++++++++++++ .../drivers/roboarm/tests/test_state.py | 48 +++++++++++++++++ positronic/drivers/roboarm/yam.py | 47 ++-------------- 3 files changed, 106 insertions(+), 42 deletions(-) create mode 100644 positronic/drivers/roboarm/state.py create mode 100644 positronic/drivers/roboarm/tests/test_state.py diff --git a/positronic/drivers/roboarm/state.py b/positronic/drivers/roboarm/state.py new file mode 100644 index 000000000..853e2433f --- /dev/null +++ b/positronic/drivers/roboarm/state.py @@ -0,0 +1,53 @@ +"""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.n_joints = n_joints + self._q = slice(0, n_joints) + self._dq = slice(n_joints, 2 * n_joints) + self._ee = slice(2 * n_joints, 2 * n_joints + 7) + self._status = 2 * n_joints + 7 + super().__init__(shape=(self._status + 1,), dtype=np.dtype(np.float32)) + + def instantiation_params(self) -> tuple[Any, ...]: + return (self.n_joints,) + + @property + def q(self) -> np.ndarray: + return self.array[self._q].copy() + + @property + def dq(self) -> np.ndarray: + return self.array[self._dq].copy() + + @property + def ee_pose(self) -> geom.Transform3D: + pose = self.array[self._ee].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])) + + def encode(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, status: RobotStatus) -> None: + self.array[self._q] = q + self.array[self._dq] = dq + self.array[self._ee.start : self._ee.start + 3] = ee_pose.translation + self.array[self._ee.start + 3 : self._ee.stop] = ee_pose.rotation.as_quat + self.array[self._status] = 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) From f37aee652dc04155892ac507db3488409d2d4caa Mon Sep 17 00:00:00 2001 From: DarksaCY Date: Tue, 8 Sep 2026 11:09:13 +0300 Subject: [PATCH 2/3] Read the joint count off the layout that carries it `n_joints` was a second copy of what the slices already say, and `instantiation_params` reads it: one that disagreed with the layout would rebuild a receiver of another size. It is a property of the layout now. `_q`, `_dq`, `_ee` and `_status` held slices and an index rather than the values they name; they say which. --- positronic/drivers/roboarm/state.py | 34 +++++++++++-------- .../drivers/roboarm/tests/test_state.py | 2 ++ 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/positronic/drivers/roboarm/state.py b/positronic/drivers/roboarm/state.py index 853e2433f..ad8b4e8b5 100644 --- a/positronic/drivers/roboarm/state.py +++ b/positronic/drivers/roboarm/state.py @@ -18,36 +18,40 @@ class PackedState(State, pimm.shared_memory.NumpySMAdapter): """ def __init__(self, n_joints: int): - self.n_joints = n_joints - self._q = slice(0, n_joints) - self._dq = slice(n_joints, 2 * n_joints) - self._ee = slice(2 * n_joints, 2 * n_joints + 7) - self._status = 2 * n_joints + 7 - super().__init__(shape=(self._status + 1,), dtype=np.dtype(np.float32)) + 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].copy() + return self.array[self._q_slice].copy() @property def dq(self) -> np.ndarray: - return self.array[self._dq].copy() + return self.array[self._dq_slice].copy() @property def ee_pose(self) -> geom.Transform3D: - pose = self.array[self._ee].copy() + 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])) + 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] = q - self.array[self._dq] = dq - self.array[self._ee.start : self._ee.start + 3] = ee_pose.translation - self.array[self._ee.start + 3 : self._ee.stop] = ee_pose.rotation.as_quat - self.array[self._status] = status.value + 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 index 0d3afbc6c..ff7cb86a9 100644 --- a/positronic/drivers/roboarm/tests/test_state.py +++ b/positronic/drivers/roboarm/tests/test_state.py @@ -34,6 +34,8 @@ def test_the_state_says_how_to_build_it_again(n_joints): assert rebuilt.n_joints == n_joints assert rebuilt.array.shape == state.array.shape + with pytest.raises(AttributeError): # the count is the layout's, and nothing else may set it + state.n_joints = n_joints + 1 def test_a_reading_is_a_copy_of_what_the_buffer_holds(): From dfefbcde45df2ef6f39a7ccd8cba86004e2e47b7 Mon Sep 17 00:00:00 2001 From: DarksaCY Date: Tue, 8 Sep 2026 11:10:40 +0300 Subject: [PATCH 3/3] Leave the read-only count to the type checker A test that assigns `n_joints` to prove it cannot be assigned needs a suppression to type-check, which says the same thing twice. --- positronic/drivers/roboarm/tests/test_state.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/positronic/drivers/roboarm/tests/test_state.py b/positronic/drivers/roboarm/tests/test_state.py index ff7cb86a9..0d3afbc6c 100644 --- a/positronic/drivers/roboarm/tests/test_state.py +++ b/positronic/drivers/roboarm/tests/test_state.py @@ -34,8 +34,6 @@ def test_the_state_says_how_to_build_it_again(n_joints): assert rebuilt.n_joints == n_joints assert rebuilt.array.shape == state.array.shape - with pytest.raises(AttributeError): # the count is the layout's, and nothing else may set it - state.n_joints = n_joints + 1 def test_a_reading_is_a_copy_of_what_the_buffer_holds():