Declare MolmoSpaces' grasp frame as the eval control frame (molmo_grasp)

.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 @@
             }
         ]
     }
-}
+}

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

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

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:

positronic/cfg/eval/sim/molmo.py

@@ -0,0 +1,103 @@
+import json
+from pathlib import Path
+
+import configuronic as cfn
+
+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 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(
+    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: 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
+    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)
+    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
+    # positronic to emit it via ``robot_meta``.
+    # ``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. 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,
+        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
+    # ``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).
+benchmark = _molmo_eval
+
+# A single-episode smoke target: the first episode of the benchmark.
+first_episode = _molmo_eval.override(episodes=0)

positronic/drivers/roboarm/ik.py

@@ -8,40 +8,70 @@
 """
 
 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'):
             for elem in link.findall('visual') + link.findall('collision'):
                 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)  # 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
 
 
 def _parse_target(target_ee_pose_vec):

positronic/drivers/roboarm/models.py

@@ -11,6 +11,18 @@
 # (a +45deg Z, i.e. 90deg off the franka ``end_effector`` frame).
 _2F85_MOUNT_RPY = '0 0 0.7853981634'
 
+# 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'
+
+# 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
@@ -62,15 +74,18 @@ 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`` 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]]
 
@@ -85,10 +100,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)