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/.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/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/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py new file mode 100644 index 000000000..46b525076 --- /dev/null +++ b/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) diff --git a/positronic/drivers/roboarm/ik.py b/positronic/drivers/roboarm/ik.py index 425b72487..0eda09782 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) # 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): diff --git a/positronic/drivers/roboarm/models.py b/positronic/drivers/roboarm/models.py index 7f087ab27..d5eafed30 100644 --- a/positronic/drivers/roboarm/models.py +++ b/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) diff --git a/positronic/drivers/roboarm/tests/test_ik.py b/positronic/drivers/roboarm/tests/test_ik.py index 23ebc6b2f..b0dbb73b4 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,50 @@ 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) # 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) # 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]) # pyright: ignore[reportAttributeAccessIssue] + 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_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) + 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/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/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 9df303653..fcee2171a 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 6cb994505..c5860191a 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.dataset.transforms.episode import Derive, EpisodeTransform, Group, Identity +from positronic import geom +from positronic.dataset.transforms import Elementwise, lazy_sequence +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 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.""" @@ -373,3 +385,84 @@ 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. + + 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 + """ + + 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: + 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 — 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 + self._pose_keys = pose_keys + self._derive_pose = derive_pose + + def __call__(self, episode): + derived = {'control_frame': FromValue(self._to)} + # ``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] + + @property + def meta(self): + return {'ee_frame': self._to} diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index a7d3d5d27..c5110c3e0 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -291,6 +291,16 @@ 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) + # 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: + 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/observation.py b/positronic/policy/observation.py index 9d424b7a0..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 @@ -18,18 +18,27 @@ 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()}) - 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(): @@ -47,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 {} @@ -54,7 +67,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/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..09211b0c4 --- /dev/null +++ b/positronic/policy/tests/test_change_ee_frame.py @@ -0,0 +1,150 @@ +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.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 +from positronic.policy.codec import ChangeEEFrame + +QUAT = Rotation.Representation.QUAT +# 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' + + +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_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(): + 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) + # ``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 + + +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) 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..96595c85d --- /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]) -> None: + 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/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 new file mode 100644 index 000000000..b1e60a79e --- /dev/null +++ b/positronic/simulator/molmo_spaces/env.py @@ -0,0 +1,204 @@ +"""MolmoSpaces — AllenAI's MuJoCo manipulation benchmark — behind the env-server protocol. + +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``/ +``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``. +""" + +# 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 +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 # pyright: ignore[reportAttributeAccessIssue] + package = types.ModuleType('mujoco.cgl') + package.cgl = cgl # pyright: ignore[reportAttributeAccessIssue] + 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 + +import mapping # noqa: E402 -- positronic-free wire mappings, on PYTHONPATH +import mujoco # noqa: E402 +import numpy as np # 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 +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) -> None: + self._episodes = load_all_episodes(Path(benchmark_dir)) + 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] = [] + + 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) + # 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 + self._control_dt = cfg.policy_dt_ms / 1000.0 + # 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')) + 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}) + # 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: + 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) # 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), + '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..961c94984 --- /dev/null +++ b/positronic/simulator/molmo_spaces/launcher.py @@ -0,0 +1,91 @@ +"""Launches the MolmoSpaces env server as a subprocess and owns its lifetime. + +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 +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: + 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. 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', 'pip', 'install', '-e', f'.[{_MOLMO_EXTRA}]'], + cwd=str(src), + env={**os.environ, 'VIRTUAL_ENV': str(venv)}, + check=True, + ) + command = [ + str(venv / 'bin' / 'python'), + 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 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) + + +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..52a3a15ef --- /dev/null +++ b/positronic/simulator/molmo_spaces/mapping.py @@ -0,0 +1,91 @@ +"""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); 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', '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). +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 000000000..217586f43 Binary files /dev/null and b/positronic/simulator/molmo_spaces/tests/droid_obs.npz differ 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..052287792 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/make_fixture.py @@ -0,0 +1,54 @@ +# /// 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 +from typing import Any + +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, 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), + '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()) # pyright: ignore[reportArgumentType] -- numpy's savez **kwds stub + 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..260f02b93 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_mapping.py @@ -0,0 +1,95 @@ +"""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, ()) + + +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 diff --git a/positronic/simulator/robolab/launcher.py b/positronic/simulator/robolab/launcher.py index a887d31ef..86fe6a72c 100644 --- a/positronic/simulator/robolab/launcher.py +++ b/positronic/simulator/robolab/launcher.py @@ -93,13 +93,15 @@ 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())) + # ``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, 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..a66f45803 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", @@ -221,14 +222,17 @@ 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 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", ] -# `exclude` only drops vendors as analysis roots; a checked module that imports them -# still pulls their diagnostics. `ignore` suppresses those transitive vendor errors too. +# `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"] [tool.setuptools.packages.find]