From 13095a9c90f398a4443f988ccdbc109915097763 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 19 Jul 2026 20:34:44 +0300 Subject: [PATCH 01/23] Add `ChangeEEFrame(to=)` codec and `droid_eef` site for per-policy EE frames Convert poses between the robot's canonical `control_frame` and a policy's own EE frame (DROID's `droid_eef` = gripper base composed with `EEF_OFFSET_ROT`). The codec runs client-side: the harness injects the robot model into the local obs and `RemoteSession` keeps it off the wire. Declare RoboLab's `robot_meta` `control_frame` as `droid_eef` truthfully and drop the `#469` HACK. --- positronic/cfg/codecs.py | 27 +++++- positronic/drivers/roboarm/ik.py | 56 ++++++++--- positronic/drivers/roboarm/models.py | 24 +++-- positronic/drivers/roboarm/tests/test_ik.py | 45 ++++++++- .../offboard/tests/test_remote_policy.py | 7 ++ positronic/policy/codec.py | 71 +++++++++++++- positronic/policy/harness.py | 7 ++ positronic/policy/remote.py | 6 +- .../policy/tests/test_change_ee_frame.py | 94 +++++++++++++++++++ positronic/simulator/robolab/launcher.py | 10 +- 10 files changed, 316 insertions(+), 31 deletions(-) create mode 100644 positronic/policy/tests/test_change_ee_frame.py diff --git a/positronic/cfg/codecs.py b/positronic/cfg/codecs.py index 9d84bff18..b46335cac 100644 --- a/positronic/cfg/codecs.py +++ b/positronic/cfg/codecs.py @@ -3,10 +3,16 @@ import configuronic as cfn from positronic import geom +from positronic.policy.codec import ChangeEEFrame from positronic.policy.observation import ObservationCodec RotRep = geom.Rotation.Representation +# The border codec that converts poses into a policy's own EE frame (e.g. DROID's ``droid_eef``). Applied +# client-side at serving via ``remote(codec=change_ee_frame.override(to=...))``; the ``compose`` ``ee_frame`` +# param below folds it into the training pipeline so the recorded data is encoded in the same frame. +change_ee_frame = cfn.Config(ChangeEEFrame) + @cfn.config() def general_obs( @@ -43,15 +49,26 @@ def general_obs( ) -@cfn.config(fps=15.0, horizon=None, binarize_grip=None, flip_grip=False) -def compose(obs, action, fps: float, horizon: float | None, binarize_grip: tuple[str, ...] | None, flip_grip: bool): +@cfn.config(fps=15.0, horizon=None, binarize_grip=None, flip_grip=False, ee_frame=None) +def compose( + obs, + action, + fps: float, + horizon: float | None, + binarize_grip: tuple[str, ...] | None, + flip_grip: bool, + ee_frame: str | None, +): """Compose observation and action codecs with timing and optional grip binarization. - ``flip_grip`` serves checkpoints that speak the inverted grip convention (see ``FlipGrip``). + ``flip_grip`` serves checkpoints that speak the inverted grip convention (see ``FlipGrip``). ``ee_frame`` + converts poses into the policy's own EE frame (see ``ChangeEEFrame``); folded in for training so the recorded + data is encoded in that frame. At serving the frame conversion runs client-side instead (``change_ee_frame``), + so leave ``ee_frame`` unset on the server's frame-agnostic codec. Layout:: - [ActionHorizon] | ActionTimestamp | [BinarizeGripTraining | BinarizeGripInference] | [FlipGrip] | obs & action + [ActionHorizon] | ActionTimestamp | [BinarizeGrip*] | [FlipGrip] | [ChangeEEFrame] | obs & action """ from positronic.policy.codec import ( ActionHorizon, @@ -62,6 +79,8 @@ def compose(obs, action, fps: float, horizon: float | None, binarize_grip: tuple ) result = obs & action + if ee_frame is not None: + result = ChangeEEFrame(to=ee_frame) | result if flip_grip: result = FlipGrip() | result if binarize_grip: diff --git a/positronic/drivers/roboarm/ik.py b/positronic/drivers/roboarm/ik.py index 425b72487..f147c378b 100644 --- a/positronic/drivers/roboarm/ik.py +++ b/positronic/drivers/roboarm/ik.py @@ -8,21 +8,36 @@ """ import xml.etree.ElementTree as ET +from functools import lru_cache import mujoco as mj import numpy as np from scipy.optimize import lsq_linear from scipy.spatial.transform import Rotation as ScipyRotation +from positronic import geom from positronic.dataset import transforms -def _prepare_spec(urdf_xml, control_frame): - """Parse URDF or MJCF into an MjSpec, stripping meshes and resolving the control frame site. +def _ensure_site(spec, frame): + """Ensure ``frame`` is a site in ``spec``, adding one at the body origin when it names a body. - The control frame must exist in the model as a site or body. For bodies (e.g. real URDF - with ``end_effector`` link baked in by positronic-franka), a site is added at its origin. + Frames resolve against the model as a site or a body (e.g. the ``end_effector`` link the real URDF bakes in, + or the ``droid_eef`` frame graft) — the single registry every frame lookup uses. """ + all_sites = {s.name for b in spec.bodies for s in b.sites} + if frame in all_sites: + return + body_names = {b.name for b in spec.bodies} + if frame in body_names: + site = spec.body(frame).add_site() + site.name = frame + return + raise ValueError(f'Frame {frame!r} not found as site or body in model') + + +def _prepare_spec(urdf_xml, control_frame): + """Parse URDF or MJCF into an MjSpec, stripping meshes and resolving the control frame site.""" root = ET.fromstring(urdf_xml) if root.tag == 'robot': for link in root.findall('.//link'): @@ -30,18 +45,33 @@ def _prepare_spec(urdf_xml, control_frame): link.remove(elem) urdf_xml = ET.tostring(root, encoding='unicode') spec = mj.MjSpec.from_string(urdf_xml) + _ensure_site(spec, control_frame) + return spec - all_sites = {s.name for b in spec.bodies for s in b.sites} - if control_frame in all_sites: - return spec - body_names = {b.name for b in spec.bodies} - if control_frame in body_names: - site = spec.body(control_frame).add_site() - site.name = control_frame - return spec +def _site_transform(data, site_id): + """The world pose of a site as a ``Transform3D``.""" + rotation = geom.Rotation.from_rotation_matrix(data.site_xmat[site_id].reshape(3, 3)) + return geom.Transform3D(data.site_xpos[site_id].copy(), rotation) + - raise ValueError(f'Control frame {control_frame!r} not found as site or body in model') +@lru_cache(maxsize=8) +def frame_transform(urdf_xml, from_frame, to_frame): + """The rigid transform expressing ``to_frame`` relative to ``from_frame`` in a robot model. + + A pose measured in ``from_frame`` (e.g. the recorded ``ee_pose`` at ``control_frame``) composes to + ``to_frame`` via ``pose * frame_transform(...)``. Both frames must be rigidly connected (fixed joints) for the + result to be config-independent, so it is read from a single forward pass at the zero configuration; frames + resolve as sites or bodies, the same registry ``_prepare_spec`` uses for the control frame. + """ + spec = _prepare_spec(urdf_xml, from_frame) + _ensure_site(spec, to_frame) + model = spec.compile() + data = mj.MjData(model) + mj.mj_forward(model, data) + from_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, from_frame)) + to_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, to_frame)) + return from_pose.inv * to_pose def _parse_target(target_ee_pose_vec): diff --git a/positronic/drivers/roboarm/models.py b/positronic/drivers/roboarm/models.py index 7f087ab27..e235eb5bd 100644 --- a/positronic/drivers/roboarm/models.py +++ b/positronic/drivers/roboarm/models.py @@ -6,11 +6,20 @@ from functools import lru_cache from pathlib import Path +from positronic import geom + _FLANGE_LINK = 'link8' # Seat the gripper on the flange, rotated about the approach axis to match the real 2F-85 coupler # (a +45deg Z, i.e. 90deg off the franka ``end_effector`` frame). _2F85_MOUNT_RPY = '0 0 0.7853981634' +# The DROID/RoboLab end-effector frame ``droid_eef``: the gripper base rotated by a fixed offset. RoboLab reports +# and accepts Cartesian poses in this frame (``eef_frame`` = its ``Robotiq_2F_85/base_link`` ∘ ``EEF_OFFSET_ROT``, +# a pure rotation). ``EEF_OFFSET_ROT`` lives in RoboLab's ``robolab/robots/droid.py`` as a wxyz quaternion; its +# euler form is the URDF ``rpy`` of the ``droid_eef`` frame graft below. +_EEF_OFFSET_ROT = (0.5, -0.5, 0.5, -0.5) +_DROID_EEF_RPY = ' '.join(f'{angle:.12g}' for angle in geom.Rotation.from_quat(_EEF_OFFSET_ROT).as_euler) + def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: """One 2F-85 finger as URDF rows. ``sign`` mirrors the y-offsets and ``base_rpy`` (180deg Z on the @@ -62,15 +71,17 @@ def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: ] -# Rows: (link, parent, joint | None, origin xyz, origin rpy, axis | None, mesh, visual xyz | None). +# Rows: (link, parent, joint | None, origin xyz, origin rpy, axis | None, mesh | None, visual xyz | None). # A row with an axis is a revolute joint whose axis sign sets its closing direction, so one positive # ``grip`` drives the whole 4-bar: driver/spring_link swing the finger in (+X), coupler/follower -# counter-rotate (-X) to keep the pad parallel. Rows without an axis are fixed. +# counter-rotate (-X) to keep the pad parallel. Rows without an axis are fixed. A row with ``mesh`` None +# is a pure frame (no visual) — ``droid_eef`` is such a frame, the DROID control frame offset from the base. _ROBOTIQ_2F85 = [ ('gripper_base_mount', _FLANGE_LINK, None, '0 0 0.007', _2F85_MOUNT_RPY, None, 'base_mount.stl', None), ('gripper_base', 'gripper_base_mount', None, '0 0 0.0038', '0 0 -1.5707963268', None, 'base.stl', None), *_2f85_finger('right', 1, '0 0 0'), *_2f85_finger('left', -1, '0 0 3.1415926536'), + ('droid_eef', 'gripper_base', None, '0 0 0', _DROID_EEF_RPY, None, None, None), ] _ROBOTIQ_2F85_JOINTS = [row[2] for row in _ROBOTIQ_2F85 if row[2]] @@ -85,10 +96,11 @@ def _build_2f85_elements() -> list[ET.Element]: inertial = ET.SubElement(link_el, 'inertial') ET.SubElement(inertial, 'mass', value='0.01') ET.SubElement(inertial, 'inertia', ixx='1e-5', iyy='1e-5', izz='1e-5', ixy='0', ixz='0', iyz='0') - visual = ET.SubElement(link_el, 'visual') - if visual_xyz is not None: - ET.SubElement(visual, 'origin', xyz=visual_xyz, rpy='0 0 0') - ET.SubElement(ET.SubElement(visual, 'geometry'), 'mesh', filename=mesh) + if mesh is not None: + visual = ET.SubElement(link_el, 'visual') + if visual_xyz is not None: + ET.SubElement(visual, 'origin', xyz=visual_xyz, rpy='0 0 0') + ET.SubElement(ET.SubElement(visual, 'geometry'), 'mesh', filename=mesh) joint_el = ET.Element('joint', name=joint or f'{link}_fixed', type='revolute' if axis else 'fixed') ET.SubElement(joint_el, 'origin', xyz=xyz, rpy=rpy) ET.SubElement(joint_el, 'parent', link=parent) diff --git a/positronic/drivers/roboarm/tests/test_ik.py b/positronic/drivers/roboarm/tests/test_ik.py index 23ebc6b2f..489723ce3 100644 --- a/positronic/drivers/roboarm/tests/test_ik.py +++ b/positronic/drivers/roboarm/tests/test_ik.py @@ -5,9 +5,18 @@ import numpy as np import pytest +from positronic import geom from positronic.dataset.episode import EpisodeContainer from positronic.dataset.tests.utils import DummySignal -from positronic.drivers.roboarm.ik import DLSIKSolver, DLSIKSolverWithLimits, LMIKSolver, ik_joints_from_episode +from positronic.drivers.roboarm.ik import ( + DLSIKSolver, + DLSIKSolverWithLimits, + LMIKSolver, + _prepare_spec, + frame_transform, + ik_joints_from_episode, +) +from positronic.drivers.roboarm.models import bundled_franka_model from positronic.utils import package_assets_path URDF = Path(package_assets_path('assets/mujoco/panda_ik.xml')).read_text() @@ -107,6 +116,40 @@ def test_ik_joints_from_episode(): np.testing.assert_allclose(reconstructed_pose[:3], ee_poses[i, :3], atol=1e-3) +def _fk_site(urdf_xml, q, frame): + """FK a named frame (site or body) to [tx,ty,tz,w,x,y,z] through the ik spec preparation.""" + model = _prepare_spec(urdf_xml, frame).compile() + data = mj.MjData(model) + qpos_ids = [model.joint(n).qposadr.item() for n in JOINT_NAMES] + data.qpos[qpos_ids] = q + mj.mj_forward(model, data) + sid = mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, frame) + quat = np.empty(4) + mj.mju_mat2Quat(quat, data.site_xmat[sid]) + return np.concatenate([data.site_xpos[sid].copy(), quat]) + + +def test_frame_transform_reproduces_droid_eef_across_configs(): + """``frame_transform`` yields one config-independent transform such that the canonical ``end_effector`` pose + composed with it reproduces the ``droid_eef`` site pose at every joint configuration — the transform the + ``ChangeEEFrame`` codec applies to observations.""" + urdf = bundled_franka_model()['urdf'] + transform = frame_transform(urdf, 'end_effector', 'droid_eef') + for q in TEST_CONFIGS: + ee = geom.Transform3D.from_vector(_fk_site(urdf, q, 'end_effector'), geom.Rotation.Representation.QUAT) + want = _fk_site(urdf, q, 'droid_eef') + got = (ee * transform).as_vector(geom.Rotation.Representation.QUAT) + np.testing.assert_allclose(got[:3], want[:3], atol=1e-9) + q_diff = min(np.linalg.norm(got[3:] - want[3:]), np.linalg.norm(got[3:] + want[3:])) + assert q_diff < 1e-9, f'rotation mismatch: {q_diff}' + + +def test_frame_transform_identity_when_frames_match(): + transform = frame_transform(bundled_franka_model()['urdf'], 'end_effector', 'end_effector') + np.testing.assert_allclose(transform.translation, 0.0, atol=1e-12) + assert transform.rotation == geom.Rotation.identity + + @pytest.mark.parametrize('solver_cls', [DLSIKSolver, DLSIKSolverWithLimits]) def test_pickle_roundtrip(solver_cls): solver = solver_cls(URDF, JOINT_NAMES, CONTROL_FRAME) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 1299ae37c..054774c9b 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -56,6 +56,13 @@ def test_no_resize_without_server_sizes_or_fallback(self): result = session._prepare_obs({'cam': img}) assert result['cam'] is img + def test_drops_client_only_model_keys(self): + """The harness injects ``urdf``/``control_frame`` for client-side frame codecs; they must never wire out.""" + session = RemoteSession(_mock_ws_session(), resize=None) + result = session._prepare_obs({'state': np.array([1.0]), 'urdf': '', 'control_frame': 'end_effector'}) + assert 'urdf' not in result and 'control_frame' not in result + np.testing.assert_array_equal(result['state'], np.array([1.0])) + def test_normalizes_list_to_tuple(self): """Wire format (msgpack) turns tuples into lists — must normalize.""" session = RemoteSession(_mock_ws_session({'image_sizes': [64, 48]}), resize=None) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 1b3d8207c..083171536 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -9,15 +9,27 @@ (e.g. observation encoder & action decoder). """ +from functools import partial from typing import Any, final import numpy as np -from positronic.dataset.transforms import Elementwise +from positronic import geom +from positronic.dataset.transforms import Elementwise, lazy_sequence from positronic.dataset.transforms.episode import Derive, EpisodeTransform, Group, Identity +from positronic.drivers.roboarm import command +from positronic.drivers.roboarm.ik import frame_transform from positronic.policy.base import DelegatingSession, PolicyWrapper, Session, _Pipeline from positronic.utils import merge_dicts +_QUAT = geom.Rotation.Representation.QUAT + + +def _to_policy_frame(pose_vec, transform: geom.Transform3D) -> np.ndarray: + """Recompose a canonical ``[tx,ty,tz,qw,qx,qy,qz]`` pose into the policy frame via ``pose * transform``.""" + pose = geom.Transform3D.from_vector(np.asarray(pose_vec, dtype=np.float64), _QUAT) + return (pose * transform).as_vector(_QUAT) + def lerobot_state(dim: int, names: list[str] | None = None) -> dict[str, Any]: """LeRobot feature descriptor for a state vector.""" @@ -343,3 +355,60 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: if 'target_grip' in data: data['target_grip'] = 1.0 - data['target_grip'] return data + + +class ChangeEEFrame(Codec): + """Cross the border between the robot's canonical EE frame and a policy's own EE frame. + + A policy trained to speak a different end-effector frame (e.g. DROID's ``droid_eef``) is served on a rig whose + canonical frame is whatever its model declares via ``control_frame``. This codec converts the pose at that + boundary by pure composition with ``T`` = the fixed transform from ``control_frame`` to ``to`` in the episode's + model: observations go ``pose * T`` into the policy frame, actions come back ``pose * T⁻¹``. ``T`` is read from + the model the same way IK reads it — ``control_frame`` and ``urdf`` from episode statics at training, from the + obs at inference — so a dataset mixing embodiments just yields a different ``T`` per episode. A pipeline without + this codec is unchanged; ``to == control_frame`` makes ``T`` the identity. + + Runs client-side, reading the robot model the harness injects into the local obs. That model is client-only and + never reaches a frame-agnostic server: ``RemoteSession`` drops it at the wire boundary. + + Compose to the left of the observation/action codecs:: + + ChangeEEFrame(to='droid_eef') | ee + """ + + def __init__(self, to: str, ee_pose_key: str = 'robot_state.ee_pose', command_pose_key: str = 'robot_command.pose'): + self._to = to + self._ee_pose_key = ee_pose_key + self._command_pose_key = command_pose_key + + def _transform(self, source) -> geom.Transform3D: + """``T`` from the ``urdf``/``control_frame`` carried by an obs dict (inference) or episode (training).""" + return frame_transform(source['urdf'], source['control_frame'], self._to) + + def encode(self, data): + return {**data, self._ee_pose_key: _to_policy_frame(data[self._ee_pose_key], self._transform(data))} + + def _decode_single(self, data: dict, context: dict | None) -> dict: + cmd = data.get('robot_command') + if not isinstance(cmd, command.CartesianPosition): + return data + return {**data, 'robot_command': command.CartesianPosition(pose=cmd.pose * self._transform(context).inv)} + + def _derive_pose(self, key: str): + def derive(episode): + transform = self._transform(episode) + return Elementwise(episode[key], lazy_sequence(partial(_to_policy_frame, transform=transform))) + + return derive + + @property + def training_encoder(self) -> EpisodeTransform: + poses = { + self._ee_pose_key: self._derive_pose(self._ee_pose_key), + self._command_pose_key: self._derive_pose(self._command_pose_key), + } + return Group(Derive(**poses), Identity()) + + @property + def meta(self): + return {'ee_frame': self._to} diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index a7d3d5d27..80c3f67be 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -291,6 +291,13 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['wall_time_ns'] = time.time_ns() inputs['obs_time_ns'] = clock.now_ns() inputs.update(self.context) + model = self.robot_meta_in.value + if 'control_frame' in model: + # The robot model backs any codec that resolves an EE frame against it (``ChangeEEFrame``), mirroring + # how training reads ``urdf``/``control_frame`` from episode statics. It is client-local — ``RemoteSession`` + # drops it before any observation crosses the wire, so it never reaches an inference server. + inputs['urdf'] = model['urdf'] + inputs['control_frame'] = model['control_frame'] inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it return inputs diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 60f73aa02..4d3c1e631 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -10,6 +10,10 @@ from .base import Policy, Session +# Client-only robot-model metadata the harness puts in the observation for client-side frame codecs +# (``ChangeEEFrame``); it is never a model input, so it is dropped before an observation crosses the wire. +_LOCAL_ONLY_KEYS = frozenset({'urdf', 'control_frame'}) + class RemoteSession(Session): """Per-episode session that forwards observations to a remote inference server.""" @@ -41,7 +45,7 @@ def _fit(image: np.ndarray, tw: int, th: int) -> np.ndarray: return RemoteSession._resize_to(image, int(w * scale), int(h * scale)) def _prepare_obs(self, obs: dict[str, Any]) -> dict[str, Any]: - return {key: self._prepare_value(key, value) for key, value in obs.items()} + return {key: self._prepare_value(key, value) for key, value in obs.items() if key not in _LOCAL_ONLY_KEYS} def _prepare_value(self, key: str, value: Any) -> Any: # Client-side codecs (e.g. GR00T) nest images inside dicts/lists, so recurse to reach every diff --git a/positronic/policy/tests/test_change_ee_frame.py b/positronic/policy/tests/test_change_ee_frame.py new file mode 100644 index 000000000..3c76dc085 --- /dev/null +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -0,0 +1,94 @@ +import numpy as np + +import positronic.drivers.roboarm.command as cmd_module +from positronic.dataset.episode import EpisodeContainer +from positronic.dataset.tests.utils import DummySignal +from positronic.drivers.roboarm.ik import frame_transform +from positronic.drivers.roboarm.models import bundled_franka_model +from positronic.geom import Rotation, Transform3D +from positronic.policy.codec import ChangeEEFrame + +QUAT = Rotation.Representation.QUAT +# RoboLab's DROID end-effector control frame: eef_frame = Robotiq_2F_85/base_link composed with a pure rotation +# EEF_OFFSET_ROT (wxyz) and zero translation (robolab/robots/droid.py). ``droid_eef`` reproduces it. +EEF_OFFSET_ROT = (0.5, -0.5, 0.5, -0.5) +URDF = bundled_franka_model()['urdf'] +CONTROL_FRAME = 'end_effector' + + +def _pose(t, euler): + return Transform3D(np.asarray(t, dtype=np.float64), Rotation.from_euler(euler)) + + +def _quat_close(a, b, atol=1e-9): + # Quaternion double cover: q and -q are the same rotation. + return min(np.linalg.norm(a - b), np.linalg.norm(a + b)) < atol + + +def test_droid_eef_realizes_the_eef_offset_rotation(): + """The canonical->droid transform is the DROID ``EEF_OFFSET_ROT`` rotation plus the base offset, so our + ``droid_eef`` site reproduces RoboLab's ``eef_frame = base_link ∘ EEF_OFFSET_ROT``.""" + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + assert _quat_close(transform.rotation.as_quat, np.array(EEF_OFFSET_ROT)) + + +def test_encode_maps_obs_to_policy_frame(): + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + pose_c = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + obs = {'robot_state.ee_pose': pose_c.as_vector(QUAT), 'urdf': URDF, 'control_frame': CONTROL_FRAME, 'grip': 0.5} + + encoded = ChangeEEFrame(to='droid_eef').encode(obs) + + np.testing.assert_allclose(encoded['robot_state.ee_pose'], (pose_c * transform).as_vector(QUAT), atol=1e-9) + assert encoded['grip'] == 0.5, 'unrelated obs keys pass through' + + +def test_decode_maps_action_back_to_canonical(): + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + pose_c = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + obs = {'robot_state.ee_pose': pose_c.as_vector(QUAT), 'urdf': URDF, 'control_frame': CONTROL_FRAME} + # The policy emits its command in the droid frame (canonical composed with the transform); decode must invert it. + action = {'robot_command': cmd_module.CartesianPosition(pose=pose_c * transform), 'target_grip': 1.0} + + decoded = ChangeEEFrame(to='droid_eef')._decode_single(dict(action), context=obs) + + np.testing.assert_allclose(decoded['robot_command'].pose.as_vector(QUAT), pose_c.as_vector(QUAT), atol=1e-9) + assert decoded['target_grip'] == 1.0 + + +def test_decode_passes_non_cartesian_commands_through(): + obs = {'urdf': URDF, 'control_frame': CONTROL_FRAME} + action = {'robot_command': cmd_module.JointPosition(positions=np.zeros(7)), 'target_grip': 0.0} + decoded = ChangeEEFrame(to='droid_eef')._decode_single(dict(action), context=obs) + assert isinstance(decoded['robot_command'], cmd_module.JointPosition) + + +def test_identity_when_target_equals_control_frame(): + pose_c = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + obs = {'robot_state.ee_pose': pose_c.as_vector(QUAT), 'urdf': URDF, 'control_frame': CONTROL_FRAME} + encoded = ChangeEEFrame(to=CONTROL_FRAME).encode(obs) + np.testing.assert_allclose(encoded['robot_state.ee_pose'], pose_c.as_vector(QUAT), atol=1e-9) + + +def test_training_encoder_maps_both_poses_forward(): + """At training both the observed and the commanded pose map forward ``* T`` (both are canonical->policy), + the deliberate dual of the inference asymmetry (obs ``* T``, action ``* T``-inverse).""" + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + obs_pose = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + cmd_pose = _pose([0.2, 0.0, 0.5], [0.0, 0.1, -0.2]) + ts = [1000, 2000] + episode = EpisodeContainer( + data={ + 'urdf': URDF, + 'control_frame': CONTROL_FRAME, + 'robot_state.ee_pose': DummySignal(ts, np.stack([obs_pose.as_vector(QUAT)] * 2)), + 'robot_command.pose': DummySignal(ts, np.stack([cmd_pose.as_vector(QUAT)] * 2)), + 'grip': DummySignal(ts, np.array([0.0, 1.0])), + } + ) + + out = ChangeEEFrame(to='droid_eef').training_encoder(episode) + + np.testing.assert_allclose(out['robot_state.ee_pose'][0][0], (obs_pose * transform).as_vector(QUAT), atol=1e-9) + np.testing.assert_allclose(out['robot_command.pose'][0][0], (cmd_pose * transform).as_vector(QUAT), atol=1e-9) + assert out['control_frame'] == CONTROL_FRAME and 'grip' in out, 'statics and unrelated signals pass through' diff --git a/positronic/simulator/robolab/launcher.py b/positronic/simulator/robolab/launcher.py index a887d31ef..79473c241 100644 --- a/positronic/simulator/robolab/launcher.py +++ b/positronic/simulator/robolab/launcher.py @@ -93,13 +93,13 @@ def _spawn(host: str, port: int) -> subprocess.Popen: # interpreter and cannot build it, so it is serialized here (wire codec) and env.py emits it as # ``robot_meta``. A fresh temp file per spawn, in the container-local tmpdir — never the shared cache # filesystem, where a concurrent eval's rewrite could tear a reader mid-decode. - # HACK: the model declares ``control_frame='end_effector'`` (the flange), but the RoboLab env reports and - # accepts poses in DROID's eef frame (gripper base ∘ EEF_OFFSET_ROT), so episode statics mislabel the frame - # and offline IK over RoboLab episodes would solve the wrong target. Fixed by ``ChangeEEFrame`` + a - # ``droid_eef`` site: https://github.com/Positronic-Robotics/positronic/issues/483 + # RoboLab reports and accepts poses in DROID's eef frame (``droid_eef`` = gripper base ∘ EEF_OFFSET_ROT), so + # the model's canonical ``control_frame`` is declared as that frame — its poses and offline IK over RoboLab + # episodes then live in the frame the env actually uses. + robot_meta = {**bundled_franka_model(), 'control_frame': 'droid_eef'} fd, meta_path = tempfile.mkstemp(prefix='robolab_robot_meta_', suffix='.bin') with os.fdopen(fd, 'wb') as meta_file: - meta_file.write(encode(bundled_franka_model())) + meta_file.write(encode(robot_meta)) # Isaac Sim prompts for its EULA on stdin at first launch; the server is headless, so accept it here. env = { **os.environ, From b3b93d3c4ec94be1d1854d1fa752905ce1bcb8e8 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 19 Jul 2026 20:54:54 +0300 Subject: [PATCH 02/23] Fix `droid_eef` geometry to match RoboLab's `eef_frame` Our `gripper_base` (MuJoCo Menagerie) does not coincide with RoboLab's `Robotiq_2F_85/base_link`, so routing `droid_eef` through it was wrong. Define it directly on the flange from RoboLab's DROID USD: 18.17mm along Z, +90 deg Z. Verified against the USD at 0.000mm / 0.00000 quaternion difference. --- positronic/drivers/roboarm/models.py | 17 +++++++---------- positronic/policy/tests/test_change_ee_frame.py | 17 +++++++++-------- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/positronic/drivers/roboarm/models.py b/positronic/drivers/roboarm/models.py index e235eb5bd..9e6e6cf12 100644 --- a/positronic/drivers/roboarm/models.py +++ b/positronic/drivers/roboarm/models.py @@ -6,19 +6,16 @@ from functools import lru_cache from pathlib import Path -from positronic import geom - _FLANGE_LINK = 'link8' # Seat the gripper on the flange, rotated about the approach axis to match the real 2F-85 coupler # (a +45deg Z, i.e. 90deg off the franka ``end_effector`` frame). _2F85_MOUNT_RPY = '0 0 0.7853981634' -# The DROID/RoboLab end-effector frame ``droid_eef``: the gripper base rotated by a fixed offset. RoboLab reports -# and accepts Cartesian poses in this frame (``eef_frame`` = its ``Robotiq_2F_85/base_link`` ∘ ``EEF_OFFSET_ROT``, -# a pure rotation). ``EEF_OFFSET_ROT`` lives in RoboLab's ``robolab/robots/droid.py`` as a wxyz quaternion; its -# euler form is the URDF ``rpy`` of the ``droid_eef`` frame graft below. -_EEF_OFFSET_ROT = (0.5, -0.5, 0.5, -0.5) -_DROID_EEF_RPY = ' '.join(f'{angle:.12g}' for angle in geom.Rotation.from_quat(_EEF_OFFSET_ROT).as_euler) +# The DROID/RoboLab end-effector control frame ``droid_eef``, relative to the franka flange (``link8``). RoboLab +# reports and accepts Cartesian poses in it (``eef_frame`` = its ``Robotiq_2F_85/base_link`` ∘ ``EEF_OFFSET_ROT``). +# Measured off RoboLab's DROID USD (``franka_robotiq_2f_85_flattened.usd``): 18.17mm along the flange Z, +90deg Z. +_DROID_EEF_XYZ = '0 0 0.01817402261' +_DROID_EEF_RPY = '0 0 1.5707963268' def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: @@ -75,13 +72,13 @@ def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: # A row with an axis is a revolute joint whose axis sign sets its closing direction, so one positive # ``grip`` drives the whole 4-bar: driver/spring_link swing the finger in (+X), coupler/follower # counter-rotate (-X) to keep the pad parallel. Rows without an axis are fixed. A row with ``mesh`` None -# is a pure frame (no visual) — ``droid_eef`` is such a frame, the DROID control frame offset from the base. +# is a pure frame (no visual) — ``droid_eef`` is such a frame, the DROID control frame on the flange. _ROBOTIQ_2F85 = [ ('gripper_base_mount', _FLANGE_LINK, None, '0 0 0.007', _2F85_MOUNT_RPY, None, 'base_mount.stl', None), ('gripper_base', 'gripper_base_mount', None, '0 0 0.0038', '0 0 -1.5707963268', None, 'base.stl', None), *_2f85_finger('right', 1, '0 0 0'), *_2f85_finger('left', -1, '0 0 3.1415926536'), - ('droid_eef', 'gripper_base', None, '0 0 0', _DROID_EEF_RPY, None, None, None), + ('droid_eef', _FLANGE_LINK, None, _DROID_EEF_XYZ, _DROID_EEF_RPY, None, None, None), ] _ROBOTIQ_2F85_JOINTS = [row[2] for row in _ROBOTIQ_2F85 if row[2]] diff --git a/positronic/policy/tests/test_change_ee_frame.py b/positronic/policy/tests/test_change_ee_frame.py index 3c76dc085..00d52af0c 100644 --- a/positronic/policy/tests/test_change_ee_frame.py +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -9,9 +9,9 @@ from positronic.policy.codec import ChangeEEFrame QUAT = Rotation.Representation.QUAT -# RoboLab's DROID end-effector control frame: eef_frame = Robotiq_2F_85/base_link composed with a pure rotation -# EEF_OFFSET_ROT (wxyz) and zero translation (robolab/robots/droid.py). ``droid_eef`` reproduces it. -EEF_OFFSET_ROT = (0.5, -0.5, 0.5, -0.5) +# RoboLab's DROID end-effector control frame ``eef_frame`` = Robotiq_2F_85/base_link ∘ EEF_OFFSET_ROT with zero +# position (robolab/robots/droid.py). Measured off RoboLab's DROID USD, relative to the flange it is 18.17mm along +# Z and a +90deg Z rotation; ``droid_eef`` reproduces it. URDF = bundled_franka_model()['urdf'] CONTROL_FRAME = 'end_effector' @@ -25,11 +25,12 @@ def _quat_close(a, b, atol=1e-9): return min(np.linalg.norm(a - b), np.linalg.norm(a + b)) < atol -def test_droid_eef_realizes_the_eef_offset_rotation(): - """The canonical->droid transform is the DROID ``EEF_OFFSET_ROT`` rotation plus the base offset, so our - ``droid_eef`` site reproduces RoboLab's ``eef_frame = base_link ∘ EEF_OFFSET_ROT``.""" - transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') - assert _quat_close(transform.rotation.as_quat, np.array(EEF_OFFSET_ROT)) +def test_droid_eef_matches_robolab_eef_frame(): + """``droid_eef`` relative to the flange is RoboLab's ``eef_frame``, measured from its DROID USD: 18.17mm along + the flange Z and a +90deg Z rotation (``link8`` -> ``Robotiq_2F_85/base_link`` ∘ EEF_OFFSET_ROT).""" + transform = frame_transform(URDF, 'link8', 'droid_eef') + np.testing.assert_allclose(transform.translation, [0.0, 0.0, 0.01817402261], atol=1e-9) + assert _quat_close(transform.rotation.as_quat, Rotation.from_euler([0.0, 0.0, np.pi / 2]).as_quat) def test_encode_maps_obs_to_policy_frame(): From 89601db75c935bfc46d8f9e77988bd1fdb030c68 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 19 Jul 2026 21:07:37 +0300 Subject: [PATCH 03/23] Relabel `control_frame` to the policy frame in `ChangeEEFrame` training The training dual moves both poses into `to`; relabel `control_frame` too so a downstream `IKJointsAction` solves the command pose against the matching site instead of the canonical one. Document that context-reading action decoders (`RelativePositionAction`) are not yet frame-aware. --- positronic/policy/codec.py | 13 ++++++++++--- positronic/policy/tests/test_change_ee_frame.py | 3 ++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 083171536..332a4b1a2 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -16,7 +16,7 @@ from positronic import geom from positronic.dataset.transforms import Elementwise, lazy_sequence -from positronic.dataset.transforms.episode import Derive, EpisodeTransform, Group, Identity +from positronic.dataset.transforms.episode import Derive, EpisodeTransform, FromValue, Group, Identity from positronic.drivers.roboarm import command from positronic.drivers.roboarm.ik import frame_transform from positronic.policy.base import DelegatingSession, PolicyWrapper, Session, _Pipeline @@ -371,6 +371,10 @@ class ChangeEEFrame(Codec): Runs client-side, reading the robot model the harness injects into the local obs. That model is client-only and never reaches a frame-agnostic server: ``RemoteSession`` drops it at the wire boundary. + Composes with absolute-pose action codecs (``AbsolutePositionAction``). A decoder that reconstructs its command + from the observation pose in the decode context (``RelativePositionAction``) would read the canonical pose, not + the policy-frame one — such context-reading decoders need frame handling that is not yet wired here. + Compose to the left of the observation/action codecs:: ChangeEEFrame(to='droid_eef') | ee @@ -403,11 +407,14 @@ def derive(episode): @property def training_encoder(self) -> EpisodeTransform: - poses = { + # Relabel ``control_frame`` alongside the poses: once both are in ``to``, a later codec that reads the + # frame from statics (``IKJointsAction`` solving the command pose) resolves it against the right site. + derived = { + 'control_frame': FromValue(self._to), self._ee_pose_key: self._derive_pose(self._ee_pose_key), self._command_pose_key: self._derive_pose(self._command_pose_key), } - return Group(Derive(**poses), Identity()) + return Group(Derive(**derived), Identity()) @property def meta(self): diff --git a/positronic/policy/tests/test_change_ee_frame.py b/positronic/policy/tests/test_change_ee_frame.py index 00d52af0c..671b129b5 100644 --- a/positronic/policy/tests/test_change_ee_frame.py +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -92,4 +92,5 @@ def test_training_encoder_maps_both_poses_forward(): np.testing.assert_allclose(out['robot_state.ee_pose'][0][0], (obs_pose * transform).as_vector(QUAT), atol=1e-9) np.testing.assert_allclose(out['robot_command.pose'][0][0], (cmd_pose * transform).as_vector(QUAT), atol=1e-9) - assert out['control_frame'] == CONTROL_FRAME and 'grip' in out, 'statics and unrelated signals pass through' + # ``control_frame`` is relabeled to the policy frame so downstream IK reads the transformed poses correctly. + assert out['control_frame'] == 'droid_eef' and 'grip' in out, 'frame relabeled; unrelated signals pass through' From fbf7bc8267a2e3df034880047553129e483c3e98 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 19 Jul 2026 21:20:21 +0300 Subject: [PATCH 04/23] Convert only the EE pose signals an episode has in `ChangeEEFrame` Joint-only datasets carry no `robot_command.pose`; registering a derived signal for it made the LeRobot conversion dereference a missing key. Derive a pose transform only when the episode has that signal, so `ee_frame` composed with a joint action still converts the observation pose. --- positronic/policy/codec.py | 26 +++++++++++++------ .../policy/tests/test_change_ee_frame.py | 22 ++++++++++++++++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 332a4b1a2..acd657c2e 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -407,14 +407,24 @@ def derive(episode): @property def training_encoder(self) -> EpisodeTransform: - # Relabel ``control_frame`` alongside the poses: once both are in ``to``, a later codec that reads the - # frame from statics (``IKJointsAction`` solving the command pose) resolves it against the right site. - derived = { - 'control_frame': FromValue(self._to), - self._ee_pose_key: self._derive_pose(self._ee_pose_key), - self._command_pose_key: self._derive_pose(self._command_pose_key), - } - return Group(Derive(**derived), Identity()) + return _ChangeEEFrameTraining(self._to, (self._ee_pose_key, self._command_pose_key), self._derive_pose) + + +class _ChangeEEFrameTraining(EpisodeTransform): + """Move the EE pose signals an episode has into ``to`` and relabel ``control_frame`` to match, so a later codec + that reads the frame from statics (``IKJointsAction`` solving the command pose) resolves it against the right + site. A pose key the episode lacks — ``robot_command.pose`` under a joint-only action — is skipped rather than + dereferenced, so joint-only training still converts its observation pose.""" + + def __init__(self, to: str, pose_keys: tuple[str, ...], derive_pose): + self._to = to + self._pose_keys = pose_keys + self._derive_pose = derive_pose + + def __call__(self, episode): + derived = {'control_frame': FromValue(self._to)} + derived.update({key: self._derive_pose(key) for key in self._pose_keys if key in episode}) + return Group(Derive(**derived), Identity())(episode) @property def meta(self): diff --git a/positronic/policy/tests/test_change_ee_frame.py b/positronic/policy/tests/test_change_ee_frame.py index 671b129b5..d1f513bf4 100644 --- a/positronic/policy/tests/test_change_ee_frame.py +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -94,3 +94,25 @@ def test_training_encoder_maps_both_poses_forward(): np.testing.assert_allclose(out['robot_command.pose'][0][0], (cmd_pose * transform).as_vector(QUAT), atol=1e-9) # ``control_frame`` is relabeled to the policy frame so downstream IK reads the transformed poses correctly. assert out['control_frame'] == 'droid_eef' and 'grip' in out, 'frame relabeled; unrelated signals pass through' + + +def test_training_encoder_skips_absent_command_pose(): + """A joint-only dataset has no ``robot_command.pose``; the training dual converts the observation pose and + relabels the frame without registering (and dereferencing) the missing command pose.""" + obs_pose = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + ts = [1000, 2000] + episode = EpisodeContainer( + data={ + 'urdf': URDF, + 'control_frame': CONTROL_FRAME, + 'robot_state.ee_pose': DummySignal(ts, np.stack([obs_pose.as_vector(QUAT)] * 2)), + 'robot_command.joints': DummySignal(ts, np.zeros((2, 7), dtype=np.float32)), + } + ) + + out = ChangeEEFrame(to='droid_eef').training_encoder(episode) + + assert 'robot_command.pose' not in list(out), 'absent command pose must not be materialized' + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + np.testing.assert_allclose(out['robot_state.ee_pose'][0][0], (obs_pose * transform).as_vector(QUAT), atol=1e-9) + assert out['control_frame'] == 'droid_eef' and 'robot_command.joints' in out From eb89f73cdb991cdfa7a078faa5c655212ecb7506 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 21:24:12 +0000 Subject: [PATCH 05/23] Add the MolmoSpaces env-server integration Serve AllenAI's MolmoSpaces MuJoCo benchmark behind the env-server wire so positronic owns the control loop: `env.py` drives a single `BaseMujocoTask` per episode in MolmoSpaces' own venv (pinned checkout, `uv run --project`), the client-side `MolmoAdapter` maps the raw payload into the canonical contract, and `cfg/eval/sim/molmo.py` sweeps a benchmark's episodes into one CLI. Pure wire mappings live in `mapping.py` (framework-free, fixture-tested). Also lowercase the DROID prompt in the observation codec (the checkpoints train on lowercased language) and document the adapter/codec/wire-client separation of responsibilities. Ticket: Positronic-Robotics/internal#91 #refs --- .github/workflows/unit-test.yaml | 2 +- CLAUDE.md | 2 + positronic/cfg/eval/sim/molmo.py | 63 +++++++ positronic/offboard/README.md | 28 +++ positronic/policy/observation.py | 12 +- positronic/simulator/molmo_spaces/__init__.py | 0 positronic/simulator/molmo_spaces/adapter.py | 56 ++++++ positronic/simulator/molmo_spaces/env.py | 162 ++++++++++++++++++ positronic/simulator/molmo_spaces/launcher.py | 92 ++++++++++ positronic/simulator/molmo_spaces/mapping.py | 88 ++++++++++ .../simulator/molmo_spaces/tests/__init__.py | 0 .../molmo_spaces/tests/droid_obs.npz | Bin 0 -> 1528 bytes .../molmo_spaces/tests/make_fixture.py | 53 ++++++ .../molmo_spaces/tests/test_adapter.py | 70 ++++++++ .../molmo_spaces/tests/test_mapping.py | 86 ++++++++++ positronic/vendors/openpi/codecs.py | 4 +- pyproject.toml | 1 + 17 files changed, 715 insertions(+), 4 deletions(-) create mode 100644 positronic/cfg/eval/sim/molmo.py create mode 100644 positronic/simulator/molmo_spaces/__init__.py create mode 100644 positronic/simulator/molmo_spaces/adapter.py create mode 100644 positronic/simulator/molmo_spaces/env.py create mode 100644 positronic/simulator/molmo_spaces/launcher.py create mode 100644 positronic/simulator/molmo_spaces/mapping.py create mode 100644 positronic/simulator/molmo_spaces/tests/__init__.py create mode 100644 positronic/simulator/molmo_spaces/tests/droid_obs.npz create mode 100644 positronic/simulator/molmo_spaces/tests/make_fixture.py create mode 100644 positronic/simulator/molmo_spaces/tests/test_adapter.py create mode 100644 positronic/simulator/molmo_spaces/tests/test_mapping.py diff --git a/.github/workflows/unit-test.yaml b/.github/workflows/unit-test.yaml index fa5a545c5..57a3b9040 100644 --- a/.github/workflows/unit-test.yaml +++ b/.github/workflows/unit-test.yaml @@ -36,7 +36,7 @@ jobs: PYTHONPATH: ${{ env.PYTHONPATH }}:$PWD run: >- uv run pytest --cov-report=html - --override-ini='testpaths=pimm/tests positronic/cfg/tests positronic/dataset/tests positronic/geom/tests positronic/offboard/tests positronic/policy/tests positronic/tests positronic/utils/tests' + --override-ini='testpaths=pimm/tests positronic/cfg/tests positronic/dataset/tests positronic/geom/tests positronic/offboard/tests positronic/policy/tests positronic/simulator/molmo_spaces/tests positronic/tests positronic/utils/tests' - name: Coverage summary (job summary) if: always() diff --git a/CLAUDE.md b/CLAUDE.md index 0e68f1db8..f2cea7ab9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -73,4 +73,6 @@ # Infrastructure - Machines, Docker contexts and images: `docker/CONTEXTS.md` - Model-specific workflows: `positronic/vendors/{lerobot,gr00t,openpi}/README.md` +- Inference serving, and the adapter/codec/wire-client separation of responsibilities (read BEFORE + writing a sim/rig adapter): `positronic/offboard/README.md` - Reconstructing previous runs: read `run_metadata_*.yaml` and episode `static.json` from output directory diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py new file mode 100644 index 000000000..97cfbc844 --- /dev/null +++ b/positronic/cfg/eval/sim/molmo.py @@ -0,0 +1,63 @@ +import json +from pathlib import Path + +import configuronic as cfn + +from positronic.cfg.eval import build_trials +from positronic.drivers.roboarm.models import bundled_franka_model +from positronic.eval import Eval, Task +from positronic.simulator.env_server.proxy import RemoteEnvControlSystem, remote_franka_embodiment +from positronic.simulator.molmo_spaces.adapter import MolmoAdapter +from positronic.simulator.molmo_spaces.launcher import serve_molmo_spaces + + +def _episode_count(benchmark_dir: str) -> int: + """The number of episodes in a MolmoSpaces ``benchmark.json`` (a JSON list of episode specs).""" + return len(json.loads((Path(benchmark_dir) / 'benchmark.json').read_text())) + + +@cfn.config( + camera_dict={'image.wrist': 'wrist_camera', 'image.exterior': 'exo_camera_1'}, + benchmark_dir=None, + episodes=None, + trial_count=1, + timeout=60.0, + seed=None, +) +def _molmo_eval(benchmark_dir, episodes, trial_count, timeout, camera_dict, seed): + """A MolmoSpaces eval: the embodiment proxies a remote MolmoSpaces env, the task carries the scenario. + + MolmoSpaces (https://github.com/allenai/molmospaces) is AllenAI's MuJoCo manipulation benchmark on the DROID + rig (Franka arm + Robotiq 2F-85) across ProcTHOR scenes; a benchmark is a ``benchmark.json`` of episode specs + (house, task, exact object poses, cameras, language goal), so ``--eval.benchmark_dir`` names the benchmark to + run and ``--eval.episodes`` optionally pins a subset of episode indices (default: the whole benchmark). The + asset packs live under ``MLSPACES_ASSETS_DIR``. + + positronic launches a single task-agnostic env server in MolmoSpaces' own interpreter; the proxy drives it + over the socket and the episode index rides each trial's reset token, so one embodiment serves every episode. + The instruction is never pinned: the task reads its language live from the env, which reports the episode's + resolved goal in every reset's meta. Episodes are exact-pose deterministic, so ``trial_count`` defaults to 1. + """ + if benchmark_dir is None: + raise ValueError('MolmoSpaces eval needs --eval.benchmark_dir pointing at a dir with benchmark.json') + if episodes is None: + indices = list(range(_episode_count(benchmark_dir))) + else: + indices = [episodes] if isinstance(episodes, int) else list(episodes) + proxy = RemoteEnvControlSystem(MolmoAdapter(camera_dict), serve_molmo_spaces(benchmark_dir)) + # MolmoSpaces drives a Franka DROID rig; recordings carry the same model (URDF + meshes + joint names + + # control frame) for the 3D viewer and offline IK, supplied here since the molmo server can't import + # positronic to emit it via ``robot_meta``. + embodiment = remote_franka_embodiment( + proxy, camera_dict, descriptor='remote.molmo_spaces.droid', static_meta=bundled_franka_model() + ) + task = Task(instruction=lambda: proxy.meta['task'], timeout=timeout, reset=proxy.reset, done=proxy.done) + scenes = [{'eval.episode_index': i} for i in indices] + return Eval(embodiment, task, build_trials(seed, trial_count, scenes)) + + +# The whole benchmark in one run (every episode in ``--eval.benchmark_dir``'s benchmark.json). +benchmark = _molmo_eval + +# A single-episode smoke target: the first episode of the benchmark. +first_episode = _molmo_eval.override(episodes=0) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ce1cbbedc..074ab813a 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -2,6 +2,34 @@ This package implements the protocol and utilities for offboard policy inference, allowing robots or simulators to stream observations to a remote server and receive actions. +## Separation of responsibilities: adapter vs codec vs wire client + +Three layers touch an observation on its way to a model, and each owns exactly one concern. +When writing a new sim/rig adapter, check this table before adding any transform to it: + +| Layer | Owns | Examples | +|---|---|---| +| **Adapter** (per sim/rig, e.g. `simulator/molmo_spaces/adapter.py`) | Rig semantics ONLY: mapping the rig's observation/action vocabulary onto positronic's raw keys | Camera-key mapping, gripper qpos → `[0, 1]` closure, decoded commands → the rig's action format | +| **Codec** (per model family, `policy/codec.py` subclasses) | Model preprocessing: everything the checkpoint's input distribution requires | Resize-with-pad to model resolution, prompt normalization (e.g. DROID lowercasing), state assembly | +| **Wire client** (`InferenceClient` / `RemotePolicy`) | Transport optimization, negotiated — never semantics | Downscaling frames to the server-advertised `image_sizes` (aspect-preserving, never upscaling), optional JPEG compression | + +Consequences: + +- **An adapter never resizes, pads, normalizes prompts, or otherwise preprocesses for the model.** + It passes frames and text through at native fidelity. If the same transform appears in an adapter + and a codec, the adapter's copy is the bug: a drifted duplicate silently changes eval inputs. +- **Bandwidth is not the adapter's problem.** The client already downsizes to what the server says + it needs: every `Codec` advertises its expected input sizes via the reserved `image_sizes` meta + key (see `Codec.meta`), the server returns it in the session handshake, and the client fits + frames to it before sending. This is default-on — an adapter that resizes "to keep the wire + payload small" is duplicating it. +- **Codecs run on either side of the wire.** positronic-native evals compose the codec around + `RemotePolicy` on the client (`cfg/policy.py` — the wire then carries model-sized encoded inputs, + and the client-side resize is disabled since `codec.meta` already reports `image_sizes`). + Thin-client deployments (a sim adapter in a foreign venv talking to a serverless endpoint) host + the codec on the server — the wire carries raw positronic keys, downsized by the negotiation + above. Both placements are supported; pick by where the dependencies can live. + ## Protocol v1 The unified WebSocket protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index 9d424b7a0..e7709fd27 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -18,14 +18,21 @@ class ObservationCodec(Codec): state: mapping from output state key to an ordered dict of {episode_key: dim} to concatenate. images: mapping from output image name to tuple (input_key, (width, height)). task_field: output key carrying the language prompt at inference; LeRobot training always uses ``task``. + lowercase_task: lowercase the task text at inference, for checkpoints trained on lowercased language + (the pretrained DROID models; MolmoSpaces' Pi baseline applies the same normalization). """ def __init__( - self, state: dict[str, dict[str, int]], images: dict[str, tuple[str, tuple[int, int]]], task_field: str = 'task' + self, + state: dict[str, dict[str, int]], + images: dict[str, tuple[str, tuple[int, int]]], + task_field: str = 'task', + lowercase_task: bool = False, ): self._state = state self._image_configs = images self._task_field = task_field + self._lowercase_task = lowercase_task self._derive_transforms = {k: partial(self._derive_state, k) for k in state.keys()} self._derive_transforms.update({k: partial(self._derive_image, k) for k in images.keys()}) @@ -54,7 +61,8 @@ def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: obs: dict[str, Any] = {} if 'task' in inputs: - obs[self._task_field] = inputs['task'] + task = inputs['task'] + obs[self._task_field] = task.lower() if self._lowercase_task else task for out_name, (input_key, (width, height)) in self._image_configs.items(): if input_key not in inputs: diff --git a/positronic/simulator/molmo_spaces/__init__.py b/positronic/simulator/molmo_spaces/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/positronic/simulator/molmo_spaces/adapter.py b/positronic/simulator/molmo_spaces/adapter.py new file mode 100644 index 000000000..1d084e8ae --- /dev/null +++ b/positronic/simulator/molmo_spaces/adapter.py @@ -0,0 +1,56 @@ +"""``MolmoAdapter``: the canonical embodiment contract <-> MolmoSpaces' raw obs/command payloads, client-side. + +Runs in positronic's interpreter (the ``MolmoSpacesEnv`` server runs in MolmoSpaces' own). Mirrors the LIBERO +adapter on the observation side; the command side is ``WireCommandAdapter``'s forwarding. All action encoding — +the wire command into MolmoSpaces' per-move-group joint targets — lives server-side in ``env.py`` where the +MuJoCo model is, so the adapter holds no model and stays geometry-only. +""" + +from typing import Any + +import pimm +from positronic import geom +from positronic.simulator.env_server.adapter import WireCommandAdapter +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.mujoco.sim import MujocoFrankaState + +# Per default MolmoSpaces DROID camera, the benchmark-variant keys the upstream Pi policy falls back to; an +# explicitly configured non-default camera key is read as-is (no variants). +_CAMERA_VARIANTS = { + mapping.MOLMO_WRIST_CAMERA: mapping.MOLMO_WRIST_CAMERA_VARIANTS, + mapping.MOLMO_EXTERIOR_CAMERA: mapping.MOLMO_EXTERIOR_CAMERA_VARIANTS, +} + + +class MolmoAdapter(WireCommandAdapter): + def __init__(self, camera_dict: dict[str, str]): + super().__init__() + self._camera_dict = camera_dict # logical observation name -> the MolmoSpaces obs camera key + + def _reset_token(self, context: dict[str, Any]) -> Any: + # The benchmark episode selector rides the token: env.py loads the benchmark once and builds the task + # for this episode index, seeding from the spec (``eval.seed`` overrides the spec's own seed when set). + return {'episode_index': context['eval.episode_index'], 'seed': context.get('eval.seed')} + + def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: + # env.py reports the eef pose in the grasp-site world frame; ``eef_quat`` is scalar-first (wxyz, from + # ``mju_mat2Quat``), so ``from_quat`` is the matching decode. + ee_pose = geom.Transform3D(raw_obs['eef_pos'], geom.Rotation.from_quat(raw_obs['eef_quat'])) + state = MujocoFrankaState() + state.encode(raw_obs['joint_pos'], raw_obs['joint_vel'], ee_pose) + obs: dict[str, Any] = {'robot_state': state, 'grip': float(raw_obs['grip'])} + for logical, molmo_key in self._camera_dict.items(): + env_key = mapping.resolve_camera_key(raw_obs, molmo_key, molmo_key, _CAMERA_VARIANTS.get(molmo_key, ())) + frame = raw_obs[env_key] # MolmoSpaces renders top-down already — no flip + adapter = pimm.shared_memory.NumpySMAdapter(shape=frame.shape, dtype=frame.dtype) + adapter.array[:] = frame + obs[logical] = adapter + return obs + + def privileged(self, raw_obs: dict[str, Any]) -> dict[str, Any]: + return {} + + def terminal(self, result: dict[str, Any]) -> dict[str, Any] | None: + # ``done`` covers termination and timeout; ``success`` is the task's judged success, so a timeout stays + # honest. + return {'eval.success': bool(result['success'])} if result['done'] else None diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py new file mode 100644 index 000000000..9c73c0c59 --- /dev/null +++ b/positronic/simulator/molmo_spaces/env.py @@ -0,0 +1,162 @@ +"""MolmoSpaces — AllenAI's MuJoCo manipulation benchmark — behind the env-server protocol. + +MolmoSpaces pins ``mujoco ~=3.5`` + its asset stack into its own uv project on Python 3.11, so this never shares +positronic's venv: the launcher runs ``uv run --project env.py --host ... --port ...`` with the +positronic-free ``server``/``protocol`` and this package's ``mapping`` module on ``PYTHONPATH``. It imports only +``molmo_spaces`` (+ mujoco/numpy) and those, never ``positronic``. + +positronic owns the control loop: this server drives a single MolmoSpaces ``BaseMujocoTask`` per episode directly +(``JsonEvalTaskSampler.sample_task`` builds the full sim/scene/renderer; ``reset``/``step``/``is_done``/ +``judge_success`` drive it), replacing MolmoSpaces' own ``JsonEvalRunner`` loop. The reset token selects the +benchmark episode (index into ``benchmark.json``) and an optional seed; the client-side ``MolmoAdapter`` maps the +raw payload this server reports into the canonical embodiment contract. + +Command side: the ``MolmoAdapter`` forwards a joint command (the DROID rig runs the joint-position controller); +this server integrates it onto the measured joints and steps the per-move-group ``{arm, gripper}`` action. +Observation side: MolmoSpaces' obs carries the joint positions/velocities and camera frames, but the +end-effector *world* pose is read from the robot view's grasp-site frame here, alongside the gripper closure, into +the raw payload the adapter assembles into a ``MujocoFrankaState``. +""" + +import argparse +import os + +# MolmoSpaces renders MuJoCo scenes, so the GL backend must be selected before any mujoco/molmo_spaces import. +# The launcher sets MUJOCO_GL in the subprocess env (egl by default); default it here too so a direct invocation +# (e.g. a validate/e2e run) still boots. Set before the imports below. +os.environ.setdefault('MUJOCO_GL', 'egl') + +from pathlib import Path # noqa: E402 +from typing import Any # noqa: E402 + +import mapping # noqa: E402 -- positronic-free wire mappings, on PYTHONPATH +import mujoco # noqa: E402 +import numpy as np # noqa: E402 +from server import EnvProtocol, EnvServer # noqa: E402 + +from molmo_spaces.configs.policy_configs import DummyPolicyConfig # noqa: E402 +from molmo_spaces.configs.robot_configs import ActionNoiseConfig, FrankaRobotConfig # noqa: E402 +from molmo_spaces.evaluation.benchmark_schema import load_all_episodes # noqa: E402 +from molmo_spaces.evaluation.configs.evaluation_configs import JsonBenchmarkEvalConfig # noqa: E402 +from molmo_spaces.tasks.json_eval_task_sampler import JsonEvalTaskSampler # noqa: E402 + + +class _DroidPickEvalConfig(JsonBenchmarkEvalConfig): + """The minimal eval config to build a Franka DROID pick task standalone. + + ``JsonBenchmarkEvalConfig`` defaults every ``MlSpacesExpConfig`` field except the robot and policy configs; + the sampler overrides ``task_type``/``scene_dataset``/``data_split``/``camera_config``/``house_inds`` from the + episode spec, so only these two are supplied. The policy config is a ``DummyPolicyConfig`` — positronic owns + the policy, and ``sample_task`` never calls the framework's ``policy_factory`` (only reads + ``force_enable_depth``). + """ + + robot_config: FrankaRobotConfig = FrankaRobotConfig() + policy_config: DummyPolicyConfig = DummyPolicyConfig() + + def model_post_init(self, __context: Any) -> None: + super().model_post_init(__context) + self.robot_config.action_noise_config = ActionNoiseConfig(enabled=False) + + +class MolmoSpacesEnv(EnvProtocol): + """A MolmoSpaces benchmark episode behind the gym-style ``reset``/``step``/``close`` the env server serves. + + Built per reset from the token's episode index (into the loaded ``benchmark.json``) and its seed. MolmoSpaces' + ``task.reset()`` does not re-place the scene — ``sample_task`` does — so each reset rebuilds the task for a + clean, deterministic scene (benchmark episodes are exact-pose deterministic, so a rebuild reproduces them). + ``step`` integrates the forwarded joint command onto the measured joints, drives the per-move-group action, + and reports MolmoSpaces' ``is_done``/``judge_success``. + """ + + def __init__(self, benchmark_dir: str): + self._episodes = load_all_episodes(Path(benchmark_dir)) + self._sampler = None + self._task = None + self._robot_view = None + self._control_dt = None + self._meta = None + # The RGB camera keys the current episode renders — emitted every frame; the client's ``camera_dict`` + # selects which the policy sees. + self._camera_names: list[str] = [] + + def _build(self, episode_index: int, seed: int | None) -> None: + if self._sampler is not None: + self._sampler.close() # release the prior episode's sim/renderer before building the next + episode = self._episodes[episode_index] + cfg = _DroidPickEvalConfig() + # Determinism enters at sampler construction (seed_task_sampling); the token's seed overrides the spec's. + cfg.seed = int(seed) if seed is not None else (episode.seed if episode.seed is not None else 42) + horizon_sec = episode.task.get('task_horizon_sec') + if horizon_sec is not None: + cfg.task_horizon = round(float(horizon_sec) * 1000.0 / cfg.policy_dt_ms) + self._sampler = JsonEvalTaskSampler(cfg, episode) + self._task = self._sampler.sample_task(house_index=episode.house_index) + self._robot_view = self._task.env.current_robot.robot_view + self._control_dt = cfg.policy_dt_ms / 1000.0 + self._meta = {'task': self._task.get_task_description(), 'house_index': episode.house_index} + + def reset(self, token: dict[str, Any]) -> dict[str, Any]: + self._build(token['episode_index'], token.get('seed')) + obs, _info = self._task.reset() # obs is a list, one dict per env; n_batch == 1 + env_obs = obs[0] + self._camera_names = [k for k, v in env_obs.items() if _is_rgb_frame(v)] + # robot_meta is empty: this venv cannot import positronic to emit the Franka model, so the eval supplies + # it via ``static_meta`` (``bundled_franka_model``). ``meta`` carries the scene/task identity. + return {'obs': self._observe(env_obs), 'meta': self._meta, 'robot_meta': {}, 'control_dt': self._control_dt} + + def step(self, action: dict[str, Any]) -> dict[str, Any]: + arm = mapping.wire_command_to_arm_action(action['command'], self._measured_arm_q()) + gripper = np.array([mapping.grip_command_to_actuator(action['grip'])], dtype=np.float32) + obs, _reward, _term, _trunc, _infos = self._task.step({'arm': arm, 'gripper': gripper}) + done = bool(self._task.is_done()) + # ``judge_success`` is the task's scored success; only meaningful once the episode is done, so a timeout + # stays honest (a running episode is never a success). + success = bool(self._task.judge_success()) if done else False + return {'obs': self._observe(obs[0]), 'done': done, 'success': success, 'control_dt': self._control_dt} + + def _measured_arm_q(self) -> np.ndarray: + return np.asarray(self._robot_view.get_move_group('arm').joint_pos, dtype=np.float32) + + def _observe(self, env_obs: dict[str, Any]) -> dict[str, Any]: + # MolmoSpaces' obs carries the joint positions/velocities and camera frames; the eef *world* pose is read + # from the arm move group's grasp-site frame (obs only exposes a robot-relative tcp pose). + arm = self._robot_view.get_move_group('arm') + eef_world = np.asarray(arm.leaf_frame_to_world, dtype=np.float64) # 4x4 grasp-site world transform + eef_quat = np.zeros(4) + mujoco.mju_mat2Quat(eef_quat, np.ascontiguousarray(eef_world[:3, :3].reshape(9))) # -> wxyz + payload = { + 'joint_pos': np.asarray(arm.joint_pos, dtype=np.float32), + 'joint_vel': np.asarray(arm.joint_vel, dtype=np.float32), + 'eef_pos': eef_world[:3, 3].astype(np.float32), + 'eef_quat': eef_quat.astype(np.float32), + 'grip': np.float32(mapping.normalize_grip_qpos(env_obs['qpos']['gripper'])), + } + for name in self._camera_names: + payload[name] = np.ascontiguousarray(env_obs[name]) + return payload + + def close(self) -> None: + if self._sampler is not None: + self._sampler.close() + self._sampler = None + self._task = None + + +def _is_rgb_frame(value: Any) -> bool: + return isinstance(value, np.ndarray) and value.ndim == 3 and value.shape[2] == 3 and value.dtype == np.uint8 + + +def main() -> None: + parser = argparse.ArgumentParser(description='Serve MolmoSpaces over the env-server protocol.') + parser.add_argument('--host', default='localhost') + parser.add_argument('--port', type=int, required=True) + parser.add_argument('--benchmark_dir', required=True, help='dir containing benchmark.json') + args = parser.parse_args() + if not os.environ.get('MLSPACES_ASSETS_DIR'): + parser.error('MLSPACES_ASSETS_DIR must point at the MolmoSpaces asset packs') + EnvServer(MolmoSpacesEnv(args.benchmark_dir), args.host, args.port).serve_forever() + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/launcher.py b/positronic/simulator/molmo_spaces/launcher.py new file mode 100644 index 000000000..f6d346b3b --- /dev/null +++ b/positronic/simulator/molmo_spaces/launcher.py @@ -0,0 +1,92 @@ +"""Launches the MolmoSpaces env server as a subprocess and owns its lifetime. + +positronic starts the server: the env runs in MolmoSpaces' own interpreter via ``uv run --project ``, +which resolves the ``molmospaces[mujoco]`` stack (mujoco ~=3.5, the resource-manager asset layer, torch) from the +pinned checkout — far too heavy, and Python-version-pinned (3.11), to share positronic's venv. The positronic-free +``env_server`` package and this package's ``mapping`` module ride ``PYTHONPATH`` so ``env.py`` imports the dumb +``server``/``protocol`` and the pure wire mappings without dragging in positronic; ``molmo_spaces`` itself resolves +from the uv project. + +MolmoSpaces renders MuJoCo scenes, so the server needs a GL backend (``MUJOCO_GL``) and its asset packs +(``MLSPACES_ASSETS_DIR``): ``MUJOCO_GL`` defaults to ``egl`` (GPU) here and both env vars pass through from the +caller, so a GPU-less box can override ``MUJOCO_GL=osmesa`` for CPU software rendering. +""" + +import fcntl +import os +import subprocess +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager +from pathlib import Path + +from positronic.simulator.env_server.launcher import ensure_pinned_checkout, serve_subprocess + +_ENV_SCRIPT = Path(__file__).parent / 'env.py' +_ENV_SERVER_DIR = Path(__file__).parents[1] / 'env_server' +_MAPPING_DIR = Path(__file__).parent # ``mapping.py`` — imported flat by env.py, positronic-free + +_MOLMO_REPO = 'https://github.com/allenai/molmospaces.git' +_MOLMO_COMMIT = 'c2f1b583f087e1d3994e1377574843b759d9d0f8' +_MOLMO_SRC = Path.home() / '.cache' / 'positronic' / 'molmospaces' / 'src' + +# MolmoSpaces pins Python 3.11 and installs its MuJoCo renderer stack via the ``mujoco`` extra (classic renderer, +# mujoco ~=3.5). ``mujoco-filament`` is the alternative for bench-v2 filament scenes; the classic renderer is the +# eval default. +_MOLMO_PYTHON = '3.11' +_MOLMO_EXTRA = 'mujoco' + + +@contextmanager +def _checkout_lock() -> Iterator[None]: + """Serialize checkout + ``uv sync`` across processes sharing the cache, so a warm-cache fan-out of eval jobs + mounting one ``~/.cache/positronic/molmospaces`` filesystem does not race a forced checkout against a sync.""" + _MOLMO_SRC.parent.mkdir(parents=True, exist_ok=True) + with open(_MOLMO_SRC.parent / 'setup.lock', 'w') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + yield + + +def _spawn(host: str, port: int, benchmark_dir: str) -> subprocess.Popen: + with _checkout_lock(): + src = ensure_pinned_checkout(_MOLMO_REPO, _MOLMO_COMMIT, _MOLMO_SRC) + # Install the stack before spawning: a cold first install far exceeds the client's connect deadline, + # which should only cover the sim's boot. Idempotent and fast when warm; the spawn passes ``--no-sync`` + # so no resolve or install ever runs outside this lock. MolmoSpaces ships no uv.lock, so this re-resolves + # on every fresh box. + subprocess.run( + ['uv', 'sync', '--project', str(src), '--python', _MOLMO_PYTHON, '--extra', _MOLMO_EXTRA], check=True + ) + command = [ + 'uv', + 'run', + '--no-sync', + '--project', + str(src), + '--python', + _MOLMO_PYTHON, + '--extra', + _MOLMO_EXTRA, + str(_ENV_SCRIPT), + '--host', + host, + '--port', + str(port), + '--benchmark_dir', + str(benchmark_dir), + ] + env = { + **os.environ, + 'PYTHONPATH': os.pathsep.join([str(_ENV_SERVER_DIR), str(_MAPPING_DIR)]), + # GPU OpenGL by default; a caller on a GPU-less box exports MUJOCO_GL=osmesa for CPU software rendering. + 'MUJOCO_GL': os.environ.get('MUJOCO_GL', 'egl'), + } + return subprocess.Popen(command, env=env) + + +def serve_molmo_spaces(benchmark_dir: str, host: str = 'localhost') -> AbstractContextManager[tuple[str, int]]: + """The MolmoSpaces env server as a ``serve`` context manager (the ``serve_subprocess`` contract). + + ``benchmark_dir`` (a dir holding ``benchmark.json``) is fixed for the run; the reset token selects the + episode within it, so one task-agnostic server serves every trial. + """ + return serve_subprocess(lambda host, port: _spawn(host, port, benchmark_dir), host) diff --git a/positronic/simulator/molmo_spaces/mapping.py b/positronic/simulator/molmo_spaces/mapping.py new file mode 100644 index 000000000..055e6619b --- /dev/null +++ b/positronic/simulator/molmo_spaces/mapping.py @@ -0,0 +1,88 @@ +"""Pure MolmoSpaces <-> positronic-wire mappings, free of both molmo_spaces and positronic. + +Imported from two interpreters: the client-side ``MolmoAdapter`` (positronic) resolves camera keys with +it, and the molmo-venv ``env.py`` builds its raw observation payload and decodes wire commands with it. It +imports only numpy, so it loads under a bare pytest and inside the molmo venv alike — the fixture tests +exercise it without either framework. The MuJoCo reads that need the live model (joint velocities, the +end-effector world pose) stay in ``env.py``; only the framework-independent arithmetic lives here. +""" + +from typing import Any + +import numpy as np + +# The DROID rig runs 7 Franka arm joints; the reset token's per-move-group action names them 'arm'/'gripper'. +NUM_ARM_JOINTS = 7 +MOLMO_ARM_GROUP = 'arm' +MOLMO_GRIPPER_GROUP = 'gripper' + +# MolmoSpaces DROID rig camera names (FrankaDroidCameraSystem); the Zed-wrist / light-randomized benchmark +# variants replace the defaults and MolmoSpaces' own Pi policy prefers them when present, so the adapter must too. +MOLMO_WRIST_CAMERA = 'wrist_camera' +MOLMO_EXTERIOR_CAMERA = 'exo_camera_1' +MOLMO_WRIST_CAMERA_VARIANTS = ('wrist_camera_zed_mini',) +MOLMO_EXTERIOR_CAMERA_VARIANTS = ('droid_shoulder_light_randomization',) + +# The Robotiq 2F-85 finger qpos saturates at this closure; the DROID observation's grip is normalized against +# it into the [0, 1] closure the policy was trained on (molmospaces pi_policy.py:126). +GRIPPER_QPOS_CLOSED = 0.824033 + +# The Robotiq gripper actuator is a single command, 0 fully open .. 255 fully closed (franka_droid_view.py:43). +ROBOTIQ_OPEN = 0.0 +ROBOTIQ_CLOSED = 255.0 + + +def normalize_grip_qpos(gripper_qpos: Any, gripper_qpos_closed: float = GRIPPER_QPOS_CLOSED) -> float: + """A Robotiq finger qpos -> the [0, 1] closure the observation reports (0 open, 1 closed).""" + value = float(np.asarray(gripper_qpos).reshape(-1)[0]) + return float(np.clip(value / gripper_qpos_closed, 0.0, 1.0)) + + +def grip_command_to_actuator(grip: float) -> float: + """A wire grip closure ([0, 1], 1 = closed) -> the Robotiq actuator command ([0, 255], 255 = closed). + + Continuous: the pi05 codec already binarizes the grip channel (``binarize_grip``), so the rig maps the + closure straight through rather than re-thresholding it here. + """ + return float(np.clip(grip, 0.0, 1.0)) * ROBOTIQ_CLOSED + + +def wire_command_to_arm_action(command: dict[str, Any], current_q: Any) -> np.ndarray: + """A tagged wire command + the live measured arm joints -> the 7 absolute joint targets molmo steps. + + MolmoSpaces' Franka runs the joint-position controller, so every command resolves to absolute joint + targets: ``joint_pos`` passes through, ``joint_vel`` integrates the per-step delta onto the measured + joints (positronic applies ``JointDelta`` as ``q + dq``), and ``hold`` re-commands the measured joints. + Cartesian commands would need IK against the live model, which this jointpos substrate does not run. + """ + current = np.asarray(current_q, dtype=np.float32).reshape(-1) + match command['type']: + case 'joint_pos': + target = np.asarray(command['q'], dtype=np.float32).reshape(-1) + case 'joint_vel': + dq = np.asarray(command['dq'], dtype=np.float32).reshape(-1) + if dq.shape[0] != current.shape[0]: + raise ValueError(f'joint delta {dq.shape[0]} vs measured joints {current.shape[0]}') + target = current + dq + case 'hold': + target = current + case other: + raise ValueError(f'MolmoSpaces jointpos substrate cannot map command {other!r}') + return target.astype(np.float32) + + +def resolve_camera_key(available: Any, key: str, default: str, variants: tuple[str, ...]) -> str: + """The MolmoSpaces observation key to read for a camera role, mirroring the upstream policy's precedence. + + An explicitly configured non-default key is read as-is; for the default role a present benchmark-variant + key wins over the default (matching molmo_spaces pi_policy). Raises with the candidate list on a miss. + """ + keys = set(available) + if key != default: + if key not in keys: + raise KeyError(f'observation has no camera {key!r}; available: {sorted(keys)}') + return key + for candidate in (*variants, key): + if candidate in keys: + return candidate + raise KeyError(f'observation has none of {(*variants, key)}; available: {sorted(keys)}') diff --git a/positronic/simulator/molmo_spaces/tests/__init__.py b/positronic/simulator/molmo_spaces/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/positronic/simulator/molmo_spaces/tests/droid_obs.npz b/positronic/simulator/molmo_spaces/tests/droid_obs.npz new file mode 100644 index 0000000000000000000000000000000000000000..217586f4399b14d46e903a4b04e74ff378cb05a5 GIT binary patch literal 1528 zcmWIWW@gc4U|`??Vnv1n5APiM4+XpoA`DsinRzAg1^LB#c?Fe>3<3;ufXYB13PLf= zek1-$THxe~fMWq`=S_;56SE|5>4Ln;X^Yk^UNCQ3d`Qrg8RGNfr%d7!D)&C4>Brh7 zo~|Ocgz2Z$GpkuvTwA$B81&c_EdE?gw(q!C=-I6QGXprqnTNFvsYzwXs#R zsW6~9B>RZWG?+s;fDTD7$}E66q6KQX8>Nnbq#s5-Ha50uH2dExap%MA7XaE{UX)o} z5}%xyn_83zb4PilHSGJ6tIbV z#?)SM}%aT{+9O$eS&<_6~b&XW#XhwHepjy$hDd%17D1&o1w& zyY=|s+`nIDFWM5Mw^n54)ZKrIUcTl#4t)Q%qUF%G^ zR%pgY*9yvvpqvKEb=*J<0}YHuKzot0Cb~vY&H?2LP!8dN=>pOXjQgM(VfhGMCn%|d zk}@b+gUV1CXkc__Vt^-nbX}kX3QBmOL?{5$0^&0s2kJskuIRc!$pgFYc4nYA(Gp95 TH!B;+3U(lz4fKsW3y22*h1lLz literal 0 HcmV?d00001 diff --git a/positronic/simulator/molmo_spaces/tests/make_fixture.py b/positronic/simulator/molmo_spaces/tests/make_fixture.py new file mode 100644 index 000000000..a3b34976c --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/make_fixture.py @@ -0,0 +1,53 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +"""Regenerate the synthetic MolmoSpaces raw-observation fixture used by test_adapter.py. + +``env.py`` reports the clean raw payload the adapter maps into the canonical contract — measured joints and +velocities, the eef world pose, the grip closure, and one frame per camera. MolmoSpaces renders real MuJoCo +scenes needing the full asset stack and a GPU, so committing a real payload is impractical; the adapter under +test only touches observation *structure* (keys, shapes, dtypes, the MujocoFrankaState assembly, camera key +mapping), which a tiny synthetic payload exercises exactly. Frames are small (36x64, the DROID 16:9 aspect) and +color-marked so a wrist/exterior swap is visible. + +Run: uv run --no-project positronic/simulator/molmo_spaces/tests/make_fixture.py +Output: droid_obs.npz next to this script (well under 100 KB) +""" + +from pathlib import Path + +import numpy as np + +RIG_HEIGHT, RIG_WIDTH = 36, 64 # (H, W); DROID exo/wrist cameras are 16:9. + + +def _marked_frame(base_rgb: tuple[int, int, int]) -> np.ndarray: + """A solid-color frame with a white top-left block — an orientation marker a flip or swap would move.""" + frame = np.zeros((RIG_HEIGHT, RIG_WIDTH, 3), dtype=np.uint8) + frame[:] = base_rgb + frame[:8, :12] = 255 + return frame + + +def build_payload() -> dict[str, np.ndarray]: + return { + 'joint_pos': np.array([0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785], dtype=np.float32), + 'joint_vel': np.linspace(-0.2, 0.2, 7, dtype=np.float32), + 'eef_pos': np.array([0.4, 0.0, 0.35], dtype=np.float32), + # Identity orientation, scalar-first (wxyz) as env.py reports via mju_mat2Quat. + 'eef_quat': np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + 'grip': np.float32(0.5), + 'wrist_camera': _marked_frame((200, 40, 40)), # reddish wrist view + 'exo_camera_1': _marked_frame((40, 160, 40)), # greenish exterior view + } + + +def main() -> None: + out = Path(__file__).parent / 'droid_obs.npz' + np.savez_compressed(out, **build_payload()) + print(f'Wrote {out} ({out.stat().st_size} bytes)') + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/tests/test_adapter.py b/positronic/simulator/molmo_spaces/tests/test_adapter.py new file mode 100644 index 000000000..7a80b8edc --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_adapter.py @@ -0,0 +1,70 @@ +"""Unit tests for ``MolmoAdapter``: the raw env-server payload -> canonical embodiment contract. + +Runs without molmo_spaces (the env server lives in its own interpreter); it needs positronic, which is where +the adapter runs. Exercises the observation mapping against a synthetic raw payload (``droid_obs.npz``), the +terminal, and the reset token. + +Run: uv run --locked pytest positronic/simulator/molmo_spaces/tests/test_adapter.py --no-cov +""" + +from pathlib import Path + +import numpy as np + +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.molmo_spaces.adapter import MolmoAdapter + +FIXTURE = Path(__file__).parent / 'droid_obs.npz' +CAMERA_DICT = {'image.wrist': mapping.MOLMO_WRIST_CAMERA, 'image.exterior': mapping.MOLMO_EXTERIOR_CAMERA} + + +def _payload() -> dict: + return dict(np.load(FIXTURE).items()) + + +def test_observations_assemble_robot_state(): + payload = _payload() + obs = MolmoAdapter(CAMERA_DICT).observations(payload) + state = obs['robot_state'] + assert np.allclose(state.q, payload['joint_pos']) + assert np.allclose(state.dq, payload['joint_vel']) + assert np.allclose(state.ee_pose.translation, payload['eef_pos']) + assert np.allclose(state.ee_pose.rotation.as_quat, payload['eef_quat']) # wxyz round-trips + assert obs['grip'] == 0.5 + + +def test_observations_camera_passthrough_no_swap(): + payload = _payload() + obs = MolmoAdapter(CAMERA_DICT).observations(payload) + # Frames pass through untouched (no resize/flip — the codec/client own preprocessing/transport). + assert np.array_equal(obs['image.wrist'].array, payload['wrist_camera']) + assert np.array_equal(obs['image.exterior'].array, payload['exo_camera_1']) + # Fixture marks wrist reddish, exterior greenish; a swap would flip the dominant channel. + wrist_mean = obs['image.wrist'].array.reshape(-1, 3).mean(axis=0) + exterior_mean = obs['image.exterior'].array.reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] + assert exterior_mean[1] > exterior_mean[0] + + +def test_observations_resolve_benchmark_variant_camera(): + # A Zed-wrist benchmark replaces the default key; the adapter must still land the reddish wrist view on + # image.wrist (regression: hard indexing KeyErrored on those observations). + payload = _payload() + payload['wrist_camera_zed_mini'] = payload.pop('wrist_camera') + obs = MolmoAdapter(CAMERA_DICT).observations(payload) + wrist_mean = obs['image.wrist'].array.reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] + + +def test_terminal_reports_success_only_when_done(): + adapter = MolmoAdapter(CAMERA_DICT) + assert adapter.terminal({'done': True, 'success': True}) == {'eval.success': True} + assert adapter.terminal({'done': True, 'success': False}) == {'eval.success': False} + assert adapter.terminal({'done': False, 'success': False}) is None + + +def test_reset_token_carries_episode_and_seed(): + adapter = MolmoAdapter(CAMERA_DICT) + assert adapter.reset_token({'eval.episode_index': 3, 'eval.seed': 7}) == {'episode_index': 3, 'seed': 7} + # An absent seed falls back to the spec's own (None here). + assert adapter.reset_token({'eval.episode_index': 2}) == {'episode_index': 2, 'seed': None} diff --git a/positronic/simulator/molmo_spaces/tests/test_mapping.py b/positronic/simulator/molmo_spaces/tests/test_mapping.py new file mode 100644 index 000000000..a045512f6 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_mapping.py @@ -0,0 +1,86 @@ +"""Unit tests for the pure MolmoSpaces <-> wire mappings. + +Runs with NEITHER molmo_spaces nor positronic's heavy stack: ``mapping`` imports only numpy, so these pin the +gripper normalization, the wire-command -> joint-target integration, and the camera-key precedence without a +sim or a GPU. + +Run: uv run --locked pytest positronic/simulator/molmo_spaces/tests/test_mapping.py --no-cov +""" + +import numpy as np +import pytest + +from positronic.simulator.molmo_spaces import mapping + + +def test_grip_qpos_normalization(): + closed = mapping.GRIPPER_QPOS_CLOSED + assert mapping.normalize_grip_qpos(0.0) == 0.0 + assert abs(mapping.normalize_grip_qpos(closed / 2) - 0.5) < 1e-6 + assert abs(mapping.normalize_grip_qpos(closed) - 1.0) < 1e-6 + assert mapping.normalize_grip_qpos(closed * 2) == 1.0 # saturates, never exceeds 1 + # A two-finger qpos reads the first finger. + assert abs(mapping.normalize_grip_qpos(np.array([closed / 2, closed / 2])) - 0.5) < 1e-6 + + +def test_grip_command_to_actuator(): + assert mapping.grip_command_to_actuator(0.0) == mapping.ROBOTIQ_OPEN == 0.0 + assert mapping.grip_command_to_actuator(1.0) == mapping.ROBOTIQ_CLOSED == 255.0 + assert mapping.grip_command_to_actuator(0.5) == 127.5 # continuous — the codec owns binarization + assert mapping.grip_command_to_actuator(2.0) == 255.0 # clipped + + +def test_wire_command_joint_pos_passthrough(): + current = np.arange(mapping.NUM_ARM_JOINTS, dtype=np.float32) + q = np.full(mapping.NUM_ARM_JOINTS, 0.3, dtype=np.float32) + out = mapping.wire_command_to_arm_action({'type': 'joint_pos', 'q': q}, current) + assert out.dtype == np.float32 and out.shape == (mapping.NUM_ARM_JOINTS,) + assert np.array_equal(out, q) # absolute target, independent of the measured joints + + +def test_wire_command_joint_vel_integrates_onto_measured(): + current = np.arange(mapping.NUM_ARM_JOINTS, dtype=np.float32) + dq = np.full(mapping.NUM_ARM_JOINTS, 0.1, dtype=np.float32) + out = mapping.wire_command_to_arm_action({'type': 'joint_vel', 'dq': dq}, current) + assert np.allclose(out, current + dq) # positronic applies JointDelta as q + dq + + +def test_wire_command_hold_recommands_measured(): + current = np.linspace(-1.0, 1.0, mapping.NUM_ARM_JOINTS, dtype=np.float32) + out = mapping.wire_command_to_arm_action({'type': 'hold'}, current) + assert np.array_equal(out, current) + + +def test_wire_command_joint_count_mismatch_raises(): + current = np.zeros(mapping.NUM_ARM_JOINTS, dtype=np.float32) + with pytest.raises(ValueError): + mapping.wire_command_to_arm_action({'type': 'joint_vel', 'dq': np.zeros(6, dtype=np.float32)}, current) + + +def test_wire_command_cartesian_unsupported(): + current = np.zeros(mapping.NUM_ARM_JOINTS, dtype=np.float32) + with pytest.raises(ValueError): + mapping.wire_command_to_arm_action({'type': 'cartesian', 'pose': np.zeros(12)}, current) + + +def test_camera_key_default_and_variant_precedence(): + default = mapping.MOLMO_WRIST_CAMERA + variants = mapping.MOLMO_WRIST_CAMERA_VARIANTS + # Default present, no variant -> the default. + assert mapping.resolve_camera_key({default: 1}, default, default, variants) == default + # A benchmark-variant key present wins over the default (matches molmo_spaces pi_policy precedence). + both = {default: 1, variants[0]: 1} + assert mapping.resolve_camera_key(both, default, default, variants) == variants[0] + # Variant only (default absent) -> the variant. + assert mapping.resolve_camera_key({variants[0]: 1}, default, default, variants) == variants[0] + + +def test_camera_key_explicit_nondefault_read_as_is(): + # An explicitly configured non-default key is read as-is, never shadowed by a variant decoy. + obs = {'my_cam': 1, mapping.MOLMO_WRIST_CAMERA_VARIANTS[0]: 1} + assert mapping.resolve_camera_key(obs, 'my_cam', mapping.MOLMO_WRIST_CAMERA, ()) == 'my_cam' + + +def test_camera_key_miss_raises(): + with pytest.raises(KeyError): + mapping.resolve_camera_key({'other': 1}, mapping.MOLMO_WRIST_CAMERA, mapping.MOLMO_WRIST_CAMERA, ()) diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 9ebf7957a..9e2cc31b6 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -138,7 +138,8 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came # Pretrained DROID models read joints and gripper as separate observation keys and the language -# prompt under `prompt` (see openpi `droid_policy.DroidInputs`). +# prompt under `prompt` (see openpi `droid_policy.DroidInputs`), lowercased — the checkpoints were +# trained on lowercased language and MolmoSpaces' Pi baseline normalizes the same way. droid_obs = cfn.Config( GenericObservationCodec, state={'observation/joint_position': {'robot_state.q': 7}, 'observation/gripper_position': {'grip': 1}}, @@ -147,6 +148,7 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came 'observation/exterior_image_1_left': ('image.exterior', (224, 224)), }, task_field='prompt', + lowercase_task=True, ) ee = codecs.compose.override(obs=ee_obs, action=codecs.absolute_pos_action) diff --git a/pyproject.toml b/pyproject.toml index 457578f01..fcdfd8d7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,7 @@ testpaths = [ "positronic/offboard/tests", "positronic/policy/tests", "positronic/simulator/env_server/tests", + "positronic/simulator/molmo_spaces/tests", "positronic/tests", "positronic/utils/tests", "positronic/vendors/gr00t/tests", From 4e1e345190091879271133fb1823914bd62caa28 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 21:33:48 +0000 Subject: [PATCH 06/23] Preserve benchmark episode seeds when no eval seed is supplied `build_trials` injected a random `eval.seed` per trial when the catalog seed was unset, clobbering each benchmark episode's own spec seed in the env (the `episode.seed` fallback never fired) and making runs non-reproducible. Build the trials directly: omit `eval.seed` unless the catalog seed is set, so the env uses the episode's spec seed; an explicit seed still overrides. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/cfg/eval/sim/molmo.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index 97cfbc844..8c13b2f21 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -3,7 +3,6 @@ import configuronic as cfn -from positronic.cfg.eval import build_trials from positronic.drivers.roboarm.models import bundled_franka_model from positronic.eval import Eval, Task from positronic.simulator.env_server.proxy import RemoteEnvControlSystem, remote_franka_embodiment @@ -52,8 +51,18 @@ def _molmo_eval(benchmark_dir, episodes, trial_count, timeout, camera_dict, seed proxy, camera_dict, descriptor='remote.molmo_spaces.droid', static_meta=bundled_franka_model() ) task = Task(instruction=lambda: proxy.meta['task'], timeout=timeout, reset=proxy.reset, done=proxy.done) - scenes = [{'eval.episode_index': i} for i in indices] - return Eval(embodiment, task, build_trials(seed, trial_count, scenes)) + # Benchmark episodes are exact-pose deterministic and carry their own seed. An unset ``seed`` leaves + # ``eval.seed`` off the trial, so the env falls back to the episode's spec seed (reproducing the benchmark); + # an explicit ``seed`` overrides it, sweeping ``seed .. seed + trial_count - 1``. (``build_trials`` injects a + # random seed when ``seed`` is None, which would clobber the spec seed and make the run non-reproducible.) + trials = [ + {'eval.episode_index': i, **({'eval.seed': seed + t} if seed is not None else {})} + for i in indices + for t in range(trial_count) + ] + for j, ctx in enumerate(trials): + ctx.update({'eval.trial_index': j, 'eval.trial_count': len(trials)}) + return Eval(embodiment, task, trials) # The whole benchmark in one run (every episode in ``--eval.benchmark_dir``'s benchmark.json). From 47f7c163c63e1afb87d6be9ad681c05818ba831e Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 21:57:01 +0000 Subject: [PATCH 07/23] Count both benchmark layouts and lowercase the training task too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Codex on #504: - `_episode_count` mirrors `load_all_episodes`' two layouts (benchmark.json, else legacy `house_*/episode_*.json`), so a whole-benchmark run no longer fails in the client for legacy dirs. - `ObservationCodec` lowercases the derived training task when `lowercase_task` is set, matching the served prompt — so a DROID codec trains and infers on one text distribution (the same-keys contract). Ticket: Positronic-Robotics/internal#91 #refs --- positronic/cfg/eval/sim/molmo.py | 13 +++++++++++-- positronic/policy/observation.py | 10 ++++++++-- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index 8c13b2f21..6409fab77 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -11,8 +11,17 @@ def _episode_count(benchmark_dir: str) -> int: - """The number of episodes in a MolmoSpaces ``benchmark.json`` (a JSON list of episode specs).""" - return len(json.loads((Path(benchmark_dir) / 'benchmark.json').read_text())) + """The episode count of a MolmoSpaces benchmark dir, mirroring ``load_all_episodes``' two layouts. + + positronic cannot import ``molmo_spaces`` here (it lives in the env server's own venv), so this counts the + benchmark files directly: a single ``benchmark.json`` (a JSON list of episode specs) when present, else the + legacy ``house_*/episode_*.json`` layout the loader also accepts. + """ + base = Path(benchmark_dir) + manifest = base / 'benchmark.json' + if manifest.exists(): + return len(json.loads(manifest.read_text())) + return sum(1 for _ in base.glob('house_*/episode_*.json')) @cfn.config( diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index e7709fd27..0467ee29d 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -7,7 +7,7 @@ from positronic.dataset import Signal, transforms from positronic.dataset.episode import Episode from positronic.dataset.transforms import image -from positronic.dataset.transforms.episode import Derive, Get +from positronic.dataset.transforms.episode import Derive from positronic.policy.codec import Codec, lerobot_image, lerobot_state @@ -36,7 +36,9 @@ def __init__( self._derive_transforms = {k: partial(self._derive_state, k) for k in state.keys()} self._derive_transforms.update({k: partial(self._derive_image, k) for k in images.keys()}) - self._derive_transforms['task'] = Get('task', '') + # Lowercase the training task the same way ``encode`` lowercases the served prompt, so a codec with + # ``lowercase_task`` trains and infers on one text distribution (the ``Codec`` same-keys contract). + self._derive_transforms['task'] = self._derive_task lerobot_features: dict[str, Any] = {} for name, features in state.items(): @@ -54,6 +56,10 @@ def _derive_image(self, out_name: str, episode: Episode) -> Signal[Any]: input_key, (width, height) = self._image_configs[out_name] return image.resize_with_pad(width, height, signal=episode[input_key]) + def _derive_task(self, episode: Episode) -> Any: + task = episode['task'] if 'task' in episode else '' + return task.lower() if self._lowercase_task else task + def _decode_single(self, data: dict, context: dict | None) -> dict: return {} From 612bde39dbc9f89103e4dccd5987bc9cd61de805 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 21:57:01 +0000 Subject: [PATCH 08/23] Install the molmo mujoco extra without curobo; fix env.py import order Two problems the CPU smoke surfaced: - `uv sync --extra mujoco` also resolves the `curobo` extra (a CUDA build needing a GPU toolchain, not on the eval path), which fails on a GPU-less box. Install the `mujoco` extra into a venv with `uv pip install -e .[mujoco]` the way MolmoSpaces' own image does, and run env.py with that venv's python. - Importing `json_eval_task_sampler` directly hits a circular import; import `json_eval_runner` first to break it. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/env.py | 9 +++-- positronic/simulator/molmo_spaces/launcher.py | 39 +++++++++---------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index 9c73c0c59..acc4e3671 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -1,9 +1,9 @@ """MolmoSpaces — AllenAI's MuJoCo manipulation benchmark — behind the env-server protocol. -MolmoSpaces pins ``mujoco ~=3.5`` + its asset stack into its own uv project on Python 3.11, so this never shares -positronic's venv: the launcher runs ``uv run --project env.py --host ... --port ...`` with the -positronic-free ``server``/``protocol`` and this package's ``mapping`` module on ``PYTHONPATH``. It imports only -``molmo_spaces`` (+ mujoco/numpy) and those, never ``positronic``. +MolmoSpaces pins ``mujoco ~=3.5`` + its asset stack on Python 3.11, so this never shares positronic's venv: the +launcher runs it with the molmospaces ``.venv``'s python (``env.py --host ... --port ... --benchmark_dir ...``), +with the positronic-free ``server``/``protocol`` and this package's ``mapping`` module on ``PYTHONPATH``. It +imports only ``molmo_spaces`` (+ mujoco/numpy) and those, never ``positronic``. positronic owns the control loop: this server drives a single MolmoSpaces ``BaseMujocoTask`` per episode directly (``JsonEvalTaskSampler.sample_task`` builds the full sim/scene/renderer; ``reset``/``step``/``is_done``/ @@ -34,6 +34,7 @@ import numpy as np # noqa: E402 from server import EnvProtocol, EnvServer # noqa: E402 +import molmo_spaces.evaluation.json_eval_runner # noqa: E402, F401 -- load first: breaks a circular import that importing json_eval_task_sampler directly hits from molmo_spaces.configs.policy_configs import DummyPolicyConfig # noqa: E402 from molmo_spaces.configs.robot_configs import ActionNoiseConfig, FrankaRobotConfig # noqa: E402 from molmo_spaces.evaluation.benchmark_schema import load_all_episodes # noqa: E402 diff --git a/positronic/simulator/molmo_spaces/launcher.py b/positronic/simulator/molmo_spaces/launcher.py index f6d346b3b..961c94984 100644 --- a/positronic/simulator/molmo_spaces/launcher.py +++ b/positronic/simulator/molmo_spaces/launcher.py @@ -1,11 +1,10 @@ """Launches the MolmoSpaces env server as a subprocess and owns its lifetime. -positronic starts the server: the env runs in MolmoSpaces' own interpreter via ``uv run --project ``, -which resolves the ``molmospaces[mujoco]`` stack (mujoco ~=3.5, the resource-manager asset layer, torch) from the -pinned checkout — far too heavy, and Python-version-pinned (3.11), to share positronic's venv. The positronic-free -``env_server`` package and this package's ``mapping`` module ride ``PYTHONPATH`` so ``env.py`` imports the dumb -``server``/``protocol`` and the pure wire mappings without dragging in positronic; ``molmo_spaces`` itself resolves -from the uv project. +positronic starts the server: the env runs in MolmoSpaces' own interpreter — a per-checkout ``.venv`` with the +``molmospaces[mujoco]`` stack (mujoco ~=3.5, the resource-manager asset layer, torch) installed into it, far too +heavy and Python-version-pinned (3.11) to share positronic's venv. The positronic-free ``env_server`` package and +this package's ``mapping`` module ride ``PYTHONPATH`` so ``env.py`` imports the dumb ``server``/``protocol`` and +the pure wire mappings without dragging in positronic; ``molmo_spaces`` resolves from the venv. MolmoSpaces renders MuJoCo scenes, so the server needs a GL backend (``MUJOCO_GL``) and its asset packs (``MLSPACES_ASSETS_DIR``): ``MUJOCO_GL`` defaults to ``egl`` (GPU) here and both env vars pass through from the @@ -47,25 +46,24 @@ def _checkout_lock() -> Iterator[None]: def _spawn(host: str, port: int, benchmark_dir: str) -> subprocess.Popen: + venv = _MOLMO_SRC / '.venv' with _checkout_lock(): src = ensure_pinned_checkout(_MOLMO_REPO, _MOLMO_COMMIT, _MOLMO_SRC) # Install the stack before spawning: a cold first install far exceeds the client's connect deadline, - # which should only cover the sim's boot. Idempotent and fast when warm; the spawn passes ``--no-sync`` - # so no resolve or install ever runs outside this lock. MolmoSpaces ships no uv.lock, so this re-resolves - # on every fresh box. + # which should only cover the sim's boot. Install the ``mujoco`` extra explicitly into a venv the way + # MolmoSpaces' own image does, rather than ``uv sync`` — which also resolves the ``curobo`` extra, a CUDA + # build that needs a GPU toolchain and is not on the eval task path. Both steps are idempotent and fast + # when warm. MolmoSpaces ships no uv.lock, so the install re-resolves on every fresh box. + if not venv.exists(): + subprocess.run(['uv', 'venv', '--python', _MOLMO_PYTHON, str(venv)], check=True) subprocess.run( - ['uv', 'sync', '--project', str(src), '--python', _MOLMO_PYTHON, '--extra', _MOLMO_EXTRA], check=True + ['uv', 'pip', 'install', '-e', f'.[{_MOLMO_EXTRA}]'], + cwd=str(src), + env={**os.environ, 'VIRTUAL_ENV': str(venv)}, + check=True, ) command = [ - 'uv', - 'run', - '--no-sync', - '--project', - str(src), - '--python', - _MOLMO_PYTHON, - '--extra', - _MOLMO_EXTRA, + str(venv / 'bin' / 'python'), str(_ENV_SCRIPT), '--host', host, @@ -77,7 +75,8 @@ def _spawn(host: str, port: int, benchmark_dir: str) -> subprocess.Popen: env = { **os.environ, 'PYTHONPATH': os.pathsep.join([str(_ENV_SERVER_DIR), str(_MAPPING_DIR)]), - # GPU OpenGL by default; a caller on a GPU-less box exports MUJOCO_GL=osmesa for CPU software rendering. + # GPU OpenGL by default; a GPU-less box exports MUJOCO_GL=osmesa, or relies on mesa's software EGL, for + # CPU rendering. 'MUJOCO_GL': os.environ.get('MUJOCO_GL', 'egl'), } return subprocess.Popen(command, env=env) From be622e51bfbafdb5d114589ecf9f647b5bb7b5f3 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 22:07:07 +0000 Subject: [PATCH 09/23] Fail fast when a MolmoSpaces benchmark has no episodes An empty or mislaid `--eval.benchmark_dir` (no benchmark.json and no legacy `house_*/episode_*.json`) produced an empty episode sweep and a silent zero-trial eval. Raise a clear error instead. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/cfg/eval/sim/molmo.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index 6409fab77..4513b7919 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -52,6 +52,11 @@ def _molmo_eval(benchmark_dir, episodes, trial_count, timeout, camera_dict, seed indices = list(range(_episode_count(benchmark_dir))) else: indices = [episodes] if isinstance(episodes, int) else list(episodes) + if not indices: + raise ValueError( + f'no benchmark episodes found under {benchmark_dir!r}; expected a benchmark.json or a legacy ' + 'house_*/episode_*.json layout (or pass --eval.episodes explicitly)' + ) proxy = RemoteEnvControlSystem(MolmoAdapter(camera_dict), serve_molmo_spaces(benchmark_dir)) # MolmoSpaces drives a Franka DROID rig; recordings carry the same model (URDF + meshes + joint names + # control frame) for the 3D viewer and offline IK, supplied here since the molmo server can't import From 373041be55218ce18bd7334313bed012b31129d4 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 22:17:56 +0000 Subject: [PATCH 10/23] Resolve the RandCam exterior camera variant The default camera_dict KeyErrored on RandCam benchmarks, which record the exterior as `randomized_zed2_analogue_1` (verified: camera_configs.py, and the `--camera_names` exterior in mb-bench.md), not `exo_camera_1` or the light-randomization variant. Add it to the exterior fallback so the default mapping covers the light-randomization and RandCam suites alike. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/mapping.py | 9 ++++++--- positronic/simulator/molmo_spaces/tests/test_mapping.py | 9 +++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/positronic/simulator/molmo_spaces/mapping.py b/positronic/simulator/molmo_spaces/mapping.py index 055e6619b..52a3a15ef 100644 --- a/positronic/simulator/molmo_spaces/mapping.py +++ b/positronic/simulator/molmo_spaces/mapping.py @@ -16,12 +16,15 @@ MOLMO_ARM_GROUP = 'arm' MOLMO_GRIPPER_GROUP = 'gripper' -# MolmoSpaces DROID rig camera names (FrankaDroidCameraSystem); the Zed-wrist / light-randomized benchmark -# variants replace the defaults and MolmoSpaces' own Pi policy prefers them when present, so the adapter must too. +# MolmoSpaces DROID rig camera names (FrankaDroidCameraSystem); a benchmark's own variants replace the defaults +# and the adapter resolves them so the default camera_dict works across the benchmarks: the light-randomization +# suite records the exterior as ``droid_shoulder_light_randomization`` (MolmoSpaces' Pi policy prefers it), and +# the RandCam suite records it as ``randomized_zed2_analogue_1`` (its ``--camera_names`` exterior); the Zed wrist +# variant is ``wrist_camera_zed_mini``. MOLMO_WRIST_CAMERA = 'wrist_camera' MOLMO_EXTERIOR_CAMERA = 'exo_camera_1' MOLMO_WRIST_CAMERA_VARIANTS = ('wrist_camera_zed_mini',) -MOLMO_EXTERIOR_CAMERA_VARIANTS = ('droid_shoulder_light_randomization',) +MOLMO_EXTERIOR_CAMERA_VARIANTS = ('droid_shoulder_light_randomization', 'randomized_zed2_analogue_1') # The Robotiq 2F-85 finger qpos saturates at this closure; the DROID observation's grip is normalized against # it into the [0, 1] closure the policy was trained on (molmospaces pi_policy.py:126). diff --git a/positronic/simulator/molmo_spaces/tests/test_mapping.py b/positronic/simulator/molmo_spaces/tests/test_mapping.py index a045512f6..260f02b93 100644 --- a/positronic/simulator/molmo_spaces/tests/test_mapping.py +++ b/positronic/simulator/molmo_spaces/tests/test_mapping.py @@ -84,3 +84,12 @@ def test_camera_key_explicit_nondefault_read_as_is(): def test_camera_key_miss_raises(): with pytest.raises(KeyError): mapping.resolve_camera_key({'other': 1}, mapping.MOLMO_WRIST_CAMERA, mapping.MOLMO_WRIST_CAMERA, ()) + + +def test_exterior_camera_variants_cover_light_randomization_and_randcam(): + # The default exterior mapping must resolve both benchmark exterior names (regression: RandCam records + # randomized_zed2_analogue_1, not exo_camera_1, so hard indexing KeyErrored). + default = mapping.MOLMO_EXTERIOR_CAMERA + variants = mapping.MOLMO_EXTERIOR_CAMERA_VARIANTS + for name in ('droid_shoulder_light_randomization', 'randomized_zed2_analogue_1'): + assert mapping.resolve_camera_key({name: 1}, default, default, variants) == name From d3f8224af3ff194ff5451640649034ff4fb8699b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 22:25:06 +0000 Subject: [PATCH 11/23] Stub macOS CGL so the molmo env server renders on a CPU box MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MolmoSpaces' renderer hardcodes a macOS CGL context on the CPU render path and dlopens Apple's OpenGL.framework, crashing at renderer init on Linux — so the launcher's advertised CPU rendering (MUJOCO_GL=osmesa / mesa software EGL) died before the first observation. Stub the no-op CGL module so it resolves; untouched on a GPU box. Validated end-to-end: a FrankaPickDroidMiniBench episode drives reset + steps on CPU (mesa software EGL), obs payload + camera frames intact. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/env.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index acc4e3671..a824180af 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -20,12 +20,33 @@ import argparse import os +import sys +import types # MolmoSpaces renders MuJoCo scenes, so the GL backend must be selected before any mujoco/molmo_spaces import. # The launcher sets MUJOCO_GL in the subprocess env (egl by default); default it here too so a direct invocation # (e.g. a validate/e2e run) still boots. Set before the imports below. os.environ.setdefault('MUJOCO_GL', 'egl') + +def _install_cgl_noop_stub() -> None: + # HACK: MolmoSpaces' renderer hardcodes a macOS CGL context on the CPU (device_id=None) render path + # (opengl_rendering.py does ``from mujoco.cgl import cgl``), which dlopens Apple's OpenGL.framework and + # crashes at renderer init on Linux — so a CPU-rendered server (MUJOCO_GL=osmesa or mesa software EGL) + # dies before the first observation. CGL locking is a no-op off macOS, so stub the module: the import + # resolves and the (un)lock does nothing. Untouched on a GPU box, where the EGL path never imports it. + if 'mujoco.cgl' in sys.modules: + return + cgl = types.ModuleType('mujoco.cgl.cgl') + cgl.CGLLockContext = cgl.CGLUnlockContext = lambda *args, **kwargs: None + package = types.ModuleType('mujoco.cgl') + package.cgl = cgl + sys.modules['mujoco.cgl'] = package + sys.modules['mujoco.cgl.cgl'] = cgl + + +_install_cgl_noop_stub() + from pathlib import Path # noqa: E402 from typing import Any # noqa: E402 From f55da7a6c193f9e5dcce5d2296644b04845cd56c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 22:39:00 +0000 Subject: [PATCH 12/23] Let positronic own the episode deadline; add the wire e2e MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Disable MolmoSpaces' internal step horizon (Codex): `JsonBenchmarkEvalConfig` defaults it to 500 steps (~33s at the 66ms policy period), which would self-terminate the task before the harness timeout and truncate the score — the env stealing the loop's authority over episode length. `task_horizon=None` runs to an infinite horizon; `is_done` then reports only the task's own success, and positronic's harness stops the trial at its `--eval.timeout`. - Add `e2e.py`: the launcher spawns the env server and a client resets + steps it over the real socket, mapping each frame through `MolmoAdapter`. Validated end-to-end on CPU (mesa software EGL) against FrankaPickDroidMiniBench. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/e2e.py | 61 ++++++++++++++++++++++++ positronic/simulator/molmo_spaces/env.py | 9 ++-- 2 files changed, 67 insertions(+), 3 deletions(-) create mode 100644 positronic/simulator/molmo_spaces/e2e.py diff --git a/positronic/simulator/molmo_spaces/e2e.py b/positronic/simulator/molmo_spaces/e2e.py new file mode 100644 index 000000000..c226f9c26 --- /dev/null +++ b/positronic/simulator/molmo_spaces/e2e.py @@ -0,0 +1,61 @@ +"""End-to-end check that the MolmoSpaces env server works over the socket + the adapter maps its payload. + +The mapping/adapter unit tests exercise the transforms in-process; this drives the **real boundary**: the +launcher spawns the env-server subprocess in MolmoSpaces' own venv, and a client resets + steps it over the +actual socket, then feeds the wire payload through ``MolmoAdapter`` — validating the launcher, the wire +protocol, ``env.py``'s task drive, and the observation mapping together. A hold command holds the arm, so a +healthy server keeps the joints steady and reports frames of the right shape; a wrong obs key, quaternion +order, or a broken wire codec fails the mapping. + +Needs the MolmoSpaces asset packs (``MLSPACES_ASSETS_DIR``) and a GL backend (``MUJOCO_GL``; a GPU-less box uses +mesa software EGL — ``EGL_PLATFORM=surfaceless LIBGL_ALWAYS_SOFTWARE=1``). Run on a box with those:: + + MLSPACES_ASSETS_DIR=... MUJOCO_GL=egl EGL_PLATFORM=surfaceless LIBGL_ALWAYS_SOFTWARE=1 \ + uv run --locked python -m positronic.simulator.molmo_spaces.e2e --benchmark_dir +""" + +import argparse + +from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.molmo_spaces.adapter import MolmoAdapter +from positronic.simulator.molmo_spaces.launcher import serve_molmo_spaces + +_CAMERA_DICT = {'image.wrist': 'wrist_camera', 'image.exterior': 'exo_camera_1'} + + +def run(benchmark_dir: str, *, episodes: int = 1, steps: int = 5, camera_dict: dict[str, str] | None = None) -> None: + """Reset + step the first ``episodes`` benchmark episodes over the socket, mapping each frame with the adapter.""" + camera_dict = camera_dict or _CAMERA_DICT + adapter = MolmoAdapter(camera_dict) + with serve_molmo_spaces(benchmark_dir) as (host, port): + conn = EnvConnection(host, port) + try: + for i in range(episodes): + frame = conn.reset({'episode_index': i, 'seed': None}) + obs = adapter.observations(frame['obs']) + assert 'robot_state' in obs and 'grip' in obs, f'missing contract keys: {sorted(obs)}' + assert all(logical in obs for logical in camera_dict), f'missing cameras: {sorted(obs)}' + q = obs['robot_state'].q + assert q.shape == (7,), f'unexpected joint shape {q.shape}' + print(f' episode {i}: reset ok — task={frame["meta"]["task"]!r} grip={obs["grip"]:.3f} q0={q[0]:.4f}') + out = {'done': False} + for _ in range(steps): + out = conn.step({'command': {'type': 'hold'}, 'grip': 0.0}) + adapter.observations(out['obs']) # the mapping round-trips on step frames too + print(f' episode {i}: {steps} steps ok (done={out["done"]})') + finally: + conn.close() + print('E2E PASSED') + + +def main() -> None: + parser = argparse.ArgumentParser(description='Drive the MolmoSpaces env server over the socket.') + parser.add_argument('--benchmark_dir', required=True, help='dir containing benchmark.json') + parser.add_argument('--episodes', type=int, default=1) + parser.add_argument('--steps', type=int, default=5) + args = parser.parse_args() + run(args.benchmark_dir, episodes=args.episodes, steps=args.steps) + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index a824180af..4acabb8bd 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -109,9 +109,12 @@ def _build(self, episode_index: int, seed: int | None) -> None: cfg = _DroidPickEvalConfig() # Determinism enters at sampler construction (seed_task_sampling); the token's seed overrides the spec's. cfg.seed = int(seed) if seed is not None else (episode.seed if episode.seed is not None else 42) - horizon_sec = episode.task.get('task_horizon_sec') - if horizon_sec is not None: - cfg.task_horizon = round(float(horizon_sec) * 1000.0 / cfg.policy_dt_ms) + # positronic's harness owns the episode deadline (the eval's timeout), so disable MolmoSpaces' internal + # step horizon: JsonBenchmarkEvalConfig defaults it to 500 steps (~33s at the 66ms policy period), which + # would self-terminate the task before the harness timeout and truncate the score. ``None`` -> the task + # runs to an infinite horizon, so ``is_done`` reports only the task's own success/termination and the + # harness stops the trial at its timeout. + cfg.task_horizon = None self._sampler = JsonEvalTaskSampler(cfg, episode) self._task = self._sampler.sample_task(house_index=episode.house_index) self._robot_view = self._task.env.current_robot.robot_view From 7aaa7e08064baf83fe4195ae7fbaa7b9825c28f7 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 22:49:17 +0000 Subject: [PATCH 13/23] Terminate the episode on task success Disabling MolmoSpaces' step horizon (prev commit) combined with the default `terminate_upon_success=False` meant a successful rollout that keeps sending joint commands never reported `done`, so the harness scored it as a timeout with no `eval.success` (Codex P1). `env.py` now ends the episode on the task's own `judge_success()` (end-on-success, the benchmark's scoring) or any molmo terminal; positronic's harness timeout stays the outer deadline. Re-validated end-to-end over the socket. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/env.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index 4acabb8bd..1054276c9 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -134,10 +134,13 @@ def step(self, action: dict[str, Any]) -> dict[str, Any]: arm = mapping.wire_command_to_arm_action(action['command'], self._measured_arm_q()) gripper = np.array([mapping.grip_command_to_actuator(action['grip'])], dtype=np.float32) obs, _reward, _term, _trunc, _infos = self._task.step({'arm': arm, 'gripper': gripper}) - done = bool(self._task.is_done()) - # ``judge_success`` is the task's scored success; only meaningful once the episode is done, so a timeout - # stays honest (a running episode is never a success). - success = bool(self._task.judge_success()) if done else False + # positronic owns the deadline (the harness timeout), so the episode ends here on the task's own scored + # success — end-on-success, the benchmark's scoring semantics — or any MolmoSpaces terminal (a done + # action). The internal step horizon is disabled (see ``_build``), so ``is_done`` reports only + # ``is_terminal``, never a timeout; without terminating on success here a successful rollout that keeps + # sending joint commands would run to the harness timeout and be scored as a non-success. + success = bool(self._task.judge_success()) + done = success or bool(self._task.is_done()) return {'obs': self._observe(obs[0]), 'done': done, 'success': success, 'control_dt': self._control_dt} def _measured_arm_q(self) -> np.ndarray: From 507c40b7329f7f1bc6e5d263dd4c7952761f3f70 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 23:05:58 +0000 Subject: [PATCH 14/23] Type the molmo eval config; baseline the foreign-venv type errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Annotate `_molmo_eval`'s params + return and add `-> None` to the env/adapter constructors — the config function was untyped (basedpyright's standard mode doesn't flag missing annotations, so the ratchet didn't catch it). - Grandfather env.py + make_fixture.py's type errors in the basedpyright baseline, matching robolab/libero: env.py runs in molmo's own venv, so its foreign imports and dynamic mujoco attrs can't resolve against positronic's deps, and make_fixture hits numpy stub quirks. The ratchet is clean (0 new errors); mapping/adapter/e2e/catalog stay fully typed (not baselined). Ticket: Positronic-Robotics/internal#91 #refs --- .basedpyright/baseline.json | 164 +++++++++++++++++++ positronic/cfg/eval/sim/molmo.py | 9 +- positronic/simulator/molmo_spaces/adapter.py | 2 +- positronic/simulator/molmo_spaces/env.py | 2 +- 4 files changed, 174 insertions(+), 3 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 1d2f48bbc..88d2ffd2a 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -10882,6 +10882,170 @@ } } ], + "./positronic/simulator/molmo_spaces/env.py": [ + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 8, + "endColumn": 22, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 29, + "endColumn": 45, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 12, + "endColumn": 15, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 19, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 32, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 7, + "endColumn": 47, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 5, + "endColumn": 40, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 5, + "endColumn": 39, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 5, + "endColumn": 45, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 5, + "endColumn": 55, + "lineCount": 1 + } + }, + { + "code": "reportMissingImports", + "range": { + "startColumn": 5, + "endColumn": 46, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 32, + "endColumn": 37, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 57, + "endColumn": 61, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 34, + "endColumn": 47, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 42, + "endColumn": 49, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 43, + "endColumn": 57, + "lineCount": 1 + } + }, + { + "code": "reportOptionalMemberAccess", + "range": { + "startColumn": 31, + "endColumn": 45, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 15, + "endColumn": 27, + "lineCount": 1 + } + } + ], + "./positronic/simulator/molmo_spaces/tests/make_fixture.py": [ + { + "code": "reportReturnType", + "range": { + "startColumn": 16, + "endColumn": 31, + "lineCount": 1 + } + }, + { + "code": "reportArgumentType", + "range": { + "startColumn": 31, + "endColumn": 46, + "lineCount": 1 + } + } + ], "./positronic/simulator/mujoco/scene_builder.py": [ { "code": "reportIndexIssue", diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index 4513b7919..edc273d72 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -32,7 +32,14 @@ def _episode_count(benchmark_dir: str) -> int: timeout=60.0, seed=None, ) -def _molmo_eval(benchmark_dir, episodes, trial_count, timeout, camera_dict, seed): +def _molmo_eval( + benchmark_dir: str | None, + episodes: int | list[int] | None, + trial_count: int, + timeout: float, + camera_dict: dict[str, str], + seed: int | None, +) -> Eval: """A MolmoSpaces eval: the embodiment proxies a remote MolmoSpaces env, the task carries the scenario. MolmoSpaces (https://github.com/allenai/molmospaces) is AllenAI's MuJoCo manipulation benchmark on the DROID diff --git a/positronic/simulator/molmo_spaces/adapter.py b/positronic/simulator/molmo_spaces/adapter.py index 1d084e8ae..96595c85d 100644 --- a/positronic/simulator/molmo_spaces/adapter.py +++ b/positronic/simulator/molmo_spaces/adapter.py @@ -23,7 +23,7 @@ class MolmoAdapter(WireCommandAdapter): - def __init__(self, camera_dict: dict[str, str]): + def __init__(self, camera_dict: dict[str, str]) -> None: super().__init__() self._camera_dict = camera_dict # logical observation name -> the MolmoSpaces obs camera key diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index 1054276c9..79e5d8242 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -91,7 +91,7 @@ class MolmoSpacesEnv(EnvProtocol): and reports MolmoSpaces' ``is_done``/``judge_success``. """ - def __init__(self, benchmark_dir: str): + def __init__(self, benchmark_dir: str) -> None: self._episodes = load_all_episodes(Path(benchmark_dir)) self._sampler = None self._task = None From 588658d9e60129ed85077abe0d6fb599e7e39e12 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 23:23:39 +0000 Subject: [PATCH 15/23] Don't grandfather molmo's new code into the type baseline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverse the baseline entries the previous commit added — new code is never baseline-eligible. env.py runs in MolmoSpaces' own venv (its molmo_spaces + PYTHONPATH-relative server/mapping imports don't resolve against positronic's deps), so exclude it like vendors/ — checked in its own interpreter, not grandfathered here. make_fixture's numpy typing is fixed properly (return type + a targeted savez ignore). basedpyright: 0 new errors, baseline pristine. Ticket: Positronic-Robotics/internal#91 #refs --- .basedpyright/baseline.json | 164 ------------------ .../molmo_spaces/tests/make_fixture.py | 5 +- pyproject.toml | 12 +- 3 files changed, 11 insertions(+), 170 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 88d2ffd2a..1d2f48bbc 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -10882,170 +10882,6 @@ } } ], - "./positronic/simulator/molmo_spaces/env.py": [ - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 8, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 29, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 12, - "endColumn": 15, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 19, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 32, - "endColumn": 41, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 7, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 5, - "endColumn": 40, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 5, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 5, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 5, - "endColumn": 55, - "lineCount": 1 - } - }, - { - "code": "reportMissingImports", - "range": { - "startColumn": 5, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 32, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 57, - "endColumn": 61, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 34, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 42, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 43, - "endColumn": 57, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 15, - "endColumn": 27, - "lineCount": 1 - } - } - ], - "./positronic/simulator/molmo_spaces/tests/make_fixture.py": [ - { - "code": "reportReturnType", - "range": { - "startColumn": 16, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 46, - "lineCount": 1 - } - } - ], "./positronic/simulator/mujoco/scene_builder.py": [ { "code": "reportIndexIssue", diff --git a/positronic/simulator/molmo_spaces/tests/make_fixture.py b/positronic/simulator/molmo_spaces/tests/make_fixture.py index a3b34976c..052287792 100644 --- a/positronic/simulator/molmo_spaces/tests/make_fixture.py +++ b/positronic/simulator/molmo_spaces/tests/make_fixture.py @@ -16,6 +16,7 @@ """ from pathlib import Path +from typing import Any import numpy as np @@ -30,7 +31,7 @@ def _marked_frame(base_rgb: tuple[int, int, int]) -> np.ndarray: return frame -def build_payload() -> dict[str, np.ndarray]: +def build_payload() -> dict[str, Any]: # grip is a float32 scalar, the rest are arrays return { 'joint_pos': np.array([0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785], dtype=np.float32), 'joint_vel': np.linspace(-0.2, 0.2, 7, dtype=np.float32), @@ -45,7 +46,7 @@ def build_payload() -> dict[str, np.ndarray]: def main() -> None: out = Path(__file__).parent / 'droid_obs.npz' - np.savez_compressed(out, **build_payload()) + np.savez_compressed(out, **build_payload()) # pyright: ignore[reportArgumentType] -- numpy's savez **kwds stub print(f'Wrote {out} ({out.stat().st_size} bytes)') diff --git a/pyproject.toml b/pyproject.toml index fcdfd8d7b..90fc32389 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -222,15 +222,19 @@ skip-magic-trailing-comma = true typeCheckingMode = "standard" pythonVersion = "3.11" include = ["pimm", "positronic"] -# vendors/* pull untyped ML libraries; they are typed per-vendor in a later pass. +# vendors/* pull untyped ML libraries; they are typed per-vendor in a later pass. molmo_spaces/env.py runs in +# MolmoSpaces' own venv (it imports molmo_spaces + the PYTHONPATH-relative server/mapping, none resolvable +# against positronic's deps), so it is checked in that interpreter, not here — the same structural reason as +# vendors, not a grandfathered baseline entry. exclude = [ "**/.venv", "**/node_modules", "positronic/vendors", + "positronic/simulator/molmo_spaces/env.py", ] -# `exclude` only drops vendors as analysis roots; a checked module that imports them -# still pulls their diagnostics. `ignore` suppresses those transitive vendor errors too. -ignore = ["positronic/vendors"] +# `exclude` only drops these as analysis roots; a checked module that imports them still pulls their +# diagnostics. `ignore` suppresses those transitive errors too. +ignore = ["positronic/vendors", "positronic/simulator/molmo_spaces/env.py"] [tool.setuptools.packages.find] where = ["."] From bcdddfa16f6b9387204c3d29fed2862f6c02e362 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 23:30:44 +0000 Subject: [PATCH 16/23] Document the grasp-site vs flange control-frame caveat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `env.py` reports `robot_state.ee_pose` at MolmoSpaces' grasp site, but the supplied `bundled_franka_model` declares the physical FR3 flange as its control frame — so recorded episodes mislabel the pose frame and offline IK over them would solve the wrong target, the same known issue as robolab (#483). The eval is unaffected (the pi05 DROID policy is joint-driven, not eef-pose driven). A grasp-site DROID model would fix both. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/cfg/eval/sim/molmo.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index edc273d72..dd85fb295 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -68,6 +68,11 @@ def _molmo_eval( # MolmoSpaces drives a Franka DROID rig; recordings carry the same model (URDF + meshes + joint names + # control frame) for the 3D viewer and offline IK, supplied here since the molmo server can't import # positronic to emit it via ``robot_meta``. + # HACK: ``bundled_franka_model``'s control frame is the physical FR3 flange, but ``env.py`` reports + # ``robot_state.ee_pose`` at MolmoSpaces' gripper grasp site, so recorded episodes mislabel the pose frame + # by the flange↔grasp offset and offline IK over them would solve the wrong target — the same known issue as + # robolab (https://github.com/Positronic-Robotics/positronic/issues/483). It does not affect the eval: the + # pi05 DROID policy is driven by joint commands, not the eef pose. A grasp-site DROID model would fix both. embodiment = remote_franka_embodiment( proxy, camera_dict, descriptor='remote.molmo_spaces.droid', static_meta=bundled_franka_model() ) From f9e6900c0c066e6d394a2798d992b7d7d6faa766 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 23 Jul 2026 23:33:15 +0000 Subject: [PATCH 17/23] Take the reset prompt from the episode spec, not a reconstruction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `task.get_task_description()` reconstructs the instruction per task type upstream (OpeningTask emits "Open the ..." even for a close episode), so it can diverge from the benchmark's authoritative goal. Read the prompt straight from `episode.language.task_description` — the same field JsonEvalTaskSampler itself uses — so recorded episodes and the served prompt always carry the benchmark's own language. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/env.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index 79e5d8242..9d954603c 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -119,7 +119,10 @@ def _build(self, episode_index: int, seed: int | None) -> None: self._task = self._sampler.sample_task(house_index=episode.house_index) self._robot_view = self._task.env.current_robot.robot_view self._control_dt = cfg.policy_dt_ms / 1000.0 - self._meta = {'task': self._task.get_task_description(), 'house_index': episode.house_index} + # The authoritative benchmark prompt, straight from the episode spec — not + # ``task.get_task_description()``, which upstream reconstructs per task type (e.g. OpeningTask emits + # "Open the ..." even for a close episode), so a reconstruction could diverge from the benchmark goal. + self._meta = {'task': episode.language.task_description, 'house_index': episode.house_index} def reset(self, token: dict[str, Any]) -> dict[str, Any]: self._build(token['episode_index'], token.get('seed')) From b9003e1affceeba16cc46f5fce8f018e41da66dc Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 00:16:14 +0000 Subject: [PATCH 18/23] Type-check molmo_spaces/env.py instead of excluding it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit env.py was excluded from basedpyright wholesale — a blanket file grandfather that also silenced the missing-annotation and wrong-type errors that are ours to fix. Check it: type the foreign-venv attributes (`_task`, `_robot_view`, ...) as `Any` so their member access is honest, scope-suppress just `reportMissingImports` file-wide (the only category that can fire in this positronic-free module — its imports are exclusively the molmo venv's), and carry targeted per-line `reportAttributeAccessIssue` ignores for the server symbols, the mujoco C-binding stub gap, and the dynamic CGL module stub. basedpyright is clean with the file included; the ratchet baseline is unchanged. Ticket: Positronic-Robotics/internal#91 #refs --- positronic/simulator/molmo_spaces/env.py | 31 ++++++++++++++++-------- pyproject.toml | 11 ++++----- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py index 9d954603c..b1e60a79e 100644 --- a/positronic/simulator/molmo_spaces/env.py +++ b/positronic/simulator/molmo_spaces/env.py @@ -18,6 +18,12 @@ the raw payload the adapter assembles into a ``MujocoFrankaState``. """ +# molmo_spaces (+ its transitive configs/tasks) resolves only inside MolmoSpaces' own venv, where the launcher +# runs this module; pyright checks it against positronic's deps, which cannot see it. This module imports no +# positronic packages, so missing-import errors here are exclusively those foreign imports — suppress just that +# category file-wide; every other type check (wrong types, optional access, ...) stays active. +# pyright: reportMissingImports=false + import argparse import os import sys @@ -38,9 +44,9 @@ def _install_cgl_noop_stub() -> None: if 'mujoco.cgl' in sys.modules: return cgl = types.ModuleType('mujoco.cgl.cgl') - cgl.CGLLockContext = cgl.CGLUnlockContext = lambda *args, **kwargs: None + cgl.CGLLockContext = cgl.CGLUnlockContext = lambda *args, **kwargs: None # pyright: ignore[reportAttributeAccessIssue] package = types.ModuleType('mujoco.cgl') - package.cgl = cgl + package.cgl = cgl # pyright: ignore[reportAttributeAccessIssue] sys.modules['mujoco.cgl'] = package sys.modules['mujoco.cgl.cgl'] = cgl @@ -53,7 +59,10 @@ def _install_cgl_noop_stub() -> None: import mapping # noqa: E402 -- positronic-free wire mappings, on PYTHONPATH import mujoco # noqa: E402 import numpy as np # noqa: E402 -from server import EnvProtocol, EnvServer # noqa: E402 + +# server resolves to a module without these symbols under positronic's deps (the real one is on the molmo +# venv's PYTHONPATH), so the symbols read as unknown here. +from server import EnvProtocol, EnvServer # noqa: E402 # pyright: ignore[reportAttributeAccessIssue] import molmo_spaces.evaluation.json_eval_runner # noqa: E402, F401 -- load first: breaks a circular import that importing json_eval_task_sampler directly hits from molmo_spaces.configs.policy_configs import DummyPolicyConfig # noqa: E402 @@ -93,11 +102,11 @@ class MolmoSpacesEnv(EnvProtocol): def __init__(self, benchmark_dir: str) -> None: self._episodes = load_all_episodes(Path(benchmark_dir)) - self._sampler = None - self._task = None - self._robot_view = None - self._control_dt = None - self._meta = None + self._sampler: Any = None + self._task: Any = None + self._robot_view: Any = None + self._control_dt: float | None = None + self._meta: dict[str, Any] | None = None # The RGB camera keys the current episode renders — emitted every frame; the client's ``camera_dict`` # selects which the policy sees. self._camera_names: list[str] = [] @@ -154,8 +163,10 @@ def _observe(self, env_obs: dict[str, Any]) -> dict[str, Any]: # from the arm move group's grasp-site frame (obs only exposes a robot-relative tcp pose). arm = self._robot_view.get_move_group('arm') eef_world = np.asarray(arm.leaf_frame_to_world, dtype=np.float64) # 4x4 grasp-site world transform - eef_quat = np.zeros(4) - mujoco.mju_mat2Quat(eef_quat, np.ascontiguousarray(eef_world[:3, :3].reshape(9))) # -> wxyz + eef_quat = np.zeros(4) # filled wxyz below + rot9 = np.ascontiguousarray(eef_world[:3, :3].reshape(9)) + # mju_mat2Quat is a C binding absent from mujoco's type stubs, so pyright can't see the attribute. + mujoco.mju_mat2Quat(eef_quat, rot9) # pyright: ignore[reportAttributeAccessIssue] payload = { 'joint_pos': np.asarray(arm.joint_pos, dtype=np.float32), 'joint_vel': np.asarray(arm.joint_vel, dtype=np.float32), diff --git a/pyproject.toml b/pyproject.toml index 90fc32389..a66f45803 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -222,19 +222,18 @@ skip-magic-trailing-comma = true typeCheckingMode = "standard" pythonVersion = "3.11" include = ["pimm", "positronic"] -# vendors/* pull untyped ML libraries; they are typed per-vendor in a later pass. molmo_spaces/env.py runs in -# MolmoSpaces' own venv (it imports molmo_spaces + the PYTHONPATH-relative server/mapping, none resolvable -# against positronic's deps), so it is checked in that interpreter, not here — the same structural reason as -# vendors, not a grandfathered baseline entry. +# vendors/* pull untyped ML libraries; they are typed per-vendor in a later pass. molmo_spaces/env.py is +# checked here (it runs in MolmoSpaces' own venv; its foreign imports are handled in-file by a scoped +# `reportMissingImports` suppression + targeted per-line ignores), not excluded — a new module owes real +# types, not a blanket exemption. exclude = [ "**/.venv", "**/node_modules", "positronic/vendors", - "positronic/simulator/molmo_spaces/env.py", ] # `exclude` only drops these as analysis roots; a checked module that imports them still pulls their # diagnostics. `ignore` suppresses those transitive errors too. -ignore = ["positronic/vendors", "positronic/simulator/molmo_spaces/env.py"] +ignore = ["positronic/vendors"] [tool.setuptools.packages.find] where = ["."] From 03465569d523764563bc6429e8755526ad57da01 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 07:59:09 +0000 Subject: [PATCH 19/23] Type-clean #485's frame code against the type ratchet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #485 (`droid_eef` + `ChangeEEFrame`) predates this repo's basedpyright ratchet and carries no baseline, so stacking it here surfaces 13 unbaselined errors in its frame code: 11 are mujoco C-binding attribute accesses (`MjData`, `mj_forward`, `mj_name2id`, `mjtObj`, `mju_mat2Quat`) in `frame_transform` and the `_fk_site` test helper — the same foreign-binding class `env.py` already suppresses per line; 2 are loose types in `encode`'s wire signature and the `Derive(**derived)` unpacking. Each gets a targeted per-line `pyright: ignore`. These belong to #485's own review; this commit only keeps the stacked tree ratchet-green and drops out once #485 rebases onto main and resolves them. Ticket: #483 #refs --- .basedpyright/baseline.json | 12 +----------- positronic/drivers/roboarm/ik.py | 8 ++++---- positronic/drivers/roboarm/tests/test_ik.py | 8 ++++---- positronic/policy/codec.py | 4 +++- positronic/simulator/robolab/launcher.py | 4 +++- 5 files changed, 15 insertions(+), 21 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 1d2f48bbc..707c5f208 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -11690,16 +11690,6 @@ } } ], - "./positronic/simulator/robolab/launcher.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 24, - "endColumn": 54, - "lineCount": 1 - } - } - ], "./positronic/simulator/robolab/make_fixture.py": [ { "code": "reportMissingImports", @@ -12785,4 +12775,4 @@ } ] } -} \ No newline at end of file +} diff --git a/positronic/drivers/roboarm/ik.py b/positronic/drivers/roboarm/ik.py index f147c378b..0eda09782 100644 --- a/positronic/drivers/roboarm/ik.py +++ b/positronic/drivers/roboarm/ik.py @@ -67,10 +67,10 @@ def frame_transform(urdf_xml, from_frame, to_frame): spec = _prepare_spec(urdf_xml, from_frame) _ensure_site(spec, to_frame) model = spec.compile() - data = mj.MjData(model) - mj.mj_forward(model, data) - from_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, from_frame)) - to_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, to_frame)) + data = mj.MjData(model) # pyright: ignore[reportAttributeAccessIssue] + mj.mj_forward(model, data) # pyright: ignore[reportAttributeAccessIssue] + from_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, from_frame)) # pyright: ignore[reportAttributeAccessIssue] + to_pose = _site_transform(data, mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, to_frame)) # pyright: ignore[reportAttributeAccessIssue] return from_pose.inv * to_pose diff --git a/positronic/drivers/roboarm/tests/test_ik.py b/positronic/drivers/roboarm/tests/test_ik.py index 489723ce3..3e6ae4c03 100644 --- a/positronic/drivers/roboarm/tests/test_ik.py +++ b/positronic/drivers/roboarm/tests/test_ik.py @@ -119,13 +119,13 @@ def test_ik_joints_from_episode(): def _fk_site(urdf_xml, q, frame): """FK a named frame (site or body) to [tx,ty,tz,w,x,y,z] through the ik spec preparation.""" model = _prepare_spec(urdf_xml, frame).compile() - data = mj.MjData(model) + data = mj.MjData(model) # pyright: ignore[reportAttributeAccessIssue] qpos_ids = [model.joint(n).qposadr.item() for n in JOINT_NAMES] data.qpos[qpos_ids] = q - mj.mj_forward(model, data) - sid = mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, frame) + mj.mj_forward(model, data) # pyright: ignore[reportAttributeAccessIssue] + sid = mj.mj_name2id(model, mj.mjtObj.mjOBJ_SITE, frame) # pyright: ignore[reportAttributeAccessIssue] quat = np.empty(4) - mj.mju_mat2Quat(quat, data.site_xmat[sid]) + mj.mju_mat2Quat(quat, data.site_xmat[sid]) # pyright: ignore[reportAttributeAccessIssue] return np.concatenate([data.site_xpos[sid].copy(), quat]) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 82f5dcc2d..1d6c90558 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -454,7 +454,9 @@ def __init__(self, to: str, pose_keys: tuple[str, ...], derive_pose): def __call__(self, episode): derived = {'control_frame': FromValue(self._to)} derived.update({key: self._derive_pose(key) for key in self._pose_keys if key in episode}) - return Group(Derive(**derived), Identity())(episode) + # ``Derive(**derived)`` unpacks a ``FromValue`` under a keyword pyright reads as ``Derive``'s ``meta`` + # param. Inherited from #485's frame work, which predates this repo's type ratchet — resolved under #485. + return Group(Derive(**derived), Identity())(episode) # pyright: ignore[reportArgumentType] @property def meta(self): diff --git a/positronic/simulator/robolab/launcher.py b/positronic/simulator/robolab/launcher.py index 79473c241..86fe6a72c 100644 --- a/positronic/simulator/robolab/launcher.py +++ b/positronic/simulator/robolab/launcher.py @@ -99,7 +99,9 @@ def _spawn(host: str, port: int) -> subprocess.Popen: robot_meta = {**bundled_franka_model(), 'control_frame': 'droid_eef'} fd, meta_path = tempfile.mkstemp(prefix='robolab_robot_meta_', suffix='.bin') with os.fdopen(fd, 'wb') as meta_file: - meta_file.write(encode(robot_meta)) + # ``encode`` (env_server wire codec) is loosely typed as returning ``Unknown | None``; it returns bytes. + # Inherited from #485's frame work, which predates this repo's type ratchet — resolved under #485's review. + meta_file.write(encode(robot_meta)) # pyright: ignore[reportArgumentType] # Isaac Sim prompts for its EULA on stdin at first launch; the server is headless, so accept it here. env = { **os.environ, From 14fe66a72083f1a1b6e5980a65c80a4edd338dd1 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 08:00:03 +0000 Subject: [PATCH 20/23] Add `molmo_grasp` EE control frame site to the bundled franka model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MolmoSpaces' env reports and accepts Cartesian poses at its arm move group's grasp site (`gripper/grasp_site`, the 2F-85 TCP), not the flange. Add it as a pure-frame site on the flange, measured off MolmoSpaces' own DROID model (`robots/franka_droid/model.xml` @ allenai/molmospaces c2f1b58): 155mm along the flange Z, no rotation — a deeper, unrotated frame than `droid_eef` (18.17mm, +90deg Z). The validation test pins the site to reproduce that measured frame. Ticket: #483 #refs --- positronic/drivers/roboarm/models.py | 9 ++++++++- positronic/drivers/roboarm/tests/test_ik.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/positronic/drivers/roboarm/models.py b/positronic/drivers/roboarm/models.py index 9e6e6cf12..d5eafed30 100644 --- a/positronic/drivers/roboarm/models.py +++ b/positronic/drivers/roboarm/models.py @@ -17,6 +17,12 @@ _DROID_EEF_XYZ = '0 0 0.01817402261' _DROID_EEF_RPY = '0 0 1.5707963268' +# MolmoSpaces' end-effector control frame ``molmo_grasp``, relative to the franka flange (``link8``). Its env +# reports and accepts Cartesian poses at the arm move group's grasp site (``gripper/grasp_site``, the 2F-85 TCP). +# Measured off MolmoSpaces' own DROID model (``robots/franka_droid/model.xml`` @ allenai/molmospaces c2f1b58): +# 155mm along the flange Z, no rotation — a deeper, unrotated frame than ``droid_eef``. +_MOLMO_GRASP_XYZ = '0 0 0.155' + def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: """One 2F-85 finger as URDF rows. ``sign`` mirrors the y-offsets and ``base_rpy`` (180deg Z on the @@ -72,13 +78,14 @@ def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[tuple]: # A row with an axis is a revolute joint whose axis sign sets its closing direction, so one positive # ``grip`` drives the whole 4-bar: driver/spring_link swing the finger in (+X), coupler/follower # counter-rotate (-X) to keep the pad parallel. Rows without an axis are fixed. A row with ``mesh`` None -# is a pure frame (no visual) — ``droid_eef`` is such a frame, the DROID control frame on the flange. +# is a pure frame (no visual) — ``droid_eef`` and ``molmo_grasp`` are such frames, EE control frames on the flange. _ROBOTIQ_2F85 = [ ('gripper_base_mount', _FLANGE_LINK, None, '0 0 0.007', _2F85_MOUNT_RPY, None, 'base_mount.stl', None), ('gripper_base', 'gripper_base_mount', None, '0 0 0.0038', '0 0 -1.5707963268', None, 'base.stl', None), *_2f85_finger('right', 1, '0 0 0'), *_2f85_finger('left', -1, '0 0 3.1415926536'), ('droid_eef', _FLANGE_LINK, None, _DROID_EEF_XYZ, _DROID_EEF_RPY, None, None, None), + ('molmo_grasp', _FLANGE_LINK, None, _MOLMO_GRASP_XYZ, '0 0 0', None, None, None), ] _ROBOTIQ_2F85_JOINTS = [row[2] for row in _ROBOTIQ_2F85 if row[2]] diff --git a/positronic/drivers/roboarm/tests/test_ik.py b/positronic/drivers/roboarm/tests/test_ik.py index 3e6ae4c03..b0dbb73b4 100644 --- a/positronic/drivers/roboarm/tests/test_ik.py +++ b/positronic/drivers/roboarm/tests/test_ik.py @@ -144,6 +144,16 @@ def test_frame_transform_reproduces_droid_eef_across_configs(): assert q_diff < 1e-9, f'rotation mismatch: {q_diff}' +def test_molmo_grasp_matches_molmospaces_grasp_site(): + """``molmo_grasp`` relative to the flange is MolmoSpaces' arm-move-group grasp site (``gripper/grasp_site``), + measured off its own DROID model (``robots/franka_droid/model.xml`` @ allenai/molmospaces c2f1b58): 155mm + along the flange Z, no rotation. It is the frame ``env.py`` reports ``robot_state.ee_pose`` in, so the served + ``robot_meta`` declares it as the control frame — a deeper, unrotated frame than ``droid_eef``.""" + transform = frame_transform(bundled_franka_model()['urdf'], 'link8', 'molmo_grasp') + np.testing.assert_allclose(transform.translation, [0.0, 0.0, 0.155], atol=1e-9) + np.testing.assert_allclose(transform.rotation.as_quat, geom.Rotation.identity.as_quat, atol=1e-9) + + def test_frame_transform_identity_when_frames_match(): transform = frame_transform(bundled_franka_model()['urdf'], 'end_effector', 'end_effector') np.testing.assert_allclose(transform.translation, 0.0, atol=1e-12) From eaead79aa6d081ffa52135fbdb09e8063e46dbcc Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 08:00:48 +0000 Subject: [PATCH 21/23] Declare `molmo_grasp` as the MolmoSpaces eval control frame `env.py` reports `robot_state.ee_pose` at the arm move group's grasp site (`arm.leaf_frame_to_world` = `gripper/grasp_site`), so declare the served `static_meta`'s `control_frame` as `molmo_grasp` instead of the FR3 flange. Recorded episodes and offline IK then live in the frame the env actually uses, not mislabelled by the 155mm flange->grasp offset. Mirrors #485's robolab `droid_eef` flip. Removes the HACK; leaves a TODO(#483) for the per-policy frame-ownership design. Ticket: #483 #refs --- positronic/cfg/eval/sim/molmo.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index dd85fb295..6454f8a81 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -68,13 +68,17 @@ def _molmo_eval( # MolmoSpaces drives a Franka DROID rig; recordings carry the same model (URDF + meshes + joint names + # control frame) for the 3D viewer and offline IK, supplied here since the molmo server can't import # positronic to emit it via ``robot_meta``. - # HACK: ``bundled_franka_model``'s control frame is the physical FR3 flange, but ``env.py`` reports - # ``robot_state.ee_pose`` at MolmoSpaces' gripper grasp site, so recorded episodes mislabel the pose frame - # by the flange↔grasp offset and offline IK over them would solve the wrong target — the same known issue as - # robolab (https://github.com/Positronic-Robotics/positronic/issues/483). It does not affect the eval: the - # pi05 DROID policy is driven by joint commands, not the eef pose. A grasp-site DROID model would fix both. + # ``env.py`` reports ``robot_state.ee_pose`` at MolmoSpaces' arm move group grasp site + # (``arm.leaf_frame_to_world`` = ``gripper/grasp_site``, 155mm along the flange Z), so the model's canonical + # ``control_frame`` is declared as that frame (``molmo_grasp``) — recorded episodes and offline IK over them + # then live in the frame the env actually uses, not the flange. This mirrors the robolab ``droid_eef`` flip. + # TODO(#483): the frame is restated here in the serving config; the per-policy EE-frame design moves that + # ownership onto the checkpoint. See https://github.com/Positronic-Robotics/positronic/issues/483. embodiment = remote_franka_embodiment( - proxy, camera_dict, descriptor='remote.molmo_spaces.droid', static_meta=bundled_franka_model() + proxy, + camera_dict, + descriptor='remote.molmo_spaces.droid', + static_meta={**bundled_franka_model(), 'control_frame': 'molmo_grasp'}, ) task = Task(instruction=lambda: proxy.meta['task'], timeout=timeout, reset=proxy.reset, done=proxy.done) # Benchmark episodes are exact-pose deterministic and carry their own seed. An unset ``seed`` leaves From 3e0e82289013d109cf269d9ad957c994e2051d4d Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 08:04:22 +0000 Subject: [PATCH 22/23] Reference the frame PR in the molmo TODO Cite PR #507 (this stacked PR) in the `molmo_grasp` control-frame TODO now that it is open, alongside the #483 tracking issue. Ticket: #483 #refs --- positronic/cfg/eval/sim/molmo.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py index 6454f8a81..46b525076 100644 --- a/positronic/cfg/eval/sim/molmo.py +++ b/positronic/cfg/eval/sim/molmo.py @@ -73,7 +73,8 @@ def _molmo_eval( # ``control_frame`` is declared as that frame (``molmo_grasp``) — recorded episodes and offline IK over them # then live in the frame the env actually uses, not the flange. This mirrors the robolab ``droid_eef`` flip. # TODO(#483): the frame is restated here in the serving config; the per-policy EE-frame design moves that - # ownership onto the checkpoint. See https://github.com/Positronic-Robotics/positronic/issues/483. + # ownership onto the checkpoint. Declared by https://github.com/Positronic-Robotics/positronic/pull/507; + # tracked in https://github.com/Positronic-Robotics/positronic/issues/483. embodiment = remote_franka_embodiment( proxy, camera_dict, From d198e03317c46c2eceaa3f23d2181a4220a2989b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 24 Jul 2026 08:28:11 +0000 Subject: [PATCH 23/23] Resolve two Codex findings on the frame-transfer path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both surfaced by #507 routing molmo's frame through #485's `ChangeEEFrame`. 1. Read the EE frame from `static_meta` too (harness.py). `_build_obs` read the frame metadata only from the live `robot_meta` channel, but an env may ship its model client-side via `Embodiment.static_meta` while emitting `robot_meta {}` — exactly molmo (`control_frame='molmo_grasp'` in static_meta). So a molmo run with `ChangeEEFrame` got no `urdf`/`control_frame` and KeyError'd on the first obs. Merge static_meta under the live `robot_meta` (which wins), mirroring `_build_episode_meta`'s own merge order. 2. Skip null pose aliases in `ChangeEEFrame.training_encoder` (codec.py). `_RENAME_ROBOT_COMMAND` aliases a joint-only dataset's absent legacy pose to a PRESENT `robot_command.pose` whose value is `None` (`Get(..., None)`), so the `key in episode` guard wrapped `None` in a derived pose signal and joint-only training failed on dereference. Guard on the value being a real signal too. Regression tests fail on the pre-fix code, pass after: a static-meta-frame obs test and a live-robot_meta-wins test (test_harness.py); a null-alias skip test (test_change_ee_frame.py). Ticket: #483 #refs --- positronic/policy/codec.py | 11 ++- positronic/policy/harness.py | 11 +-- .../policy/tests/test_change_ee_frame.py | 32 +++++++++ positronic/policy/tests/test_harness.py | 67 ++++++++++++++++++- 4 files changed, 112 insertions(+), 9 deletions(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 1d6c90558..c5860191a 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -443,8 +443,9 @@ def training_encoder(self) -> EpisodeTransform: class _ChangeEEFrameTraining(EpisodeTransform): """Move the EE pose signals an episode has into ``to`` and relabel ``control_frame`` to match, so a later codec that reads the frame from statics (``IKJointsAction`` solving the command pose) resolves it against the right - site. A pose key the episode lacks — ``robot_command.pose`` under a joint-only action — is skipped rather than - dereferenced, so joint-only training still converts its observation pose.""" + site. A pose key the episode lacks — or carries only as a null rename alias (``robot_command.pose`` under a + joint-only action, aliased to ``None`` by ``_RENAME_ROBOT_COMMAND``'s ``Get(..., None)``) — is skipped rather + than dereferenced, so joint-only training still converts its observation pose.""" def __init__(self, to: str, pose_keys: tuple[str, ...], derive_pose): self._to = to @@ -453,7 +454,11 @@ def __init__(self, to: str, pose_keys: tuple[str, ...], derive_pose): def __call__(self, episode): derived = {'control_frame': FromValue(self._to)} - derived.update({key: self._derive_pose(key) for key in self._pose_keys if key in episode}) + # ``key in episode`` alone is not enough: ``_RENAME_ROBOT_COMMAND`` aliases a joint-only dataset's absent + # legacy pose to a present key whose value is ``None``, so guard on the value being a real signal too. + derived.update({ + key: self._derive_pose(key) for key in self._pose_keys if key in episode and episode[key] is not None + }) # ``Derive(**derived)`` unpacks a ``FromValue`` under a keyword pyright reads as ``Derive``'s ``meta`` # param. Inherited from #485's frame work, which predates this repo's type ratchet — resolved under #485. return Group(Derive(**derived), Identity())(episode) # pyright: ignore[reportArgumentType] diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 80c3f67be..c5110c3e0 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -291,11 +291,14 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['wall_time_ns'] = time.time_ns() inputs['obs_time_ns'] = clock.now_ns() inputs.update(self.context) - model = self.robot_meta_in.value + # The robot model backs any codec that resolves an EE frame against it (``ChangeEEFrame``), mirroring + # how training reads ``urdf``/``control_frame`` from episode statics. It comes from the same source as the + # episode meta (``_build_episode_meta``): ``static_meta`` — an env that ships its model client-side rather + # than over the wire (molmo declares ``control_frame`` in ``static_meta`` while emitting ``robot_meta {}``) + # — merged under the live ``robot_meta`` channel, which wins. It is client-local: ``RemoteSession`` drops + # it before any observation crosses the wire, so it never reaches an inference server. + model = {**self._embodiment.static_meta, **self._static_meta, **self.robot_meta_in.value} if 'control_frame' in model: - # The robot model backs any codec that resolves an EE frame against it (``ChangeEEFrame``), mirroring - # how training reads ``urdf``/``control_frame`` from episode statics. It is client-local — ``RemoteSession`` - # drops it before any observation crosses the wire, so it never reaches an inference server. inputs['urdf'] = model['urdf'] inputs['control_frame'] = model['control_frame'] inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it diff --git a/positronic/policy/tests/test_change_ee_frame.py b/positronic/policy/tests/test_change_ee_frame.py index d1f513bf4..09211b0c4 100644 --- a/positronic/policy/tests/test_change_ee_frame.py +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -3,6 +3,7 @@ import positronic.drivers.roboarm.command as cmd_module from positronic.dataset.episode import EpisodeContainer from positronic.dataset.tests.utils import DummySignal +from positronic.dataset.transforms.episode import Derive, Get, Group, Identity from positronic.drivers.roboarm.ik import frame_transform from positronic.drivers.roboarm.models import bundled_franka_model from positronic.geom import Rotation, Transform3D @@ -116,3 +117,34 @@ def test_training_encoder_skips_absent_command_pose(): transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') np.testing.assert_allclose(out['robot_state.ee_pose'][0][0], (obs_pose * transform).as_vector(QUAT), atol=1e-9) assert out['control_frame'] == 'droid_eef' and 'robot_command.joints' in out + + +def test_training_encoder_skips_null_command_pose_alias(): + """`_RENAME_ROBOT_COMMAND` aliases a joint-only dataset's absent legacy pose to a PRESENT + `robot_command.pose` key whose value is `None` (`Get('robot_commands.pose', None)`). The membership test + must skip that null alias, not wrap `None` in a derived pose signal — else joint-only training fails when the + pose columns are dereferenced. (`key in episode` alone was True for the alias.)""" + obs_pose = _pose([0.3, 0.1, 0.4], [0.2, -0.3, 0.5]) + ts = [1000, 2000] + base = EpisodeContainer( + data={ + 'urdf': URDF, + 'control_frame': CONTROL_FRAME, + 'robot_state.ee_pose': DummySignal(ts, np.stack([obs_pose.as_vector(QUAT)] * 2)), + 'robot_command.joints': DummySignal(ts, np.zeros((2, 7), dtype=np.float32)), + } + ) + # Reproduce `_RENAME_ROBOT_COMMAND` in its production shape (a `Group` with passthrough): alias the (absent) + # legacy `robot_commands.pose` to a present `None`, keeping the episode's other signals. `Derive` takes + # keyword transforms, but pyright reads the unpacked kwarg as its positional `meta` param. + rename = Derive(**{'robot_command.pose': Get('robot_commands.pose', None)}) # pyright: ignore[reportArgumentType] + episode = Group(rename, Identity())(base) + assert 'robot_command.pose' in episode and episode['robot_command.pose'] is None, 'null-alias precondition' + + out = ChangeEEFrame(to='droid_eef').training_encoder(episode) + + # The alias passes through untouched (`None`), not wrapped into a derived pose signal that dereferences it. + assert out['robot_command.pose'] is None + transform = frame_transform(URDF, CONTROL_FRAME, 'droid_eef') + np.testing.assert_allclose(out['robot_state.ee_pose'][0][0], (obs_pose * transform).as_vector(QUAT), atol=1e-9) + assert out['control_frame'] == 'droid_eef' and 'robot_command.joints' in out diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 9bb89de1f..3eacb2d9a 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -30,7 +30,7 @@ from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler -def make_embodiment(descriptor: str = '', cameras=('image.cam',)) -> Embodiment: +def make_embodiment(descriptor: str = '', cameras=('image.cam',), static_meta=None) -> Embodiment: """Minimal Franka-shaped embodiment for harness unit tests. The sources/dests are no-ops: these tests pair the harness ports directly @@ -47,7 +47,7 @@ def make_embodiment(descriptor: str = '', cameras=('image.cam',)) -> Embodiment: 'robot_command': Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command), 'target_grip': Command(pimm.NoOpReceiver(), 0.0, None), } - return Embodiment(descriptor, observations, commands, {}, pimm.NoOpEmitter()) + return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter()) class _SpySession(Session): @@ -328,6 +328,69 @@ def test_harness_passes_descriptor_to_policy(world): assert policy.last_obs['descriptor'] == 'mujoco.franka' +@pytest.mark.timeout(3.0) +def test_static_meta_frame_reaches_obs_when_robot_meta_empty(world): + """A frame declared in ``static_meta`` (molmo: the env emits ``robot_meta {}``) still reaches the obs, so a + ``ChangeEEFrame`` codec can resolve it. Regression: ``_build_obs`` read only the live ``robot_meta`` channel, + so molmo's frame never reached the codec and it would KeyError on the first observation.""" + policy = SpyPolicy() + embodiment = make_embodiment(static_meta={'urdf': '', 'control_frame': 'molmo_grasp'}) + harness = Harness(policy, embodiment) + harness.commands['robot_command']._bind(RecordingEmitter()) + harness.commands['target_grip']._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations['image.cam']) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations['grip']) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (None, 0.05), + ]) + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=20) + + assert policy.last_obs is not None + assert policy.last_obs['control_frame'] == 'molmo_grasp' + assert policy.last_obs['urdf'] == '' + + +@pytest.mark.timeout(3.0) +def test_live_robot_meta_overrides_static_meta_frame(world): + """The live ``robot_meta`` channel wins over ``static_meta`` for the frame, matching the episode-meta merge + order (``_build_episode_meta``) — an env that emits its own model overrides a client-side default.""" + policy = SpyPolicy() + embodiment = make_embodiment(static_meta={'urdf': '', 'control_frame': 'end_effector'}) + harness = Harness(policy, embodiment) + harness.commands['robot_command']._bind(RecordingEmitter()) + harness.commands['target_grip']._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations['image.cam']) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations['grip']) + directive_em = world.pair(harness.directive) + meta_em = world.pair(harness.robot_meta_in) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(meta_em.emit, {'urdf': '', 'control_frame': 'droid_eef'}), 0.005), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (None, 0.05), + ]) + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=20) + + assert policy.last_obs is not None + assert policy.last_obs['control_frame'] == 'droid_eef' + assert policy.last_obs['urdf'] == '' + + @pytest.mark.timeout(3.0) def test_harness_waits_for_complete_inputs(world): pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity)