diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 8e1741cd1..62fd60d24 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5628,16 +5628,6 @@ } } ], - "./positronic/simulator/env_server/proxy.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 28, - "endColumn": 32, - "lineCount": 1 - } - } - ], "./positronic/simulator/env_server/tests/mujoco_env.py": [ { "code": "reportArgumentType", @@ -5687,14 +5677,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 60, - "endColumn": 64, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -5705,14 +5687,6 @@ } ], "./positronic/simulator/env_server/tests/test_remote_env.py": [ - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 22, - "endColumn": 32, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 621a40dab..af211f2c8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -165,3 +165,25 @@ leaks a timing vocabulary into the dataset core. Hence telemetry is a set of sid process, next to the dataset but never inside it: nested spans for the phase split and a free-running machine-load sampler, wall-clock native, owned by `positronic/telemetry.py`. The pass report is an offline reduce over those raw files, so nothing is stored twice and the dataset stays clock-agnostic. + +**An adoption loses nothing.** A task is defined in real or in one simulator, and carries that +simulator with it — object poses, success criteria and horizons included. What a customer buys is +one API across all of them: they implement a single Positronic policy interface and their model +runs on every supported env, giving up nothing the env offers natively. Two requirements hold that +up, one on each side of the interface: + +- Given a policy that already drives an env directly, it must be possible to construct a Positronic + `Policy` equivalent to it. This binds policy construction as much as it binds the adoption. +- An adoption's capabilities match what its env provides natively, so a Positronic run reproduces + the env's own run: a deterministic env to byte-identical outcomes (modulo wire format), a + non-deterministic env to an identical sim/inference call sequence (same count, same order). + +Every sim-env adoption ships a native-vs-Positronic parity test that drives one pinned episode +through both stacks and asserts this, re-run on every bump of the sim's pinned version. + +The episode horizon is one case of that reproduction rather than a rule of its own: a task that +defines a horizon has it enforced by the env, which reports expiry through the same terminal `done` +a success uses; a task that defines none leaves nothing to reproduce. The harness `Task.timeout` is +only a runaway-cost safety net, so the config that knows the benchmark derives the timeout from the +horizon it declares rather than taking one on faith — a budget below the horizon would silently +truncate valid episodes and score them as failures. diff --git a/CLAUDE.md b/CLAUDE.md index 508d4934f..f06ff815f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -115,4 +115,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/docs/evaluation.md b/docs/evaluation.md index 080096c4b..cee843bfe 100644 --- a/docs/evaluation.md +++ b/docs/evaluation.md @@ -9,6 +9,7 @@ You ship a new checkpoint and want a clean answer to one question: is it actuall ## What you get - **One checkpoint, every target.** Sim: LIBERO, RoboLab (NVIDIA Isaac Lab), MolmoSpaces. Real hardware: the DROID setup (Franka FR3 + Robotiq 2F-85), bimanual next. Serve a DROID policy once and it runs across all of them, and on the rig, with nothing to port. Sim for cheap, broad iteration; real hardware as ground truth. +- **Native benchmarks, comparable scores.** Each benchmark runs with its own task definitions and horizons — the run reproduces the native benchmark rather than a re-interpretation of it — so your score is comparable with the benchmark's own published numbers. - **Blinded A/B.** Your checkpoint against your own previous checkpoints, or against our maintained baselines (π0.5, GR00T, SmolVLA, ACT) — randomized and blinded, so lighting and setup drift don't bias the result. - **Every run returned.** Multi-view video, full telemetry, and the complete run dataset — not just a success rate. Yours to analyze. - **Latency-honest execution.** On real hardware, inference and network delay are real — a slow model is scored as slow. In sim the world pauses during inference by default (as in other harnesses), but you can charge the model's measured inference time with `--charge_inference_time=True`, so sim scores reflect the delay the robot would actually feel — something sim-only harnesses can't model. diff --git a/positronic/cfg/eval/sim/molmo.py b/positronic/cfg/eval/sim/molmo.py new file mode 100644 index 000000000..c1540a107 --- /dev/null +++ b/positronic/cfg/eval/sim/molmo.py @@ -0,0 +1,99 @@ +import logging +from pathlib import Path + +import configuronic as cfn + +from positronic.cfg.eval import number_trials, spec +from positronic.drivers.roboarm.models import GRASP_SITE_LINK, bundled_franka_model +from positronic.eval import Eval, Observation, Task +from positronic.eval import keys as eval_keys +from positronic.simulator.env_server.proxy import RemoteEnvControlSystem, remote_franka_embodiment +from positronic.simulator.molmo_spaces import keys as molmo_keys +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.molmo_spaces.adapter import DEFAULT_CAMERA_DICT, MolmoAdapter +from positronic.simulator.molmo_spaces.launcher import serve_molmo_spaces + +# How far the harness deadline sits above the benchmark horizon. Being sim-time, the spare budget costs +# nothing unless the sim stops terminating, which is the only thing the deadline is there to catch. +_TIMEOUT_MARGIN_SEC = 10.0 + + +@cfn.config(camera_dict=DEFAULT_CAMERA_DICT, episodes=None, trial_count=1, timeout=None, seed=None) +def _molmo_eval( + benchmark_dir: str, + episodes: int | list[int] | None, + trial_count: int, + timeout: float | None, + 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 directory holding a ``benchmark.json`` + (a JSON list of episode specs — house, task, exact object poses, cameras, language goal), so + ``--eval.benchmark_dir`` names that directory 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, the env answers which episodes the sweep runs, and the episode index rides each trial's reset + token. 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. + + ``timeout`` is not the benchmark horizon — the sim owns that (the benchmark's ``task_horizon_sec``, enforced + env-side and delivered as a terminal ``done``). It is only a runaway-cost safety net for a sim that never + terminates, so its default is the benchmark's own horizon plus a margin. An explicit value can only lower the + deadline, never raise it, and one at or below the horizon truncates valid episodes — so any value that + differs from the default is warned about. + """ + # A non-positive count yields no trials at all, and an empty plan reads to the self-driving harness as a + # finished run — the command would exit 0 having evaluated nothing. + if trial_count < 1: + raise ValueError(f'--eval.trial_count must be at least 1, got {trial_count}') + proxy = RemoteEnvControlSystem(MolmoAdapter(camera_dict), serve_molmo_spaces(Path(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``. ``DEFAULT_FRAME`` is declared on the gripper's grasp site, + # which is where ``env.py`` reports ``robot_state.ee_pose`` and resolves Cartesian targets, so a policy + # frame reached from it via ``ChangeEEFrame`` and offline IK over a recording both anchor correctly. + embodiment = remote_franka_embodiment( + proxy, camera_dict, descriptor='remote.molmo_spaces.droid', static_meta=bundled_franka_model(GRASP_SITE_LINK) + ) + # The env's full MuJoCo state is recorded as privileged ground truth, never fed to the policy. + privileged = {mapping.OBS_SIM_STATE: Observation(proxy.privileged[mapping.OBS_SIM_STATE], None)} + + def tasks() -> list[Task]: + params = proxy.tasks(spec(episodes=episodes)) + # The benchmark declares one horizon over all its episodes (the env refuses an inconsistent one), so one + # backstop deadline covers the run. + backstop = params[0][molmo_keys.TASK_HORIZON] + _TIMEOUT_MARGIN_SEC + if timeout is not None and timeout != backstop: + logging.warning( + '--eval.timeout %ss overrides the benchmark backstop of %ss (the %ss horizon plus a margin); ' + 'running with %ss. The deadline only catches a sim that stopped terminating, and a deadline at ' + 'or below the horizon cuts valid episodes short and scores them as failures.', + timeout, + backstop, + params[0][molmo_keys.TASK_HORIZON], + min(timeout, backstop), + ) + deadline = backstop if timeout is None else min(timeout, backstop) + task = Task(instruction_source=lambda: proxy.meta[mapping.META_TASK], timeout_sec=deadline) + # 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``. + return number_trials([ + (task, {**p, **({eval_keys.SEED: seed + t} if seed is not None else {})}) + for p in params + for t in range(trial_count) + ]) + + return Eval(embodiment, tasks, privileged=privileged, done=proxy.done) + + +# 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/models.py b/positronic/drivers/roboarm/models.py index e51d665ff..3fc7d1788 100644 --- a/positronic/drivers/roboarm/models.py +++ b/positronic/drivers/roboarm/models.py @@ -3,7 +3,7 @@ real arm and the live franka driver.""" import xml.etree.ElementTree as ET -from functools import lru_cache +from functools import cache, lru_cache from pathlib import Path from typing import NamedTuple @@ -99,10 +99,17 @@ def _2f85_finger(side: str, sign: int, base_rpy: str) -> list[_UrdfRow]: ] +GRASP_SITE_LINK = 'gripper_grasp_site' +# The 2F-85's grasp point — where the closed pads meet — 155mm along the flange approach axis, sharing the +# flange's orientation. MuJoCo Menagerie's own 2F-85 places its ``grasp_site`` here, and a MuJoCo rig driving +# this gripper measures and accepts poses at that site rather than at the arm's flange. +_2F85_GRASP_XYZ = '0 0 0.155' + # The coupler stays fixed: given an axis, the outer link hangs 19 mm out at full grip. _ROBOTIQ_2F85 = [ _UrdfRow('gripper_base_mount', FLANGE_LINK, None, '0 0 0.007', _2F85_MOUNT_RPY, None, 'base_mount.stl', None), _UrdfRow('gripper_base', 'gripper_base_mount', None, '0 0 0.0038', '0 0 -1.5707963268', None, 'base.stl', None), + _UrdfRow(GRASP_SITE_LINK, FLANGE_LINK, None, _2F85_GRASP_XYZ, '0 0 0', None, None, None), *_2f85_finger('right', 1, '0 0 0'), *_2f85_finger('left', -1, '0 0 3.1415926536'), # RoboLab's ``eef_frame`` (``Robotiq_2F_85/base_link`` ∘ ``EEF_OFFSET_ROT``), measured off its DROID USD @@ -174,19 +181,21 @@ def attach_robotiq_2f85(arm_root: ET.Element, meshes: dict[str, bytes]) -> dict: return gripper[roboarm_keys.GRIPPER] -@lru_cache(maxsize=1) -def bundled_franka_model() -> dict: +@cache +def bundled_franka_model(default_frame_at: str = EE_LINK) -> dict: """The bundled real franka arm + Robotiq 2F-85 for the 3D viewer: the FR3 URDF and its collision meshes with the 2F-85 grafted onto the flange, plus the canonical joint names and control frame. - Backfills real-robot datasets recorded before they stored their own model. + Backfills real-robot datasets recorded before they stored their own model. ``default_frame_at`` names the + link ``DEFAULT_FRAME`` is declared on: a rig that measures and drives at the gripper's grasp point passes + ``GRASP_SITE_LINK``, so it publishes poses in the frame it drives rather than at the franka EE. """ here = Path(__file__).resolve() arm_root = ET.fromstring((here.parent / 'fr3.urdf').read_text()) mesh_dir = here.parents[2] / 'assets' / 'fr3_collision' meshes = {f.name: f.read_bytes() for f in sorted(mesh_dir.glob('*.stl'))} gripper = attach_robotiq_2f85(arm_root, meshes) - add_default_frame(arm_root, EE_LINK) + add_default_frame(arm_root, default_frame_at) return { roboarm_keys.URDF: ET.tostring(arm_root, encoding='unicode'), 'meshes': meshes, diff --git a/positronic/drivers/roboarm/tests/test_ik.py b/positronic/drivers/roboarm/tests/test_ik.py index 6ec1eeab5..3e0c04c74 100644 --- a/positronic/drivers/roboarm/tests/test_ik.py +++ b/positronic/drivers/roboarm/tests/test_ik.py @@ -23,6 +23,8 @@ DROID_EE_FRAME, DROID_EEF_LINK, EE_LINK, + FLANGE_LINK, + GRASP_SITE_LINK, bundled_franka_model, bundled_panda_model, ) @@ -245,3 +247,20 @@ def test_pickle_roundtrip(solver_cls): q_result = restored.solve(q_start, target_pose) result_pose = _fk(PANDA_URDF, q_result) np.testing.assert_allclose(result_pose[:3], target_pose[:3], atol=1e-3) + + +def test_grasp_site_sits_at_the_2f85_grasp_point(): + """The 2F-85's grasp point is 155mm along the flange approach axis, in the flange's own orientation — + where MolmoSpaces' franka_droid model places its ``gripper/grasp_site``.""" + transform = frame_transform(bundled_franka_model()[roboarm_keys.URDF], FLANGE_LINK, GRASP_SITE_LINK) + np.testing.assert_allclose(transform.translation, [0.0, 0.0, 0.155], atol=1e-9) + np.testing.assert_allclose(transform.rotation.as_rotation_matrix, np.eye(3), atol=1e-9) + + +def test_grasp_site_model_declares_the_frame_it_reports_in(): + """A rig measuring at the grasp point declares ``DEFAULT_FRAME`` there, so the frame it publishes poses + in is the frame it drives.""" + model = bundled_franka_model(GRASP_SITE_LINK) + assert model[roboarm_keys.CONTROL_FRAME] == DEFAULT_FRAME + transform = frame_transform(model[roboarm_keys.URDF], DEFAULT_FRAME, GRASP_SITE_LINK) + np.testing.assert_allclose(transform.as_matrix, np.eye(4), atol=1e-9) diff --git a/positronic/eval/__init__.py b/positronic/eval/__init__.py index 4e8f2872a..2ec83c276 100644 --- a/positronic/eval/__init__.py +++ b/positronic/eval/__init__.py @@ -67,7 +67,9 @@ class Task: """One trial: the goal the policy conditions on, the time budget it runs under, and what sets it up.""" instruction_source: str | Callable[[], str] - # Time budget for a rollout; ``None`` ends on ``Eval.done`` alone. + # Time budget for a rollout; ``None`` ends on ``Eval.done`` alone. A benchmark sim enforces the task's own + # horizon and reports expiry as a terminal, so there the budget is a runaway-cost net set well beyond any + # healthy horizon; a real or attended eval has no such terminal and the budget is the trial's actual bound. timeout_sec: float | None # What to ask for, keyed as ``Embodiment.prepare_handlers`` is; a handler this does not name goes unasked prepare_args: dict[str, Any] = field(default_factory=dict) diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ddf7140e1..f1c481951 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -2,6 +2,35 @@ 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** — the client being the process driving the robot or sim, + the server being the process holding the model. positronic-native evals compose the codec around + `RemotePolicy` on the client (`cfg/policy.py` — the wire then carries model-sized encoded inputs, + and the client-side resize is disabled since `codec.meta` already reports `image_sizes`). + Thin-client deployments (a sim adapter in a foreign venv talking to a serverless endpoint) host + the codec on the server — the wire carries raw positronic keys, downsized by the negotiation + above. Both placements are supported; pick by where the dependencies can live. + ## Protocol v1 The unified WebSocket protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index ff6aa9b12..16a37a046 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from functools import partial from typing import Any @@ -8,7 +9,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 # The encoded observation's language prompt, under the name LeRobot training and its policies both use. It @@ -22,7 +23,9 @@ class ObservationCodec(Codec): Args: 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. + task_field: output key carrying the language prompt at inference; training always emits ``TASK_FIELD``. + 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). """ WIRE_NAME = 'observation_codec' @@ -32,14 +35,20 @@ def __init__( state: dict[str, dict[str, int]], images: dict[str, tuple[str, tuple[int, int]]], task_field: str = TASK_FIELD, + lowercase_task: bool = False, ): self._state = state self._image_configs = images self._task_field = task_field + self._lowercase_task = lowercase_task - self._derive_transforms: dict[str, Any] = {k: partial(self._derive_state, k) for k in state.keys()} + self._derive_transforms: dict[str, Callable[[Episode], Any]] = { + 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_FIELD] = Get(keys.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_FIELD] = self._derive_task lerobot_features: dict[str, Any] = {} for name, features in state.items(): @@ -57,6 +66,12 @@ 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 _normalize_task(self, task: str) -> str: + return task.lower() if self._lowercase_task else task + + def _derive_task(self, episode: Episode) -> Any: + return self._normalize_task(episode[keys.TASK] if keys.TASK in episode else '') + def _decode_single(self, data: dict) -> dict: return {} @@ -64,7 +79,7 @@ def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: obs: dict[str, Any] = {} if keys.TASK in inputs: - obs[self._task_field] = inputs[keys.TASK] + obs[self._task_field] = self._normalize_task(inputs[keys.TASK]) for out_name, (input_key, (width, height)) in self._image_configs.items(): if input_key not in inputs: @@ -101,5 +116,10 @@ def to_spec(self): images = {name: [key, list(size)] for name, (key, size) in self._image_configs.items()} return { 'name': self.WIRE_NAME, - 'args': {'state': self._state, 'images': images, 'task_field': self._task_field}, + 'args': { + 'state': self._state, + 'images': images, + 'task_field': self._task_field, + 'lowercase_task': self._lowercase_task, + }, } diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index c6f9f1569..e3f50a8d0 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -19,6 +19,7 @@ FlipGrip, ) from positronic.policy.observation import ObservationCodec +from positronic.policy.spec import from_spec def test_observation_encode_images_and_state_shapes(): @@ -67,6 +68,14 @@ def test_observation_encode_task(): assert obs_keys.TASK not in obs_no_task +def test_observation_codec_spec_preserves_lowercase_task(): + enc = ObservationCodec(state={}, images={}, lowercase_task=True) + rebuilt = from_spec(enc.to_spec()) + assert isinstance(rebuilt, ObservationCodec) + obs = rebuilt.encode({obs_keys.TASK: 'MixedCase'}) + assert obs[obs_keys.TASK] == 'mixedcase' + + def test_absolute_position_action_encode_decode_quat(): # Identity rotation, known translation/grip ts = [1000, 2000] diff --git a/positronic/simulator/env_server/adapter.py b/positronic/simulator/env_server/adapter.py index 08be1cde8..d5c49c97f 100644 --- a/positronic/simulator/env_server/adapter.py +++ b/positronic/simulator/env_server/adapter.py @@ -15,6 +15,7 @@ import pimm from positronic import geom, keys from positronic.drivers.roboarm import command as roboarm_command +from positronic.simulator.env_server import protocol class EnvAdapter(ABC): @@ -83,23 +84,23 @@ def _wire_command(cmd: Any) -> dict[str, Any]: rep = geom.Rotation.Representation.ROTATION_MATRIX match cmd: case roboarm_command.CartesianPosition(pose): - wire = {'type': 'cartesian', 'pose': pose.as_vector(rep)} + wire = {protocol.COMMAND_TYPE: protocol.CARTESIAN, protocol.COMMAND_POSE: pose.as_vector(rep)} case roboarm_command.JointPosition(positions): - wire = {'type': 'joint_pos', 'q': positions} + wire = {protocol.COMMAND_TYPE: protocol.JOINT_POS, protocol.COMMAND_JOINT_POS: positions} case roboarm_command.JointDelta(velocities): - wire = {'type': 'joint_vel', 'dq': velocities} + wire = {protocol.COMMAND_TYPE: protocol.JOINT_VEL, protocol.COMMAND_JOINT_VEL: velocities} case roboarm_command.CartesianDelta(delta, frame): # The env anchors a delta on the pose it measures, which is its control frame and nowhere else, so # a delta still expressed somewhere else has no faithful wire form. if not np.allclose(frame.as_matrix, np.eye(4)): raise ValueError('CartesianDelta outside the env control frame cannot be sent to a remote env') - wire = {'type': 'cartesian_delta', 'delta': delta.as_vector(rep)} + wire = {protocol.COMMAND_TYPE: protocol.CARTESIAN_DELTA, protocol.COMMAND_DELTA: delta.as_vector(rep)} case None: - return {'type': 'hold'} + return {protocol.COMMAND_TYPE: protocol.HOLD} case other: raise ValueError(f'no wire encoding for robot_command {type(other).__name__}') if cmd.mode is not None: - wire['mode'] = roboarm_command.to_wire(cmd.mode) + wire[protocol.COMMAND_MODE] = roboarm_command.to_wire(cmd.mode) return wire @@ -144,4 +145,7 @@ def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: if isinstance(cmd, roboarm_command.CartesianDelta | roboarm_command.JointDelta): self._held.pop(keys.ROBOT_COMMAND) grip = float(self._held.get(keys.TARGET_GRIP, 0.0)) - return {'command': _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), 'grip': grip} + return { + protocol.ACTION_COMMAND: _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), + protocol.ACTION_GRIP: grip, + } diff --git a/positronic/simulator/env_server/launcher.py b/positronic/simulator/env_server/launcher.py index e2708d7ff..0993ddf7f 100644 --- a/positronic/simulator/env_server/launcher.py +++ b/positronic/simulator/env_server/launcher.py @@ -8,6 +8,7 @@ import shutil import socket import subprocess +import time from collections.abc import Callable, Iterator from contextlib import contextmanager from pathlib import Path @@ -57,22 +58,50 @@ def terminate(proc: subprocess.Popen) -> None: proc.kill() +# How long a server may take to bind its port, and how often the wait re-checks. The deadline must cover a +# cold first boot, where a simulator compiles shaders and loads asset packs before serving. +_BIND_DEADLINE = 1800.0 +_BIND_POLL_INTERVAL = 0.2 + + +def _await_bind(proc: subprocess.Popen, host: str, port: int, deadline: float) -> None: + """Block until *proc* accepts connections on *port*, raising if it exits or the deadline passes. + + The server's accept loop drops a connection that never handshakes, so this probe costs it nothing. + """ + end = time.monotonic() + deadline + while True: + try: + with socket.create_connection((host, port), timeout=1.0): + return + except OSError: + pass + status = proc.poll() + if status is not None: + raise RuntimeError(f'env server exited with status {status} before binding {host}:{port}') + if time.monotonic() >= end: + raise TimeoutError(f'env server did not bind {host}:{port} within {deadline:.0f}s') + time.sleep(_BIND_POLL_INTERVAL) + + @contextmanager -def serve_subprocess(spawn: Callable[[str, int], subprocess.Popen], host: str) -> Iterator[tuple[str, int]]: - """Run an env-server subprocess for the body's lifetime, yielding its ``(host, port)``. +def serve_subprocess( + spawn: Callable[[str, int], subprocess.Popen], host: str, bind_deadline: float = _BIND_DEADLINE +) -> Iterator[tuple[str, int]]: + """Run an env-server subprocess for the body's lifetime, yielding its ``(host, port)`` once it is bound. The single owner of the subprocess: ``RemoteEnvControlSystem`` enters it to tie the subprocess to the World run, and a plain client (e.g. an e2e demo replay) enters it directly to talk over the socket without a World. The task spec rides the reset token, so the subprocess needs only its address — it serves - whatever task the first reset asks for. The port is picked before the spawn; the client's connect retry - covers the gap until the server binds it. + whatever task the first reset asks for. - TODO: a subprocess that dies at startup goes unnoticed until the client's connect deadline — nothing - surfaces its exit during the retry wait. + ``bind_deadline`` must cover a cold first boot: a heavy simulator spends minutes compiling shaders and + loading assets before it binds. """ port = free_port() proc = spawn(host, port) try: + _await_bind(proc, host, port, bind_deadline) yield host, port finally: terminate(proc) diff --git a/positronic/simulator/env_server/protocol.py b/positronic/simulator/env_server/protocol.py index 53ca48079..5b7e3cfbb 100644 --- a/positronic/simulator/env_server/protocol.py +++ b/positronic/simulator/env_server/protocol.py @@ -28,6 +28,50 @@ class Command(Enum): CLOSE = 'close' +# The canonical command contract: the tag on every arm command a client puts on the wire. It is total — one +# contract carries every policy onto every embodiment — so an env adoption converts each of these into +# whatever its own controller natively takes. Owned here because both interpreters spell the tags: +# positronic's ``EnvAdapter`` writes them, an env venv's own decoder reads them, and this is the module both +# sides import. +CARTESIAN = 'cartesian' +CARTESIAN_DELTA = 'cartesian_delta' +JOINT_POS = 'joint_pos' +JOINT_VEL = 'joint_vel' +HOLD = 'hold' +CANONICAL_COMMAND_TYPES = (CARTESIAN, CARTESIAN_DELTA, JOINT_POS, JOINT_VEL, HOLD) + +# The action a client puts on the wire: the tagged arm command, and the gripper closure alongside it. +ACTION_COMMAND = 'command' +ACTION_GRIP = 'grip' + +# The tagged command's own fields: the tag, the one value each tag carries (``hold`` carries none), and the +# control law the command pins. +COMMAND_TYPE = 'type' +COMMAND_POSE = 'pose' # CARTESIAN — an absolute pose, [t(3), R(9)] +COMMAND_DELTA = 'delta' # CARTESIAN_DELTA — a relative pose, same encoding +COMMAND_JOINT_POS = 'q' # JOINT_POS — absolute joint targets +COMMAND_JOINT_VEL = 'dq' # JOINT_VEL — per-step joint deltas +COMMAND_MODE = 'mode' # any tag — the pinned control mode, absent when the command pins none + +# The address every env-server script is spawned with: its launcher builds the command in positronic's +# interpreter, its ``env.py`` parser declares it in the adoption's own, so a rename that misses one side +# fails at spawn rather than at import. +OPT_HOST = '--host' +OPT_PORT = '--port' + +# The frames an env reports back. ``reset`` carries the observation, the scene meta, the robot model identity +# and the control period; ``step`` carries the observation, the terminal, the control period, and — where the +# env judges one — its success. ``horizon`` is the episode limit the env enforces itself, in sim-seconds, +# absent when the env enforces none. +FRAME_OBS = 'obs' +FRAME_META = 'meta' +FRAME_ROBOT_META = 'robot_meta' +FRAME_CONTROL_DT = 'control_dt' +FRAME_HORIZON = 'horizon' +FRAME_DONE = 'done' +FRAME_SUCCESS = 'success' + + def _pack(obj): if isinstance(obj, np.ndarray): if obj.dtype.kind in ('V', 'O', 'c'): diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index 440d9915e..962c8ff11 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -18,6 +18,7 @@ from positronic.dataset.serializers import Serializers from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol from positronic.simulator.env_server.adapter import EnvAdapter from positronic.simulator.env_server.client import EnvConnection @@ -94,10 +95,10 @@ def reset(self, params: dict[str, Any]) -> None: for _, receiver in self.commands.items(): receiver.read() self._frame = conn.reset(self._adapter.reset_token(params)) - self._meta = self._frame['meta'] + self._meta = self._frame[protocol.FRAME_META] self._active = True - self.robot_meta.emit(self._frame['robot_meta']) - self._emit_payload(self._frame['obs']) + self.robot_meta.emit(self._frame[protocol.FRAME_ROBOT_META]) + self._emit_payload(self._frame[protocol.FRAME_OBS]) # An empty payload clears the wire: a terminal the previous trial reached would end this one at once. self.done.emit({}) @@ -111,7 +112,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p try: while not should_stop.value: # The proxy paces every turn; ``control_dt`` is known only once a reset ran. - yield pimm.Sleep(self._frame['control_dt'] if self._frame is not None else _IDLE_DT) + yield pimm.Sleep(self._frame[protocol.FRAME_CONTROL_DT] if self._frame is not None else _IDLE_DT) if (call := next(self.env_reset.incoming(), None)) is not None: with pimm.calls.raise_to(call): self.reset(dict(call.request or {})) @@ -121,11 +122,13 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p with telemetry.span(telemetry_keys.SPAN_ENV_STEP): self._frame = self._step_env() with telemetry.span(telemetry_keys.SPAN_MATERIALIZE): - self._emit_payload(self._frame['obs']) + self._emit_payload(self._frame[protocol.FRAME_OBS]) finally: self._cleanup.close() def _step_env(self) -> dict[str, Any]: + # Stepping is reachable only while ``_active``, which ``reset`` sets once the connection is up. + assert self._conn is not None, 'stepped before the first reset connected' reads = ((name, receiver.read()) for name, receiver in self.commands.items()) commands = {name: msg for name, msg in reads if msg is not None} result = self._conn.step(self._adapter.action(commands)) diff --git a/positronic/simulator/env_server/server.py b/positronic/simulator/env_server/server.py index 102070e27..4974658bd 100644 --- a/positronic/simulator/env_server/server.py +++ b/positronic/simulator/env_server/server.py @@ -4,7 +4,7 @@ Protocol (msgpack frames, see ``protocol``): client ``{'cmd': 'tasks', 'spec': ...}`` -> server ``{'tasks': [{...}, ...]}`` - client ``{'cmd': 'reset', 'token': ...}`` -> server ``{'obs', 'meta', 'robot_meta', 'control_dt'}`` + client ``{'cmd': 'reset', 'token': ...}`` -> server ``{'obs', 'meta', 'robot_meta', 'control_dt', 'horizon'?}`` client ``{'cmd': 'step', 'action': {...}}`` -> server ``{'obs', 'done', 'control_dt'}`` client ``{'cmd': 'close'}`` -> server ``{'ok': True}`` Command handling failures return ``{'error': str}`` without closing the session; the client re-raises them. @@ -46,6 +46,10 @@ def reset(self, token: Any) -> dict[str, Any]: Return ``obs``, scene ``meta``, ``robot_meta``, and ``control_dt`` in seconds. Either metadata dict may be empty when the client supplies it. + + ``horizon`` (optional) is the sim-enforced episode deadline in sim-seconds — the env's own time limit, + which it delivers as a terminal ``done`` on expiry. It is reported for observability, so a run can be + checked against the horizon the env actually resolved; omit it when the env enforces none. """ @abstractmethod diff --git a/positronic/simulator/env_server/tests/mujoco_env.py b/positronic/simulator/env_server/tests/mujoco_env.py index 18423f51d..a4adf59ec 100644 --- a/positronic/simulator/env_server/tests/mujoco_env.py +++ b/positronic/simulator/env_server/tests/mujoco_env.py @@ -22,6 +22,7 @@ from positronic.drivers.roboarm import command as roboarm_command from positronic.eval import Eval, Observation, Task from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol from positronic.simulator.env_server.adapter import WireCommandAdapter from positronic.simulator.env_server.proxy import RemoteEnvControlSystem, remote_franka_embodiment from positronic.simulator.env_server.server import EnvProtocol @@ -118,36 +119,46 @@ def reset(self, token: Any) -> dict[str, Any]: self._sim.reset(token) self._gen = self._sim.run(_NeverStop(), self._clock) next(self._gen) # loop setup + first control-period sleep + robot_meta = self._robot_meta_recv.read() + assert robot_meta is not None, 'the sim publishes robot_meta as it resets, before the sleep stepped above' # Native ``stack_cubes`` has no language scene meta (its instruction is a static client string), so ``meta`` # is empty; the sim's robot identity (URDF / joints) is the ``robot_meta``. return { - 'obs': self._read_obs(), - 'meta': {}, - 'robot_meta': dict(self._robot_meta_recv.read().data), - 'control_dt': self._timestep, + protocol.FRAME_OBS: self._read_obs(), + protocol.FRAME_META: {}, + protocol.FRAME_ROBOT_META: dict(robot_meta.data), + protocol.FRAME_CONTROL_DT: self._timestep, } def step(self, action: dict[str, Any]) -> dict[str, Any]: assert self._gen is not None, 'step() called before reset()' # real Gym envs reject step-before-reset - command = action['command'] - match command['type']: - case 'hold': + command = action[protocol.ACTION_COMMAND] + match command[protocol.COMMAND_TYPE]: + case protocol.HOLD: cmd = None - case 'joint_pos': - cmd = roboarm_command.JointPosition(np.asarray(command['q'], dtype=np.float64)) - case 'joint_vel': - cmd = roboarm_command.JointDelta(np.asarray(command['dq'], dtype=np.float64)) - case 'cartesian': - cmd = roboarm_command.CartesianPosition(geom.Transform3D.from_vector(command['pose'], _ROTMAT)) - case 'cartesian_delta': - cmd = roboarm_command.CartesianDelta(geom.Transform3D.from_vector(command['delta'], _ROTMAT)) + case protocol.JOINT_POS: + q = np.asarray(command[protocol.COMMAND_JOINT_POS], dtype=np.float64) + cmd = roboarm_command.JointPosition(q) + case protocol.JOINT_VEL: + dq = np.asarray(command[protocol.COMMAND_JOINT_VEL], dtype=np.float64) + cmd = roboarm_command.JointDelta(dq) + case protocol.CARTESIAN: + pose = geom.Transform3D.from_vector(command[protocol.COMMAND_POSE], _ROTMAT) + cmd = roboarm_command.CartesianPosition(pose) + case protocol.CARTESIAN_DELTA: + delta = geom.Transform3D.from_vector(command[protocol.COMMAND_DELTA], _ROTMAT) + cmd = roboarm_command.CartesianDelta(delta) case other: raise ValueError(f'MujocoEnv got unsupported command type {other!r}') if cmd is not None: self._cmd_emit.emit(cmd) - self._grip_emit.emit(float(action['grip'])) + self._grip_emit.emit(float(action[protocol.ACTION_GRIP])) self._advance(self._timestep) - return {'obs': self._read_obs(), 'done': False, 'control_dt': self._timestep} + return { + protocol.FRAME_OBS: self._read_obs(), + protocol.FRAME_DONE: False, + protocol.FRAME_CONTROL_DT: self._timestep, + } def close(self) -> None: if self._gen is not None: diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index bd29820c8..40625a5a2 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -1,3 +1,6 @@ +import socket +import subprocess +import sys import threading import time from contextlib import contextmanager, nullcontext @@ -31,7 +34,7 @@ from positronic.simulator.env_server import protocol from positronic.simulator.env_server.adapter import EnvAdapter, _in_env_control_frame, _wire_command from positronic.simulator.env_server.client import _CLOSE_ACK_TIMEOUT, EnvConnection -from positronic.simulator.env_server.launcher import free_port +from positronic.simulator.env_server.launcher import free_port, serve_subprocess from positronic.simulator.env_server.proxy import RemoteEnvControlSystem from positronic.simulator.env_server.server import EnvProtocol from positronic.simulator.env_server.tests.conftest import serve_env @@ -84,6 +87,26 @@ def _assert_obs_equal(a: dict, b: dict) -> None: np.testing.assert_array_equal(a['sim_state'][key], b['sim_state'][key]) +@pytest.mark.timeout(30.0) +def test_serve_subprocess_reports_a_server_that_dies_before_binding(): + """A server that raises during startup never binds, so its port stays closed exactly as a slow boot's + does.""" + spawn = lambda host, port: subprocess.Popen([sys.executable, '-c', 'raise SystemExit(3)']) # noqa: E731 + with pytest.raises(RuntimeError, match='status 3'), serve_subprocess(spawn, 'localhost'): + pass + + +@pytest.mark.timeout(30.0) +def test_serve_subprocess_yields_once_the_port_accepts(): + script = ( + 'import socket,sys,time\ns=socket.socket()\ns.bind(("localhost",int(sys.argv[1])))\ns.listen()\ntime.sleep(30)' + ) + spawn = lambda host, port: subprocess.Popen([sys.executable, '-c', script, str(port)]) # noqa: E731 + with serve_subprocess(spawn, 'localhost') as (host, port): + with socket.create_connection((host, port), timeout=1.0): + pass + + @pytest.mark.timeout(60.0) def test_transport_is_transparent(env_server): """The same seed and actions must yield identical raw observations in-process and over the socket.""" @@ -92,8 +115,17 @@ def test_transport_is_transparent(env_server): direct = make_mujoco_env(list(CAMERAS.values())) direct_reset = direct.reset(seed) - base = np.asarray(direct_reset['obs']['q']) - actions = [{'command': {'type': 'joint_pos', 'q': base + 0.03 * i}, 'grip': 0.2 * (i % 2)} for i in range(1, 6)] + base = np.asarray(direct_reset[protocol.FRAME_OBS]['q']) + actions = [ + { + protocol.ACTION_COMMAND: { + protocol.COMMAND_TYPE: protocol.JOINT_POS, + protocol.COMMAND_JOINT_POS: base + 0.03 * i, + }, + protocol.ACTION_GRIP: 0.2 * (i % 2), + } + for i in range(1, 6) + ] direct_steps = [direct.step(action) for action in actions] direct.close() @@ -102,12 +134,12 @@ def test_transport_is_transparent(env_server): socket_steps = [conn.step(action) for action in actions] conn.close() - assert direct_reset['control_dt'] == socket_reset['control_dt'] - _assert_obs_equal(direct_reset['obs'], socket_reset['obs']) + assert direct_reset[protocol.FRAME_CONTROL_DT] == socket_reset[protocol.FRAME_CONTROL_DT] + _assert_obs_equal(direct_reset[protocol.FRAME_OBS], socket_reset[protocol.FRAME_OBS]) for direct_step, socket_step in zip(direct_steps, socket_steps, strict=True): - _assert_obs_equal(direct_step['obs'], socket_step['obs']) - assert direct_step['done'] == socket_step['done'] - assert direct_step['control_dt'] == socket_step['control_dt'] + _assert_obs_equal(direct_step[protocol.FRAME_OBS], socket_step[protocol.FRAME_OBS]) + assert direct_step[protocol.FRAME_DONE] == socket_step[protocol.FRAME_DONE] + assert direct_step[protocol.FRAME_CONTROL_DT] == socket_step[protocol.FRAME_CONTROL_DT] @pytest.mark.timeout(60.0) @@ -177,7 +209,7 @@ def ignore_ping(frame): return if release.wait(timeout=0.2): return - connection.send(protocol.encode({'obs': {'ready': True}})) + connection.send(protocol.encode({protocol.FRAME_OBS: {'ready': True}})) host, port = 'localhost', free_port() with websocket_serve(handler, host, port) as server: @@ -196,7 +228,7 @@ def test_scene_reset_survives_delayed_heartbeat_replies(server_without_heartbeat host, port, ignored_pings = server_without_heartbeat conn = EnvConnection(host, port) try: - assert conn.reset({}) == {'obs': {'ready': True}} + assert conn.reset({}) == {protocol.FRAME_OBS: {'ready': True}} assert ignored_pings finally: conn.close() @@ -217,24 +249,23 @@ def test_unanswered_heartbeat_closes_pending_requests(server_without_heartbeat, conn.close() -_HOLD = {'command': {'type': 'hold'}, 'grip': 0.0} +_HOLD = {protocol.ACTION_COMMAND: {protocol.COMMAND_TYPE: protocol.HOLD}, protocol.ACTION_GRIP: 0.0} def _settle(env, action: dict, steps: int) -> np.ndarray: - """Apply the action once, hold for ``steps`` ticks, and return the settled end-effector position.""" - env.step(action) - out = {'obs': None} + """Apply ``action`` once, then idle ``steps`` ticks while the position actuators settle; return the final eef.""" + out = env.step(action) for _ in range(steps): out = env.step(_HOLD) - return np.asarray(out['obs']['ee_pos']) + return np.asarray(out[protocol.FRAME_OBS]['ee_pos']) def test_a_pinned_control_mode_rides_the_wire(): """Control modes must pass through for the environment to interpret.""" mode = roboarm_command.Impedance(kq=(40.0,) * 7, kqd=(4.0,) * 7, kx=(750.0,) * 6, kxd=(37.0,) * 6) wired = _wire_command(roboarm_command.JointPosition(np.zeros(7), mode=mode)) - assert wired['mode'] == roboarm_command.to_wire(mode) - assert 'mode' not in _wire_command(roboarm_command.JointPosition(np.zeros(7))) + assert wired[protocol.COMMAND_MODE] == roboarm_command.to_wire(mode) + assert protocol.COMMAND_MODE not in _wire_command(roboarm_command.JointPosition(np.zeros(7))) class TestEnvControlFrame: @@ -246,13 +277,13 @@ class TestEnvControlFrame: def test_an_absolute_pose_arrives_in_the_env_frame(self): pose = geom.Transform3D(np.array([0.4, 0.1, 0.3]), geom.Rotation.from_euler([0.1, 0.2, 0.3])) wired = _wire_command(_in_env_control_frame(roboarm_command.CartesianPosition(pose), self.frame)) - np.testing.assert_allclose(wired['pose'], (pose * self.frame).as_vector(self.rotmat), atol=1e-12) + np.testing.assert_allclose(wired[protocol.COMMAND_POSE], (pose * self.frame).as_vector(self.rotmat), atol=1e-12) def test_a_delta_already_in_the_env_frame_wires_bare(self): delta = geom.Transform3D(np.array([0.0, 0.0, 0.04]), geom.Rotation.identity) cmd = roboarm_command.CartesianDelta(delta, frame=self.frame) wired = _wire_command(_in_env_control_frame(cmd, self.frame)) - np.testing.assert_allclose(wired['delta'], delta.as_vector(self.rotmat), atol=1e-12) + np.testing.assert_allclose(wired[protocol.COMMAND_DELTA], delta.as_vector(self.rotmat), atol=1e-12) def test_a_command_re_expressed_for_the_env_keeps_its_mode(self): mode = roboarm_command.Impedance(kq=(40.0,) * 7, kqd=(4.0,) * 7, kx=(750.0,) * 6, kxd=(37.0,) * 6) @@ -295,15 +326,28 @@ def test_cartesian_delta_matches_absolute_target(): abs_env = make_mujoco_env(list(CAMERAS.values())) reset = abs_env.reset(seed) - ee0 = np.asarray(reset['obs']['ee_pos']) - target = geom.Transform3D(ee0 + lift, geom.Rotation.from_quat(reset['obs']['ee_quat'])) - ee_abs = _settle(abs_env, {'command': {'type': 'cartesian', 'pose': target.as_vector(rotmat)}, 'grip': 0.0}, settle) + ee0 = np.asarray(reset[protocol.FRAME_OBS]['ee_pos']) + target = geom.Transform3D(ee0 + lift, geom.Rotation.from_quat(reset[protocol.FRAME_OBS]['ee_quat'])) + absolute = { + protocol.ACTION_COMMAND: { + protocol.COMMAND_TYPE: protocol.CARTESIAN, + protocol.COMMAND_POSE: target.as_vector(rotmat), + }, + protocol.ACTION_GRIP: 0.0, + } + ee_abs = _settle(abs_env, absolute, settle) abs_env.close() delta_env = make_mujoco_env(list(CAMERAS.values())) delta_env.reset(seed) delta = geom.Transform3D(lift, geom.Rotation.identity) - delta_action = {'command': {'type': 'cartesian_delta', 'delta': delta.as_vector(rotmat)}, 'grip': 0.0} + delta_action = { + protocol.ACTION_COMMAND: { + protocol.COMMAND_TYPE: protocol.CARTESIAN_DELTA, + protocol.COMMAND_DELTA: delta.as_vector(rotmat), + }, + protocol.ACTION_GRIP: 0.0, + } ee_delta = _settle(delta_env, delta_action, settle) ee_idle = _settle(delta_env, _HOLD, 50) # the delta already fired; idling must not re-compose it delta_env.close() @@ -335,16 +379,20 @@ def reset(self, token): self._steps = 0 meta = {'task': _COUNTDOWN} return { - 'obs': {'q': np.full(7, self._steps, dtype=np.float64)}, - 'meta': meta, - 'robot_meta': {}, - 'control_dt': self._control_dt, + protocol.FRAME_OBS: {'q': np.full(7, self._steps, dtype=np.float64)}, + protocol.FRAME_META: meta, + protocol.FRAME_ROBOT_META: {}, + protocol.FRAME_CONTROL_DT: self._control_dt, } def step(self, action): self._steps += 1 done = self._done_after is not None and self._steps >= self._done_after - return {'obs': {'q': np.full(7, self._steps, dtype=np.float64)}, 'done': done, 'control_dt': self._control_dt} + return { + protocol.FRAME_OBS: {'q': np.full(7, self._steps, dtype=np.float64)}, + protocol.FRAME_DONE: done, + protocol.FRAME_CONTROL_DT: self._control_dt, + } def close(self): pass @@ -367,7 +415,7 @@ def privileged(self, raw_obs): return {} def terminal(self, result): - return {eval_keys.SUCCESS: True} if result['done'] else None + return {eval_keys.SUCCESS: True} if result[protocol.FRAME_DONE] else None @pytest.mark.timeout(60.0) @@ -446,6 +494,21 @@ def test_proxy_caches_reset_meta_as_live_instruction_source(): assert task.instruction == 'countdown' +def test_the_canonical_contract_is_exactly_what_a_client_can_emit(): + """``protocol`` owns the contract and ``_wire_command`` is the only thing that writes it, so pinning the two + against each other keeps the set an env adoption must cover equal to the set a policy can actually emit.""" + pose = geom.Transform3D(np.zeros(3), geom.Rotation.identity) + commands = [ + roboarm_command.CartesianPosition(pose), + roboarm_command.CartesianDelta(pose), + roboarm_command.JointPosition(np.zeros(7)), + roboarm_command.JointDelta(np.zeros(7)), + None, # nothing held: the arm holds where it is + ] + tags = {_wire_command(command)[protocol.COMMAND_TYPE] for command in commands} + assert tags == set(protocol.CANONICAL_COMMAND_TYPES) + + @pytest.mark.timeout(60.0) def test_remote_eval_runs_to_timeout_without_done(env_server, tmp_path): """A timed-out trial must record canonical signals without reporting termination or success.""" @@ -516,7 +579,7 @@ def test_full_chunk_executes_between_replans(env_server, tmp_path): """ host, port = env_server probe = make_mujoco_env([]) - control_dt = probe.reset(0)['control_dt'] + control_dt = probe.reset(0)[protocol.FRAME_CONTROL_DT] probe.close() chunk_len = 5 @@ -547,15 +610,21 @@ def test_full_chunk_executes_between_replans(env_server, tmp_path): 'message', [ {protocol.CMD: 'bogus'}, - {protocol.CMD: protocol.Command.STEP.value, protocol.ACTION: {'command': {'type': 'bogus'}, 'grip': 0.0}}, + { + protocol.CMD: protocol.Command.STEP.value, + protocol.ACTION: {protocol.ACTION_COMMAND: {protocol.COMMAND_TYPE: 'bogus'}, protocol.ACTION_GRIP: 0.0}, + }, ], ) def test_server_failure_crosses_as_error_frame(env_server, message): - """Rejected commands must reach the client as errors while leaving the connection usable.""" + """A command the env rejects comes back as an error the client re-raises — the connection survives + rather than dying on the server-side exception, and the next command still works.""" host, port = env_server conn = EnvConnection(host, port) conn.reset(7) with pytest.raises(RuntimeError, match='bogus'): conn._request(message) - assert 'obs' in conn.step({'command': {'type': 'joint_pos', 'q': np.zeros(7)}, 'grip': 0.0}) + # The socket is still usable after a delivered failure. + joints = {protocol.COMMAND_TYPE: protocol.JOINT_POS, protocol.COMMAND_JOINT_POS: np.zeros(7)} + assert protocol.FRAME_OBS in conn.step({protocol.ACTION_COMMAND: joints, protocol.ACTION_GRIP: 0.0}) conn.close() diff --git a/positronic/simulator/libero/adapter.py b/positronic/simulator/libero/adapter.py index e4c9ddd72..d1534a33d 100644 --- a/positronic/simulator/libero/adapter.py +++ b/positronic/simulator/libero/adapter.py @@ -13,6 +13,7 @@ import pimm from positronic import geom, keys from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol from positronic.simulator.env_server.adapter import WireCommandAdapter from positronic.simulator.libero import keys as libero_keys from positronic.simulator.mujoco.sim import MujocoFrankaState @@ -65,4 +66,4 @@ def privileged(self, raw_obs: dict[str, Any]) -> dict[str, Any]: def terminal(self, result: dict[str, Any]) -> dict[str, Any] | None: # ``done`` is LIBERO's success check rather than a step limit, so reaching it is the success. - return {eval_keys.SUCCESS: True} if result['done'] else None + return {eval_keys.SUCCESS: True} if result[protocol.FRAME_DONE] else None 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..ccd467cb4 --- /dev/null +++ b/positronic/simulator/molmo_spaces/adapter.py @@ -0,0 +1,77 @@ +"""``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, keys +from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol +from positronic.simulator.env_server.adapter import WireCommandAdapter +from positronic.simulator.molmo_spaces import keys as molmo_keys +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, +} + +# Which MolmoSpaces camera each logical observation reads on the DROID rig — the pairing the benchmarks record, +# and the one whose variants the table above resolves. +DEFAULT_CAMERA_DICT = {keys.WRIST_IMAGE: mapping.MOLMO_WRIST_CAMERA, keys.EXTERIOR_IMAGE: mapping.MOLMO_EXTERIOR_CAMERA} + + +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 task_params(self, records: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + eval_keys.TASK: record['name'], + molmo_keys.EPISODE_INDEX: record['episode_index'], + molmo_keys.TASK_HORIZON: record['task_horizon_sec'], + } + for record in records + ] + + def _reset_token(self, params: dict[str, Any]) -> Any: + # The benchmark episode selector rides the token. An absent seed leaves the episode spec's own in force. + return { + mapping.TOKEN_EPISODE_INDEX: params[molmo_keys.EPISODE_INDEX], + mapping.TOKEN_SEED: params.get(eval_keys.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[mapping.OBS_EEF_POS], geom.Rotation.from_quat(raw_obs[mapping.OBS_EEF_QUAT])) + state = MujocoFrankaState() + state.encode(raw_obs[mapping.OBS_JOINT_POS], raw_obs[mapping.OBS_JOINT_VEL], ee_pose) + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs[mapping.OBS_GRIP])} + for logical, molmo_key in self._camera_dict.items(): + env_key = mapping.resolve_camera_key(raw_obs, 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]: + # The env's full MuJoCo state — recorded as ground truth so success can be recomputed offline, never fed + # to the policy (mirrors the libero adapter). + return {mapping.OBS_SIM_STATE: raw_obs[mapping.OBS_SIM_STATE]} + + 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_keys.SUCCESS: bool(result[protocol.FRAME_SUCCESS])} if result[protocol.FRAME_DONE] else None diff --git a/positronic/simulator/molmo_spaces/env.py b/positronic/simulator/molmo_spaces/env.py new file mode 100644 index 000000000..1522f113c --- /dev/null +++ b/positronic/simulator/molmo_spaces/env.py @@ -0,0 +1,449 @@ +"""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) and the flat ``protocol`` resolve only inside MolmoSpaces' +# own venv, where the launcher runs this module; pyright checks it against positronic's deps, which cannot see +# them. Each of those imports carries its own ``reportMissingImports`` suppression, so an import that should +# resolve here — anything from positronic, which this module must never take — still fails the check. + +import argparse +import os +import sys +import types + +import mapping # positronic-free wire mappings, on PYTHONPATH; numpy only, so it pulls in no GL + +# MolmoSpaces renders MuJoCo scenes, so the GL backend must be selected before any mujoco/molmo_spaces import. +# The launcher sets it in the subprocess env; default it here too so a direct invocation (e.g. a validate/e2e +# run) still boots. Set before the imports below. +os.environ.setdefault(mapping.GL_BACKEND_ENV, mapping.GL_BACKEND_DEFAULT) + + +# MolmoSpaces' renderer module, which the stub below stands in for on Linux. +_CGL_PACKAGE = 'mujoco.cgl' +_CGL_MODULE = f'{_CGL_PACKAGE}.cgl' + + +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. + # macOS keeps the real module, where those locks guard an actual context. + if sys.platform == 'darwin' or _CGL_PACKAGE in sys.modules: + return + cgl = types.ModuleType(_CGL_MODULE) + cgl.CGLLockContext = cgl.CGLUnlockContext = lambda *args, **kwargs: None # pyright: ignore[reportAttributeAccessIssue] + package = types.ModuleType(_CGL_PACKAGE) + package.cgl = cgl # pyright: ignore[reportAttributeAccessIssue] + sys.modules[_CGL_PACKAGE] = package + sys.modules[_CGL_MODULE] = cgl + + +_install_cgl_noop_stub() + +from pathlib import Path # noqa: E402 +from typing import Any # noqa: E402 + +import mujoco # noqa: E402 +import numpy as np # noqa: E402 +import protocol # noqa: E402 -- the positronic-free wire contract, on PYTHONPATH # pyright: ignore[reportMissingImports] + +# 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] + +# Imported for its import-time ``_assert_data_versions_match()``: MolmoSpaces pins the asset versions a +# benchmark may be evaluated against, and this is the only place upstream enforces it. Driving the sampler +# directly skips the native entrypoint, so without this a run on mismatched asset packs would score where +# MolmoSpaces itself refuses to. +import molmo_spaces.evaluation.eval_main # noqa: E402, F401 # pyright: ignore[reportMissingImports] +import molmo_spaces.evaluation.json_eval_runner # noqa: E402, F401 -- load first: breaks a circular import that importing json_eval_task_sampler directly hits # pyright: ignore[reportMissingImports] +from molmo_spaces.configs.policy_configs import DummyPolicyConfig # noqa: E402 # pyright: ignore[reportMissingImports] +from molmo_spaces.configs.robot_configs import ( # noqa: E402 # pyright: ignore[reportMissingImports] + ActionNoiseConfig, + FrankaRobotConfig, +) +from molmo_spaces.evaluation.benchmark_schema import ( # noqa: E402 # pyright: ignore[reportMissingImports] + load_all_episodes, +) +from molmo_spaces.evaluation.configs.evaluation_configs import ( # noqa: E402 # pyright: ignore[reportMissingImports] + JsonBenchmarkEvalConfig, +) +from molmo_spaces.tasks.json_eval_task_sampler import ( # noqa: E402 # pyright: ignore[reportMissingImports] + JsonEvalTaskSampler, +) + +# Damped-least-squares differential IK, matching the LIBERO rig's solver (positronic/simulator/libero/env.py): +# the same iteration budget, damping and convergence tolerance, on MuJoCo's own site/body Jacobian. +_IK_ITERS = 100 +_IK_DAMPING = 0.05 +_IK_TOL = 1e-4 + + +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) + + +# MolmoSpaces reports a move group's leaf frame as one of MuJoCo's frame kinds; the arm's is a site. +_SITE_FRAME = 'site' + + +def _discovery_hint() -> str: + """The benchmark dirs found under ``MLSPACES_ASSETS_DIR``, appended to a path that holds none.""" + assets = os.environ.get(mapping.ASSETS_DIR_ENV) + if not assets: + return f' Point {mapping.ASSETS_DIR_ENV} at the MolmoSpaces asset packs to have the available ones listed.' + root = Path(assets) / mapping.ASSETS_BENCHMARKS_DIR + found = sorted(str(p.parent) for p in root.rglob(mapping.MOLMO_BENCHMARK_MANIFEST)) + if not found: + return f' No {mapping.MOLMO_BENCHMARK_MANIFEST} found under {root}.' + return f' Available under {root}: {", ".join(found)}' + + +def _assert_measures_at_grasp_site(robot_view: Any) -> None: + """Fail unless the arm move group's leaf frame is ``mapping.MOLMO_GRASP_SITE``. + + That frame is what every pose this server reports is measured in, and the eval declares its recorded + model's control frame at the same point. Nothing else ties the two together, so a scene whose arm resolves + somewhere else would misframe every recorded pose silently — for the viewer, for offline IK and for any + frame a policy asks for. + """ + arm = robot_view.get_move_group(mapping.MOLMO_ARM_GROUP) + if arm.leaf_frame_type != _SITE_FRAME: + raise ValueError(f'arm move group measures at a {arm.leaf_frame_type}, not the expected site') + name = mujoco.mj_id2name(arm.mj_model, mujoco.mjtObj.mjOBJ_SITE, arm.leaf_frame_id) # pyright: ignore[reportAttributeAccessIssue] + if name != mapping.MOLMO_GRASP_SITE and not name.endswith(f'/{mapping.MOLMO_GRASP_SITE}'): + raise ValueError(f'arm move group measures at site {name!r}, expected {mapping.MOLMO_GRASP_SITE!r}') + + +class MolmoSpacesEnv(EnvProtocol): + """A MolmoSpaces benchmark behind the ``tasks``/``reset``/``step``/``close`` the env server serves. + + ``tasks`` answers the benchmark's episode records. Each reset builds 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: Path, task_horizon_steps: int | None = None) -> None: + self._benchmark_dir = benchmark_dir + self._episodes = load_all_episodes(benchmark_dir) + # An explicit per-run horizon override (steps), mirroring MolmoSpaces' ``--task_horizon_steps``; ``None`` + # reads the benchmark's own ``task_horizon_sec``. + self._task_horizon_override = task_horizon_steps + self._sampler: Any = None + self._task: Any = None + self._robot_view: Any = None + self._control_dt: float | None = None + # The episode's enforced horizon in sim-seconds (``task_horizon`` steps x the control period), reported + # at reset. + self._horizon_sec: float | None = None + self._meta: dict[str, Any] | None = None + # The RGB camera keys the current episode renders, emitted every frame. + self._camera_names: list[str] = [] + # Scratch ``MjData`` the kinematics probes (``_fk``/``_ik``) run on, allocated once per episode and + # refreshed from the live buffer per call — a Cartesian policy solves IK every control step, so the + # allocation stays out of the loop. Rebuilt in ``_build``, since it is sized by the episode's model. + self._scratch: Any = None + + 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). + cfg.seed = mapping.resolve_episode_seed(episode, episode_index, seed) + # With ``task_horizon`` set, the task enforces it and ``is_done`` reports expiry, so a horizon-expired + # trial ends with a terminal ``done`` exactly as the native benchmark scores it. The horizon is the + # benchmark's, not an episode's. + cfg.task_horizon = mapping.resolve_task_horizon_steps( + self._episodes, cfg.policy_dt_ms, self._task_horizon_override + ) + 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 + _assert_measures_at_grasp_site(self._robot_view) + self._scratch = None # sized by this episode's model; allocated on the first probe + self._control_dt = cfg.policy_dt_ms / 1000.0 + self._horizon_sec = cfg.task_horizon * self._control_dt + # 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 = { + mapping.META_TASK: episode.language.task_description, + mapping.META_HOUSE_INDEX: episode.house_index, + } + + def tasks(self, spec: dict[str, Any]) -> list[dict[str, Any]]: + """The episode records ``spec`` selects: ``episodes`` is one index, or a list of them; absent, the whole + benchmark. Every record carries the one horizon the benchmark enforces, converted against the config + that enforces it.""" + if not self._episodes: + raise ValueError( + f'no benchmark episodes under {self._benchmark_dir}; expected a ' + f'{mapping.MOLMO_BENCHMARK_MANIFEST} or a legacy house_*/episode_*.json layout.{_discovery_hint()}' + ) + count = len(self._episodes) + selection = spec.get('episodes') + indices = list(range(count)) if selection is None else [selection] if isinstance(selection, int) else selection + # A negative index would silently run a from-the-end episode mislabeled by its own index. + out_of_range = [i for i in indices if not 0 <= i < count] + if out_of_range: + raise ValueError( + f'episodes {out_of_range} out of range for the {count} episodes under {self._benchmark_dir}' + ) + cfg = _DroidPickEvalConfig() + steps = mapping.resolve_task_horizon_steps(self._episodes, cfg.policy_dt_ms, self._task_horizon_override) + horizon_sec = steps * cfg.policy_dt_ms / 1000.0 + return [ + {'name': self._episodes[i].language.task_description, 'episode_index': i, 'task_horizon_sec': horizon_sec} + for i in indices + ] + + def reset(self, token: dict[str, Any]) -> dict[str, Any]: + self._build(token[mapping.TOKEN_EPISODE_INDEX], token.get(mapping.TOKEN_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 mapping.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``). + return { + protocol.FRAME_OBS: self._observe(env_obs), + protocol.FRAME_META: self._meta, + protocol.FRAME_ROBOT_META: {}, + protocol.FRAME_CONTROL_DT: self._control_dt, + protocol.FRAME_HORIZON: self._horizon_sec, + } + + def step(self, action: dict[str, Any]) -> dict[str, Any]: + arm = mapping.wire_command_to_arm_action( + action[protocol.ACTION_COMMAND], self._measured_arm_q(), ik=self._ik, current_eef=self._measured_eef_pose() + ) + gripper = np.array([mapping.grip_command_to_actuator(action[protocol.ACTION_GRIP])], dtype=np.float32) + obs, _reward, _term, _trunc, _infos = self._task.step({ + mapping.MOLMO_ARM_GROUP: arm, + mapping.MOLMO_GRIPPER_GROUP: gripper, + }) + # The trial ends on the task's judged success, on a MolmoSpaces terminal, or on horizon expiry — the + # latter two through ``is_done``. ``success`` is ORed in for end-on-success, the benchmark's scoring + # semantics: without it a successful rollout that kept sending joint commands would idle to the horizon. + # It stays ``judge_success()`` alone, so a horizon expiry ends the trial with ``success=False``, as + # native scoring has it. + success = bool(self._task.judge_success()) + done = success or bool(self._task.is_done()) + return { + protocol.FRAME_OBS: self._observe(obs[0]), + protocol.FRAME_DONE: done, + protocol.FRAME_SUCCESS: success, + protocol.FRAME_CONTROL_DT: self._control_dt, + } + + def _measured_arm_q(self) -> np.ndarray: + return np.asarray(self._robot_view.get_move_group(mapping.MOLMO_ARM_GROUP).joint_pos, dtype=np.float32) + + def _measured_eef_pose(self) -> tuple[np.ndarray, np.ndarray]: + """The measured grasp-site world pose as ``(translation, 3x3 rotation)`` — the frame a Cartesian + command targets and the one ``_observe`` reports, so command and observation share a frame.""" + eef_world = np.asarray( + self._robot_view.get_move_group(mapping.MOLMO_ARM_GROUP).leaf_frame_to_world, dtype=np.float64 + ) + return eef_world[:3, 3].copy(), eef_world[:3, :3].copy() + + def _scratch_data(self, move_group: Any) -> Any: + """The scratch ``MjData``, refreshed from the live one, for off-sim kinematics probing. + + A fresh ``MjData`` seeded with ``qpos`` alone is NOT equivalent: MolmoSpaces places the robot in a scene + whose pose also rides on state outside ``qpos`` (mocap bodies among it), which a fresh buffer resets to + the model defaults — the grasp site then resolves metres away from the live one. Copying the whole + struct keeps every such field, so the probe differs from the live scene only in the joints the caller + sets, and copying into a retained buffer keeps the per-step allocation out of the control loop. + """ + if self._scratch is None: + self._scratch = mujoco.MjData(move_group.mj_model) # pyright: ignore[reportAttributeAccessIssue] + mujoco.mj_copyData(self._scratch, move_group.mj_model, move_group.mj_data) # pyright: ignore[reportAttributeAccessIssue] + return self._scratch + + def _fk(self, q: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """The grasp-site world pose a candidate arm configuration reaches. + + Evaluated on a scratch ``MjData`` seeded from the live scene (objects intact), so the live sim is never + perturbed: set the arm joints, propagate, read the leaf frame. The inverse of ``_ik``. + """ + arm = self._robot_view.get_move_group(mapping.MOLMO_ARM_GROUP) + data = self._scratch_data(arm) + data.qpos[np.asarray(arm.joint_posadr)] = np.asarray(q, dtype=np.float64).reshape(-1) + mujoco.mj_forward(arm.mj_model, data) # pyright: ignore[reportAttributeAccessIssue] + return _leaf_pose(arm, data) + + def _ik(self, target_pos: np.ndarray, target_rot: np.ndarray) -> np.ndarray: + """Absolute world grasp-site target -> the arm joint targets that reach it. + + Damped-least-squares differential IK on MuJoCo's own leaf-frame Jacobian, mirroring the LIBERO rig's + solver. It iterates on a scratch ``MjData`` seeded from the live scene (objects intact), so probing + candidate joint configurations never perturbs the sim being stepped. Joint targets stay inside the + move group's limits, and a target the arm cannot reach yields the closest configuration the iteration + reached rather than raising — an unreachable waypoint holds near the limit instead of aborting a trial. + """ + arm = self._robot_view.get_move_group(mapping.MOLMO_ARM_GROUP) + model = arm.mj_model + posadr = np.asarray(arm.joint_posadr) + veladr = np.asarray(arm.joint_veladr) + limits = np.asarray(arm.joint_pos_limits, dtype=np.float64) + data = self._scratch_data(arm) + q = np.asarray(arm.joint_pos, dtype=np.float64).copy() + for _ in range(_IK_ITERS): + data.qpos[posadr] = q + mujoco.mj_forward(model, data) # pyright: ignore[reportAttributeAccessIssue] + cur_pos, cur_rot = _leaf_pose(arm, data) + err = _pose_error(target_pos, target_rot, cur_pos, cur_rot) + if np.linalg.norm(err) < _IK_TOL: + break + jac = np.zeros((6, model.nv)) + self._leaf_jacobian(arm, model, data, jac) + jac = jac[:, veladr] + dq = jac.T @ np.linalg.solve(jac @ jac.T + _IK_DAMPING**2 * np.eye(6), err) + q = np.clip(q + dq, limits[:, 0], limits[:, 1]) + return q + + @staticmethod + def _leaf_jacobian(move_group: Any, model: Any, data: Any, out: np.ndarray) -> None: + """The ``(6, nv)`` leaf-frame Jacobian into *out*, evaluated on *data*. + + Mirrors the move group's own ``get_jacobian`` but against a caller-supplied ``MjData``, which the IK + iteration needs (the group's method is bound to the live one). + """ + if move_group.leaf_frame_type == _SITE_FRAME: + mujoco.mj_jacSite(model, data, out[:3], out[3:], move_group.leaf_frame_id) # pyright: ignore[reportAttributeAccessIssue] + else: + mujoco.mj_jacBody(model, data, out[:3], out[3:], move_group.leaf_frame_id) # pyright: ignore[reportAttributeAccessIssue] + + def _observe(self, env_obs: dict[str, Any]) -> dict[str, Any]: + """The raw observation payload for one env frame: measured joints, the eef world pose, grip, camera frames. + + 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, since obs exposes only a robot-relative tcp pose. + """ + arm = self._robot_view.get_move_group(mapping.MOLMO_ARM_GROUP) + 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 = { + mapping.OBS_JOINT_POS: np.asarray(arm.joint_pos, dtype=np.float32), + mapping.OBS_JOINT_VEL: np.asarray(arm.joint_vel, dtype=np.float32), + mapping.OBS_EEF_POS: eef_world[:3, 3].astype(np.float32), + mapping.OBS_EEF_QUAT: eef_quat.astype(np.float32), + mapping.OBS_GRIP: np.float32( + mapping.normalize_grip_qpos(env_obs[mapping.MOLMO_OBS_QPOS][mapping.MOLMO_GRIPPER_GROUP]) + ), + # The full MuJoCo generalized state: every body's pose + velocity, objects included. + mapping.OBS_SIM_STATE: self._full_physics_state(), + } + for name in self._camera_names: + payload[name] = np.ascontiguousarray(env_obs[name]) + return payload + + def _full_physics_state(self) -> np.ndarray: + """The scene's complete integrable state: ``mjSTATE_INTEGRATION``, the minimal subset a deterministic + MuJoCo sim restores from to reproduce its forward trajectory — positions and velocities, and with them + mocap bodies, actuator activation, controls and the solver warm-start. Object poses in it let analysis + recompute success. Positions start at index 1, after the scalar time.""" + data = self._robot_view.mj_data + model = data.model + spec = mujoco.mjtState.mjSTATE_INTEGRATION # pyright: ignore[reportAttributeAccessIssue] + state = np.empty(mujoco.mj_stateSize(model, spec), dtype=np.float64) # pyright: ignore[reportAttributeAccessIssue] + mujoco.mj_getState(model, data, state, spec) # pyright: ignore[reportAttributeAccessIssue] + return state + + def close(self) -> None: + if self._sampler is not None: + self._sampler.close() + self._sampler = None + self._task = None + + +# rules-allow: stranded-definition — this file keeps its pure MuJoCo helpers at module scope as a set: +# `_assert_measures_at_grasp_site`, `_leaf_pose` and `_pose_error` all take plain model/data arguments, hold no +# `self`, and are called only from `MolmoSpacesEnv`. Moving one into the class splits the set for no gain; +# moving all three is a layout decision for the file, not a fix to this definition. +def _leaf_pose(move_group: Any, data: Any) -> tuple[np.ndarray, np.ndarray]: + """A move group's leaf-frame world pose read off *data* — which may be a scratch ``MjData``, unlike the + group's own ``leaf_frame_to_world``, so IK can probe candidate joints without touching the live sim.""" + if move_group.leaf_frame_type == _SITE_FRAME: + pos, mat = data.site_xpos[move_group.leaf_frame_id], data.site_xmat[move_group.leaf_frame_id] + else: + pos, mat = data.xpos[move_group.leaf_frame_id], data.xmat[move_group.leaf_frame_id] + return np.array(pos, dtype=np.float64), np.array(mat, dtype=np.float64).reshape(3, 3) + + +def _pose_error(target_pos: np.ndarray, target_rot: np.ndarray, cur_pos: np.ndarray, cur_rot: np.ndarray) -> np.ndarray: + """The world-frame 6-vector error ``[translation, axis-angle rotation]`` from a measured to a target pose. + + Both halves are expressed in the world frame, matching the world-frame leaf Jacobian the IK step solves + against. The rotation error is the axis-angle of ``R_target @ R_cur^T``, via MuJoCo's quaternion helpers. + """ + quat = np.zeros(4) + mujoco.mju_mat2Quat(quat, np.ascontiguousarray((target_rot @ cur_rot.T).reshape(9))) # pyright: ignore[reportAttributeAccessIssue] + rot_err = np.zeros(3) + mujoco.mju_quat2Vel(rot_err, quat, 1.0) # pyright: ignore[reportAttributeAccessIssue] + return np.concatenate([np.asarray(target_pos, dtype=np.float64).reshape(3) - cur_pos, rot_err]) + + +def main() -> None: + parser = argparse.ArgumentParser(description='Serve MolmoSpaces over the env-server protocol.') + parser.add_argument(protocol.OPT_HOST, default='localhost') + parser.add_argument(protocol.OPT_PORT, type=int, required=True) + parser.add_argument( + mapping.OPT_BENCHMARK_DIR, required=True, help=f'dir containing {mapping.MOLMO_BENCHMARK_MANIFEST}' + ) + parser.add_argument( + mapping.OPT_TASK_HORIZON_STEPS, + type=int, + default=None, + help='override the benchmark horizon (steps per episode)', + ) + args = parser.parse_args() + if not os.environ.get(mapping.ASSETS_DIR_ENV): + parser.error(f'{mapping.ASSETS_DIR_ENV} must point at the MolmoSpaces asset packs') + env = MolmoSpacesEnv(Path(args.benchmark_dir), args.task_horizon_steps) + EnvServer(env, args.host, args.port).serve_forever() + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/keys.py b/positronic/simulator/molmo_spaces/keys.py new file mode 100644 index 000000000..fcbdc9bb7 --- /dev/null +++ b/positronic/simulator/molmo_spaces/keys.py @@ -0,0 +1,7 @@ +"""The keys of a MolmoSpaces trial's params: the episode the eval selects and the horizon the sim enforces.""" + +# The benchmark episode a trial runs, as ``task_params`` names it from the env's task records; ``_reset_token`` +# reads it back into the token that selects the episode. +EPISODE_INDEX = 'eval.episode_index' +# The sim-enforced episode deadline in sim-seconds; the eval config sets the trial's backstop deadline from it. +TASK_HORIZON = 'eval.task_horizon' diff --git a/positronic/simulator/molmo_spaces/launcher.py b/positronic/simulator/molmo_spaces/launcher.py new file mode 100644 index 000000000..4ac47956d --- /dev/null +++ b/positronic/simulator/molmo_spaces/launcher.py @@ -0,0 +1,128 @@ +"""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``). Both env vars pass through from the caller; unset, ``MUJOCO_GL`` takes the backend +the host platform offers — ``egl`` (GPU) on Linux, ``cgl`` on macOS, which has no EGL and rejects it. A +GPU-less Linux box overrides with ``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 import protocol +from positronic.simulator.env_server.launcher import ensure_pinned_checkout, serve_subprocess +from positronic.simulator.molmo_spaces import mapping + +_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 ships no lockfile, so a bare install re-resolves every transitive dep on each fresh box. This +# constraints file pins the full resolution (a frozen known-good venv, minus molmo-spaces' own editable line), +# fed to the install via ``-c`` so the pinned commit always builds the same environment. Regenerate it when +# ``_MOLMO_COMMIT`` bumps — see the file header. +_MOLMO_CONSTRAINTS = Path(__file__).parent / 'molmo_constraints.txt' + +# 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' + +# ``env.py`` imports positronic's ``env_server`` off PYTHONPATH, which needs ``websockets`` (the wire server) and +# ``msgpack`` (the frame codec). MolmoSpaces currently pulls both, but that is incidental to its own deps — install +# them explicitly so env_server's wire contract holds even if MolmoSpaces drops them. Constraints mirror positronic's. +_WIRE_DEPS = ('websockets>=15.0.1', 'msgpack') + + +@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 ensure_molmo_venv() -> Path: + """The MolmoSpaces venv python, after ensuring the pinned checkout and its installed stack exist. + + Install the stack before returning: a cold first install far exceeds any 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 ``molmo_constraints.txt`` pins the transitive resolution (``-c``) for a reproducible env. + """ + venv = _MOLMO_SRC / '.venv' + with _checkout_lock(): + src = ensure_pinned_checkout(_MOLMO_REPO, _MOLMO_COMMIT, _MOLMO_SRC) + if not venv.exists(): + subprocess.run(['uv', 'venv', '--python', _MOLMO_PYTHON, str(venv)], check=True) + subprocess.run( + ['uv', 'pip', 'install', '-c', str(_MOLMO_CONSTRAINTS), '-e', f'.[{_MOLMO_EXTRA}]', *_WIRE_DEPS], + cwd=str(src), + env={**os.environ, 'VIRTUAL_ENV': str(venv)}, + check=True, + ) + return venv / 'bin' / 'python' + + +def molmo_subprocess_env() -> dict[str, str]: + """The environment a molmo-venv script runs under: the positronic-free ``env_server``/``mapping`` on + PYTHONPATH and a GL backend. GPU OpenGL by default; a GPU-less box exports MUJOCO_GL=osmesa, or relies on + mesa's software EGL, for CPU rendering.""" + return { + **os.environ, + 'PYTHONPATH': os.pathsep.join([str(_ENV_SERVER_DIR), str(_MAPPING_DIR)]), + mapping.GL_BACKEND_ENV: os.environ.get(mapping.GL_BACKEND_ENV, mapping.GL_BACKEND_DEFAULT), + } + + +def _spawn(host: str, port: int, benchmark_dir: Path, task_horizon_steps: int | None) -> subprocess.Popen: + # env.py exits on these before it binds the port. Check them here, where the failure can name the missing + # precondition instead of reaching the caller as a bare pre-bind exit status. + if not os.environ.get(mapping.ASSETS_DIR_ENV): + raise ValueError(f'{mapping.ASSETS_DIR_ENV} must point at the MolmoSpaces asset packs') + if not benchmark_dir.is_dir(): + raise ValueError(f'benchmark dir {benchmark_dir} does not exist') + python = ensure_molmo_venv() + command = [ + str(python), + str(_ENV_SCRIPT), + protocol.OPT_HOST, + host, + protocol.OPT_PORT, + str(port), + mapping.OPT_BENCHMARK_DIR, + str(benchmark_dir), + ] + if task_horizon_steps is not None: + command += [mapping.OPT_TASK_HORIZON_STEPS, str(task_horizon_steps)] + return subprocess.Popen(command, env=molmo_subprocess_env()) + + +def serve_molmo_spaces( + benchmark_dir: Path, host: str = 'localhost', task_horizon_steps: int | None = None +) -> 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. ``task_horizon_steps`` optionally overrides + the benchmark's own horizon (mirroring MolmoSpaces' ``--task_horizon_steps``); ``None`` reads it per episode. + """ + return serve_subprocess(lambda host, port: _spawn(host, port, benchmark_dir, task_horizon_steps), host) diff --git a/positronic/simulator/molmo_spaces/mapping.py b/positronic/simulator/molmo_spaces/mapping.py new file mode 100644 index 000000000..02d14b9ce --- /dev/null +++ b/positronic/simulator/molmo_spaces/mapping.py @@ -0,0 +1,266 @@ +"""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 numpy plus the positronic-free ``protocol`` (which owns the wire command tags), 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. +""" + +import sys +from collections.abc import Callable, Iterable +from typing import Any, TypeAlias + +import numpy as np + +# ``protocol`` lands as a package on the positronic side and flat on ``PYTHONPATH`` inside the molmo venv, +# where ``positronic`` is not installed — the same two-shape import ``server`` uses. +try: + from positronic.simulator.env_server import protocol +except ImportError: + import protocol # pyright: ignore[reportMissingImports] + +# The DROID rig runs 7 Franka arm joints; the reset token's per-move-group action names them 'arm'/'gripper'. +NUM_ARM_JOINTS = 7 + +# An absolute world target ``(translation, 3x3 rotation)`` -> the arm joint targets that reach it. Supplied +# by ``env.py``, which holds the model this module deliberately does not. +IkSolver: TypeAlias = Callable[[np.ndarray, np.ndarray], Any] +MOLMO_ARM_GROUP = 'arm' +MOLMO_GRIPPER_GROUP = 'gripper' + +# The MolmoSpaces site the arm move group's leaf frame resolves to, and so the frame this adoption reports +# poses in and resolves Cartesian targets against. The eval declares its recorded model's control frame at the +# same physical point, and ``env.py`` checks the live scene against this name so the two cannot drift apart. +# A scene prefixes every model name with the robot's namespace (``robot_0/``), so the live name ends with this. +MOLMO_GRASP_SITE = 'gripper/grasp_site' + +# Where the MolmoSpaces asset packs live, and the subdirectory of that root holding the benchmarks. +ASSETS_DIR_ENV = 'MLSPACES_ASSETS_DIR' +ASSETS_BENCHMARKS_DIR = 'benchmarks' + +# A benchmark dir's manifest: the JSON list of episode specs ``load_all_episodes`` reads, and what marks a +# directory as a benchmark for discovery. +MOLMO_BENCHMARK_MANIFEST = 'benchmark.json' + +# The env-server subprocess CLI, spelled by the launcher building the command and by ``env.py``'s parser +# declaring it — two interpreters, so a rename that misses one fails at spawn rather than at import. +OPT_BENCHMARK_DIR = '--benchmark_dir' +OPT_TASK_HORIZON_STEPS = '--task_horizon_steps' + +# MuJoCo's backend selector, and the backend this adoption asks for. MuJoCo validates the value against the +# host platform and raises on one it does not offer there, so the default follows the platform: EGL is the +# headless-GPU path on Linux, CGL the only context macOS has. A GPU-less Linux box overrides with osmesa. +GL_BACKEND_ENV = 'MUJOCO_GL' +GL_BACKEND_DEFAULT = 'cgl' if sys.platform == 'darwin' else 'egl' + +# The reset token: which benchmark episode to build, and the seed overriding the episode spec's own. +TOKEN_EPISODE_INDEX = 'episode_index' +TOKEN_SEED = 'seed' + +# The reset frame's scene meta: the episode's resolved language goal and the ProcTHOR house it runs in. +META_TASK = 'task' +META_HOUSE_INDEX = 'house_index' + +# The MolmoSpaces observation field holding the per-move-group joint positions, which is where the +# gripper closure is read from. +MOLMO_OBS_QPOS = 'qpos' + +# The benchmark episode spec's task definition, and the horizon it declares in sim-seconds. +MOLMO_EPISODE_TASK = 'task' +MOLMO_TASK_HORIZON_SEC = 'task_horizon_sec' + +# The raw observation payload ``env.py`` reports and ``MolmoAdapter`` reads back. +OBS_JOINT_POS = 'joint_pos' +OBS_JOINT_VEL = 'joint_vel' +OBS_EEF_POS = 'eef_pos' +OBS_EEF_QUAT = 'eef_quat' +OBS_GRIP = 'grip' +OBS_SIM_STATE = 'sim_state' + +# 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 is_rgb_frame(value: Any) -> bool: + """Whether an observation entry is a rendered camera frame — which is how the camera keys are discovered, + since MolmoSpaces names them per benchmark and the obs dict carries no other HWC uint8 array.""" + return isinstance(value, np.ndarray) and value.ndim == 3 and value.shape[2] == 3 and value.dtype == np.uint8 + + +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 unpack_wire_pose(vector: Any) -> tuple[np.ndarray, np.ndarray]: + """A wire pose ``[t(3), R(9)]`` -> ``(translation, 3x3 rotation)``. + + The client encodes every pose with ``Transform3D.as_vector(ROTATION_MATRIX)``: translation first, then the + rotation matrix row-major. + """ + vec = np.asarray(vector, dtype=np.float64).reshape(-1) + if vec.shape[0] != 12: + raise ValueError(f'wire pose must be [t(3), R(9)], got {vec.shape[0]} values') + return vec[:3].copy(), vec[3:].reshape(3, 3).copy() + + +def compose_world_delta(cur_pos: Any, cur_rot: Any, delta_pos: Any, delta_rot: Any) -> tuple[np.ndarray, np.ndarray]: + """The absolute pose a world-frame ``cartesian_delta`` targets from a measured pose. + + Translation adds in the world frame and rotation left-multiplies (``goal_ori = R(delta) @ ee_ori``) — the + convention positronic's ``apply_cartesian_delta`` and LIBERO's own delta bridging both use. + """ + return ( + np.asarray(cur_pos, dtype=np.float64).reshape(3) + np.asarray(delta_pos, dtype=np.float64).reshape(3), + np.asarray(delta_rot, dtype=np.float64).reshape(3, 3) @ np.asarray(cur_rot, dtype=np.float64).reshape(3, 3), + ) + + +def _require_ik(ik: IkSolver | None, kind: str) -> IkSolver: + """The caller's IK solver, or a loud failure — a Cartesian target is unresolvable without the live model.""" + if ik is None: + raise ValueError(f'command {kind!r} needs an ik solver; none was supplied') + return ik + + +def wire_command_to_arm_action( + command: dict[str, Any], current_q: Any, *, ik: IkSolver | None = None, current_eef: tuple[Any, Any] | None = None +) -> np.ndarray: + """A tagged wire command + the live measured arm joints -> the 7 absolute joint targets molmo steps. + + This is where the adoption covers the canonical command contract: MolmoSpaces' Franka natively takes only + joint-position targets, so every canonical type is converted into one. ``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. + + The Cartesian pair needs the live model, which this module deliberately does not hold: the caller passes + ``ik`` (an absolute world target ``(pos, rot)`` -> joint targets) and, for ``cartesian_delta``, the measured + ``current_eef`` pose the delta composes onto. Both are supplied by ``env.py``, which owns the sim. + """ + current = np.asarray(current_q, dtype=np.float32).reshape(-1) + match command[protocol.COMMAND_TYPE]: + case protocol.JOINT_POS: + target = np.asarray(command[protocol.COMMAND_JOINT_POS], dtype=np.float32).reshape(-1) + case protocol.JOINT_VEL: + dq = np.asarray(command[protocol.COMMAND_JOINT_VEL], 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 protocol.HOLD: + target = current + case protocol.CARTESIAN: + solver = _require_ik(ik, protocol.CARTESIAN) + target = np.asarray(solver(*unpack_wire_pose(command[protocol.COMMAND_POSE])), dtype=np.float32).reshape(-1) + case protocol.CARTESIAN_DELTA: + solver = _require_ik(ik, protocol.CARTESIAN_DELTA) + if current_eef is None: + raise ValueError(f'command {protocol.CARTESIAN_DELTA!r} needs the measured eef pose; none supplied') + delta_pos, delta_rot = unpack_wire_pose(command[protocol.COMMAND_DELTA]) + target_pos, target_rot = compose_world_delta(*current_eef, delta_pos, delta_rot) + target = np.asarray(solver(target_pos, target_rot), dtype=np.float32).reshape(-1) + case other: + raise ValueError( + f'{other!r} is not a canonical command type; the contract is {list(protocol.CANONICAL_COMMAND_TYPES)}' + ) + return target.astype(np.float32) + + +def resolve_camera_key(available: Any, key: str, variants: tuple[str, ...] = ()) -> str: + """The MolmoSpaces observation key to read for a camera role, mirroring the upstream policy's precedence. + + A present benchmark variant wins over ``key`` (matching molmo_spaces pi_policy); with no variants ``key`` + is read as-is. Raises with the candidate list on a miss. + """ + keys = set(available) + for candidate in (*variants, key): + if candidate in keys: + return candidate + raise KeyError(f'observation has none of {(*variants, key)}; available: {sorted(keys)}') + + +def resolve_episode_seed(episode: Any, episode_index: int, override_seed: int | None = None) -> int: + """The seed an episode runs under, mirroring MolmoSpaces' own precedence. + + An explicit ``override_seed`` wins, then the episode spec's own seed. A spec carrying none falls back to + the episode index, which is what ``JsonEvalRunner.get_episode_seed`` does — a constant instead would put + every unseeded episode of a benchmark on one random stream, and none of them on the native one. + """ + if override_seed is not None: + return int(override_seed) + spec_seed = getattr(episode, 'seed', None) + return int(spec_seed) if spec_seed is not None else int(episode_index) + + +def declared_task_horizon_sec(declared: Iterable[float | None]) -> float: + """The one horizon a benchmark declares, in sim-seconds, over every episode's ``task_horizon_sec``. + + The horizon belongs to the benchmark, not to an episode within it, so a benchmark that declares none, one + that disagrees with itself, and one that declares a non-positive span all have no horizon to run at. + Callers pass the values already read from their own representation of the specs: parsed episode objects in + the molmo venv, raw JSON on the positronic side. + """ + horizons = set() + for value in declared: + if value is None: + raise ValueError( + f'benchmark episodes carry no {MOLMO_TASK_HORIZON_SEC} in their task dict — the horizon is part ' + 'of the task definition; add it to the benchmark' + ) + if value <= 0: + raise ValueError(f'benchmark declares a non-positive {MOLMO_TASK_HORIZON_SEC} of {value}s') + horizons.add(value) + if len(horizons) != 1: + raise ValueError(f'benchmark declares inconsistent {MOLMO_TASK_HORIZON_SEC} values {sorted(horizons)}') + return float(horizons.pop()) + + +def resolve_task_horizon_steps(episodes: Any, policy_dt_ms: float, override_steps: int | None = None) -> int: + """A benchmark's enforced horizon in policy steps, mirroring MolmoSpaces' own resolution. + + Upstream's ``determine_task_horizon`` (``evaluation/eval_main.py``, the entrypoint its README documents) + resolves in this order and nothing else: an explicit ``--task_horizon_steps`` override, then the benchmark's + own ``task_horizon_sec`` from the episodes' task dicts, converted with ``round(sec * 1000 / policy_dt_ms)``. + It raises when any episode declares none, and again when the episodes disagree. This reproduces all three, + raises included: ``JsonBenchmarkEvalConfig.task_horizon``'s 500-step default is a config default upstream + overwrites before the runner ever sees it. A horizon that resolves below one step is refused on either + path, so no route reaches the task with a budget it expires inside. + """ + if override_steps is not None: + if override_steps < 1: + raise ValueError(f'task_horizon_steps override must be at least 1 step, got {override_steps}') + return override_steps + sec = declared_task_horizon_sec(episode.task.get(MOLMO_TASK_HORIZON_SEC) for episode in episodes) + steps = round(sec * 1000.0 / policy_dt_ms) + if steps < 1: + raise ValueError( + f'benchmark {MOLMO_TASK_HORIZON_SEC} of {sec}s rounds to {steps} steps at a {policy_dt_ms}ms policy ' + 'period — the episode would expire before its first action' + ) + return steps diff --git a/positronic/simulator/molmo_spaces/molmo_constraints.txt b/positronic/simulator/molmo_spaces/molmo_constraints.txt new file mode 100644 index 000000000..bfa0d807f --- /dev/null +++ b/positronic/simulator/molmo_spaces/molmo_constraints.txt @@ -0,0 +1,177 @@ +# Pinned transitive dependency set for the MolmoSpaces env-server venv (ensure_molmo_venv, launcher.py). +# MolmoSpaces ships no lockfile, so a cold 'uv pip install -e .[mujoco]' re-resolves every dep afresh; this +# constraints file (uv pip install -c ...) pins that resolution so every box builds the same env. +# +# Regenerate from a known-good venv (the one the parity/e2e run validated): +# VIRTUAL_ENV= uv pip freeze | grep -v '^-e ' > molmo_constraints.txt (then re-add this header) +# molmo-spaces itself is installed editable from source (-e .), so its freeze line is excluded here. +absl-py==2.5.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aiosignal==1.4.0 +annotated-types==0.8.0 +antlr4-python3-runtime==4.9.3 +anyio==4.14.2 +attrs==26.1.0 +beaker-py==2.7.1 +boto3==1.43.55 +botocore==1.43.55 +certifi==2026.7.22 +cffi==2.1.0 +charset-normalizer==3.4.9 +click==8.4.2 +cloudpickle==3.1.2 +coacd==1.0.11 +compress-json==1.1.1 +contourpy==1.3.3 +cryptography==49.0.0 +cycler==0.12.1 +datasets==5.0.0 +decorator==5.3.1 +decord==0.6.0 +dill==0.4.1 +einops==0.8.2 +etils==1.14.0 +evdev==1.9.3 +farama-notifications==0.0.6 +fastjsonschema==2.21.2 +ffmpeg-python==0.2.0 +filelock==3.32.0 +fonttools==4.63.0 +frozenlist==1.8.0 +fsspec==2026.4.0 +ftfy==6.3.1 +future==1.0.0 +glfw==2.10.2 +google-crc32c==1.8.0 +grpcio==1.83.0 +gymnasium==1.3.0 +h11==0.16.0 +h5py==3.16.0 +hf-xet==1.5.2 +hidapi==0.15.0 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.24.0 +idna==3.18 +imageio==2.37.4 +imageio-ffmpeg==0.6.0 +iniconfig==2.3.0 +jax==0.6.2 +jaxlib==0.6.2 +jinja2==3.1.6 +jmespath==1.1.0 +joblib==1.5.3 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +jupyter-core==5.9.1 +kiwisolver==1.5.0 +lazy-loader==0.5 +lmdb==1.7.5 +lxml==6.1.1 +markdown==3.10.2 +markupsafe==3.0.3 +matplotlib==3.11.1 +ml-dtypes==0.5.4 +molmospaces-resources==0.0.1b4 +moviepy==2.2.1 +mpmath==1.3.0 +msgpack==1.2.1 +msgpack-numpy==0.4.8 +mujoco==3.5.0 +mujoco-mjx==3.5.0 +mujoco-warp==3.5.0.2 +multidict==6.7.1 +multiprocess==0.70.19 +nbformat==5.10.4 +nbstripout==0.9.1 +networkx==3.6.1 +nltk==3.9.4 +numpy==2.4.6 +numpy-quaternion==2024.0.13 +nvidia-cublas-cu12==12.6.4.1 +nvidia-cuda-cupti-cu12==12.6.80 +nvidia-cuda-nvrtc-cu12==12.6.77 +nvidia-cuda-runtime-cu12==12.6.77 +nvidia-cudnn-cu12==9.5.1.17 +nvidia-cufft-cu12==11.3.0.4 +nvidia-cufile-cu12==1.11.1.6 +nvidia-curand-cu12==10.3.7.77 +nvidia-cusolver-cu12==11.7.1.2 +nvidia-cusparse-cu12==12.5.4.2 +nvidia-cusparselt-cu12==0.6.3 +nvidia-ml-py==13.610.43 +nvidia-nccl-cu12==2.26.2 +nvidia-nvjitlink-cu12==12.6.85 +nvidia-nvtx-cu12==12.6.77 +omegaconf==2.3.1 +open-clip-torch==3.2.0 +opencv-python==5.0.0.93 +opt-einsum==3.4.0 +packaging==26.2 +pandas==3.0.5 +pillow==11.3.0 +platformdirs==4.11.0 +pluggy==1.6.0 +prior==1.0.3 +proglog==0.1.12 +propcache==0.5.2 +protobuf==6.33.6 +psutil==7.2.2 +pyarrow==25.0.0 +pycparser==3.0 +pydantic==2.13.4 +pydantic-core==2.46.4 +pygithub==2.9.1 +pygments==2.20.0 +pyjwt==2.13.0 +pynacl==1.6.2 +pynput==1.8.2 +pyopengl==3.1.10 +pyparsing==3.3.2 +pytest==9.1.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +python-xlib==0.33 +pyyaml==6.0.3 +qrcode==8.2 +referencing==0.37.0 +regex==2026.7.19 +requests==2.34.2 +rpds-py==2026.6.3 +s3transfer==0.19.2 +safetensors==0.8.0 +scikit-image==0.26.0 +scipy==1.17.1 +sentry-sdk==2.66.1 +setuptools==83.0.0 +shapely==2.1.2 +shortuuid==1.0.13 +six==1.17.0 +stringcase==1.2.0 +sympy==1.14.0 +teledex==0.0.7 +tensorboard==2.21.0 +tensorboard-data-server==0.7.2 +termcolor==3.3.0 +tifffile==2026.3.3 +timm==1.0.28 +toppra==0.6.3 +torch==2.7.1 +torchvision==0.22.1 +tqdm==4.69.0 +traitlets==5.15.1 +trimesh==4.12.2 +triton==3.3.1 +typing-extensions==4.16.0 +typing-inspection==0.4.2 +urllib3==2.7.0 +wandb==0.28.1 +warp-lang==1.15.0 +wcwidth==0.8.2 +websockets==16.1.1 +werkzeug==3.1.8 +xxhash==3.8.1 +yarl==1.24.5 +zipp==4.1.0 +zstandard==0.25.0 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/e2e.py b/positronic/simulator/molmo_spaces/tests/e2e.py new file mode 100644 index 000000000..81cb88905 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/e2e.py @@ -0,0 +1,92 @@ +"""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.tests.e2e --benchmark_dir +""" + +import argparse +from pathlib import Path + +import numpy as np + +from positronic import keys +from positronic.simulator.env_server import protocol +from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.molmo_spaces.adapter import DEFAULT_CAMERA_DICT, MolmoAdapter +from positronic.simulator.molmo_spaces.launcher import serve_molmo_spaces + + +def _check_sim_state(adapter: MolmoAdapter, raw_obs: dict) -> np.ndarray: + """The privileged full MuJoCo state must survive the wire and reach the recorder as a finite qpos+qvel vector.""" + sim_state = adapter.privileged(raw_obs)[mapping.OBS_SIM_STATE] + assert isinstance(sim_state, np.ndarray) and sim_state.ndim == 1 and sim_state.size > 0, ( + f'privileged sim_state malformed: {type(sim_state)} shape={getattr(sim_state, "shape", None)}' + ) + assert np.isfinite(sim_state).all(), 'privileged sim_state carries non-finite values' + return sim_state + + +def run( + benchmark_dir: Path, + *, + episodes: int = 1, + steps: int = 5, + camera_dict: dict[str, str] | None = None, + task_horizon_steps: int | None = None, +) -> None: + """Reset + step the first ``episodes`` benchmark episodes over the socket, mapping each frame with the adapter.""" + camera_dict = camera_dict or DEFAULT_CAMERA_DICT + adapter = MolmoAdapter(camera_dict) + with serve_molmo_spaces(benchmark_dir, task_horizon_steps=task_horizon_steps) as (host, port): + conn = EnvConnection(host, port) + try: + for i in range(episodes): + frame = conn.reset({mapping.TOKEN_EPISODE_INDEX: i, mapping.TOKEN_SEED: None}) + obs = adapter.observations(frame[protocol.FRAME_OBS]) + assert keys.ROBOT_STATE in obs and keys.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[keys.ROBOT_STATE].q + assert q.shape == (7,), f'unexpected joint shape {q.shape}' + sim_state = _check_sim_state(adapter, frame[protocol.FRAME_OBS]) + print( + f' episode {i}: reset ok — task={frame[protocol.FRAME_META][mapping.META_TASK]!r} ' + f'grip={obs[keys.GRIP]:.3f} ' + f'q0={q[0]:.4f} sim_state={sim_state.size}d' + ) + out = {protocol.FRAME_DONE: False} + for _ in range(steps): + hold = {protocol.ACTION_COMMAND: {protocol.COMMAND_TYPE: protocol.HOLD}, protocol.ACTION_GRIP: 0.0} + out = conn.step(hold) + adapter.observations(out[protocol.FRAME_OBS]) # the mapping round-trips on step frames too + _check_sim_state(adapter, out[protocol.FRAME_OBS]) + print(f' episode {i}: {steps} steps ok (done={out[protocol.FRAME_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) + parser.add_argument( + '--task_horizon_steps', type=int, default=None, help='override the benchmark horizon (steps per episode)' + ) + args = parser.parse_args() + run(Path(args.benchmark_dir), episodes=args.episodes, steps=args.steps, task_horizon_steps=args.task_horizon_steps) + + +if __name__ == '__main__': + main() 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..b6353549d --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/make_fixture.py @@ -0,0 +1,56 @@ +# /// 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 + +from positronic.simulator.molmo_spaces import mapping + +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 { + mapping.OBS_JOINT_POS: np.array([0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.785], dtype=np.float32), + mapping.OBS_JOINT_VEL: np.linspace(-0.2, 0.2, 7, dtype=np.float32), + mapping.OBS_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. + mapping.OBS_EEF_QUAT: np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + mapping.OBS_GRIP: np.float32(0.5), + mapping.MOLMO_WRIST_CAMERA: _marked_frame((200, 40, 40)), # reddish wrist view + mapping.MOLMO_EXTERIOR_CAMERA: _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/make_replay_fixture.py b/positronic/simulator/molmo_spaces/tests/make_replay_fixture.py new file mode 100644 index 000000000..2ba77c0ea --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/make_replay_fixture.py @@ -0,0 +1,214 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +"""Regenerate a deterministic-replay fixture from a recorded MolmoSpaces eval episode. + +``test_replay.py`` replays a real pi05 rollout open-loop against the sim and asserts it reproduces. This +script distils one recorded episode into the fixture that replay needs: the commanded joint targets and grip +per step, taken from the recording, plus checkpoints of the ``sim_state`` those commands produce, taken by +replaying them here. The commands are what the recording pins; the checkpoints pin the integration's current +trajectory, so a later run that drifts from it fails. Regenerate them together whenever the recorded +``sim_state`` changes shape or the pinned MolmoSpaces commit moves. + +Two properties make the distillation exact. The recorded commands are *absolute* joint targets, so the +replay never reads the measured state back — it is genuinely open-loop, and the only thing under test is the +sim rollout plus the env-server path. And the proxy applies whichever command was last received when it +steps, so sampling the command signal at each observation frame's timestamp (``Signal.time`` — the same +last-value-at-or-before semantics a pimm receiver has) reconstructs the stream the sim saw, unchanged +commands included. + +It reconstructs that stream only as far as the recording pins it: an episode's command signals stop before +its observations do (internal#130), so the fixture keeps the prefix up to the final recorded command and +counts the rest as the recording's gap. + +Commands are stored as float32, the dtype ``env.py`` casts them to, so the fixture holds the bits the sim +actually applied rather than the float64 the recorder wrote. + +Run (needs positronic for the dataset reader, and the MolmoSpaces assets for the replay — hence +``--locked``, not ``--no-project``):: + + MLSPACES_ASSETS_DIR=... MUJOCO_GL=egl EGL_PLATFORM=surfaceless LIBGL_ALWAYS_SOFTWARE=1 \ + uv run --locked python positronic/simulator/molmo_spaces/tests/make_replay_fixture.py \ + --dataset_dir ~/.cache/positronic/s3/_/inference/molmo_battle_test/2026-07-29/sweep_jp \ + --episode_index 3 --episode_index 6 + +Output: ``replay_ep.npz`` next to this script, one per episode (tens of KB — actions and checkpoints +only, never the videos). +""" + +import argparse +import os +import re +from pathlib import Path + +import numpy as np + +from positronic import keys +from positronic.dataset.local_dataset import DiskEpisode +from positronic.dataset.signal import Signal +from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol +from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.molmo_spaces import keys as molmo_keys +from positronic.simulator.molmo_spaces import launcher, mapping + +# The fixture's own fields, as a distilled episode records them. +FIELD_EPISODE_INDEX = 'episode_index' +FIELD_BENCHMARK_PATH = 'benchmark_path' +FIELD_TASK = 'task' +FIELD_COMMANDS = 'commands' +FIELD_GRIPS = 'grips' +FIELD_UNREPLAYABLE_TAIL_STEPS = 'unreplayable_tail_steps' +FIELD_CHECKPOINT_STEPS = 'checkpoint_steps' +FIELD_CHECKPOINT_SIM_STATE = 'checkpoint_sim_state' +FIELD_EXPECTED_SUCCESS = 'expected_success' + +# Checkpoint stride over the replayed steps: dense enough that drift is caught early rather than only at the +# end state, sparse enough to keep the fixture small. The final step is always included on top. +CHECKPOINT_STRIDE = 8 + +# The eval CLI records its full command line in the dataset's run metadata; the benchmark the episodes were +# recorded against is the one argument the replay must resolve on the box it runs on. +_BENCHMARK_ARG = re.compile(r'--eval\.benchmark_dir=(\S+)') + + +def find_episode_dir(dataset_dir: Path, episode_index: int) -> Path: + """The recorded episode directory whose spec carries ``episode_index``.""" + for path in sorted(dataset_dir.rglob('static.json')): + episode_dir = path.parent + if DiskEpisode(episode_dir).static.get(molmo_keys.EPISODE_INDEX) == episode_index: + return episode_dir + raise SystemExit(f'no recorded episode with {molmo_keys.EPISODE_INDEX}={episode_index} under {dataset_dir}') + + +def read_benchmark_path(dataset_dir: Path) -> str: + """The evaluated benchmark's path under the asset packs' ``benchmarks/`` root, from the run metadata. + + The path is kept from ``benchmarks/`` down — suite, scene dataset, task, benchmark — because the leaf + name alone is ambiguous: the same benchmark name exists under every scene dataset (ithor, + procthor-10k, ...) with different episodes, and replaying the wrong one silently replays a different + scene. Everything above ``benchmarks/`` is the box's own asset root and varies, so it is dropped. + """ + metadata = sorted(dataset_dir.glob('run_metadata_*.yaml')) + if not metadata: + raise SystemExit(f'no run_metadata_*.yaml in {dataset_dir} — cannot tell which benchmark was evaluated') + match = _BENCHMARK_ARG.search(metadata[-1].read_text()) + if match is None: + raise SystemExit(f'{metadata[-1]} records no --eval.benchmark_dir') + parts = Path(match.group(1)).parts + if mapping.ASSETS_BENCHMARKS_DIR not in parts: + raise SystemExit(f'evaluated benchmark {match.group(1)} is not under a benchmarks/ asset root') + return str(Path(*parts[parts.index(mapping.ASSETS_BENCHMARKS_DIR) + 1 :])) + + +def sample_at(signal: Signal, timestamps: list[int]) -> list: + """The signal's value at each timestamp — the last one at or before it, a pimm receiver's semantics.""" + sampled = signal.time[timestamps] + assert isinstance(sampled, Signal) # a sequence of timestamps samples a Signal, a single one a record + return [value for value, _ts in sampled] + + +def replay_commands( + benchmark_dir: Path, episode_index: int, commands: np.ndarray, grips: np.ndarray +) -> list[np.ndarray]: + """Step the commands open-loop through a MolmoSpaces env server, returning the sim state each produced. + + Stops early if the sim ends the trial, so a caller can tell a full replay from a truncated one by the + length of what comes back. + """ + states: list[np.ndarray] = [] + with launcher.serve_molmo_spaces(benchmark_dir) as (host, port): + conn = EnvConnection(host, port) + try: + # No seed: the benchmark episode carries its own, exactly as the recorded run left it unset. + conn.reset({mapping.TOKEN_EPISODE_INDEX: episode_index, mapping.TOKEN_SEED: None}) + for command, grip in zip(commands, grips, strict=True): + action = { + protocol.ACTION_COMMAND: { + protocol.COMMAND_TYPE: protocol.JOINT_POS, + protocol.COMMAND_JOINT_POS: command, + }, + protocol.ACTION_GRIP: float(grip), + } + out = conn.step(action) + states.append(np.asarray(out[protocol.FRAME_OBS][mapping.OBS_SIM_STATE], dtype=np.float64)) + if out[protocol.FRAME_DONE]: + break + finally: + conn.close() + return states + + +def build_fixture(episode_dir: Path, benchmark_path: str, assets_dir: Path) -> dict[str, np.ndarray]: + episode = DiskEpisode(episode_dir) + states = episode[mapping.OBS_SIM_STATE] + # rules-allow: hardcoded-keys — 'target_grip' is a canonical channel name spelled across every + # adoption and the eval configs; it belongs in positronic.keys, as its own sweep (internal#211). + commands, grips = episode[keys.TARGET_JOINTS], episode['target_grip'] + # Frame 0 is the reset observation; every later frame is one step. + frame_ts = [ts for _value, ts in states] + step_ts = frame_ts[1:] + played = [np.asarray(value, dtype=np.float32) for value in sample_at(commands, step_ts)] + grip = [float(np.asarray(value).reshape(-1)[0]) for value in sample_at(grips, step_ts)] + + # The recording's command signals stop before its observations do (internal#130), so only the steps up to + # and including the first one that reads the final recorded command are pinned by the recording; past that + # the commands the run actually applied were never written, and no substitute reproduces them. Replay that + # prefix and report the rest as the recording's gap rather than replaying commands it does not contain. + last_command_ts = commands[len(commands) - 1][1] + replayable = int(np.searchsorted(step_ts, last_command_ts, side='left')) + 1 + + steps = np.arange(1, replayable + 1) + checkpoints = np.unique(np.concatenate([steps[::CHECKPOINT_STRIDE], steps[-1:]])) + episode_index = int(episode.static[molmo_keys.EPISODE_INDEX]) + played_prefix = np.stack(played[:replayable]) + grip_prefix = np.array(grip[:replayable], dtype=np.float32) + benchmark_dir = assets_dir / mapping.ASSETS_BENCHMARKS_DIR / benchmark_path + replayed = replay_commands(benchmark_dir, episode_index, played_prefix, grip_prefix) + if len(replayed) != replayable: + raise SystemExit( + f'episode {episode_index}: the sim ended after {len(replayed)} of {replayable} replayable steps, ' + 'so the recording and the integration no longer agree on the trial length' + ) + return { + FIELD_EPISODE_INDEX: np.asarray(episode.static[molmo_keys.EPISODE_INDEX], dtype=np.int32), + FIELD_BENCHMARK_PATH: np.asarray(benchmark_path), + FIELD_TASK: np.asarray(episode.static[mapping.META_TASK]), + FIELD_COMMANDS: played_prefix, + FIELD_GRIPS: grip_prefix, + FIELD_UNREPLAYABLE_TAIL_STEPS: np.asarray(len(step_ts) - replayable, dtype=np.int32), + FIELD_CHECKPOINT_STEPS: checkpoints.astype(np.int32), + FIELD_CHECKPOINT_SIM_STATE: np.stack([replayed[int(step) - 1] for step in checkpoints]), + FIELD_EXPECTED_SUCCESS: np.asarray(episode.static[eval_keys.SUCCESS], dtype=bool), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description='Distil recorded eval episodes into replay fixtures.') + parser.add_argument('--dataset_dir', type=Path, required=True, help='recorded eval run (holds run_metadata)') + parser.add_argument( + '--episode_index', type=int, action='append', required=True, help='benchmark episode to distil; repeatable' + ) + args = parser.parse_args() + + assets = os.environ.get(mapping.ASSETS_DIR_ENV) + if not assets: + raise SystemExit( + f'{mapping.ASSETS_DIR_ENV} must point at the MolmoSpaces asset packs — the checkpoints ' + 'are taken by replaying the recorded commands, which needs the benchmark scene' + ) + + benchmark_path = read_benchmark_path(args.dataset_dir) + for episode_index in args.episode_index: + fixture = build_fixture(find_episode_dir(args.dataset_dir, episode_index), benchmark_path, Path(assets)) + if not fixture[FIELD_EXPECTED_SUCCESS]: + raise SystemExit(f'episode {episode_index} did not succeed — replay fixtures pin successful rollouts') + out = Path(__file__).parent / f'replay_ep{episode_index:02d}.npz' + np.savez_compressed(out, **fixture) # pyright: ignore[reportArgumentType] -- numpy's savez **kwds stub + steps = len(fixture[FIELD_COMMANDS]) + print(f'Wrote {out} ({out.stat().st_size} bytes, {steps} steps, {benchmark_path})') + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/tests/parity.py b/positronic/simulator/molmo_spaces/tests/parity.py new file mode 100644 index 000000000..e5f2b8f7f --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/parity.py @@ -0,0 +1,178 @@ +"""Native-vs-positronic parity check for the MolmoSpaces integration. + +The fidelity check ``docs/architecture.md`` ("Benchmarks are native; adoptions are faithful") mandates for every +sim-env integration, added here for MolmoSpaces. It drives one pinned benchmark episode twice — natively through +MolmoSpaces' own stack (``parity_native.py``: ``JsonEvalTaskSampler`` -> ``reset``/``step``/``is_done``/ +``judge_success``, MolmoSpaces' native horizon) and through the positronic path (launcher -> env server -> wire -> +the raw payload the ``MolmoAdapter`` maps) — feeding the *same* scripted actions (hold the arm, gripper open) and +asserts they agree byte-for-byte. + +MolmoSpaces benchmark episodes are exact-pose deterministic, so the strong fidelity form applies: identical call +sequence (same step count, same terminating step) and byte-identical outcomes modulo wire format (joint +positions/velocities, eef pose, gripper closure and success verdict equal at every step; camera frames equal by +content hash). Holding the arm never succeeds, so the episode runs out its horizon — exercising the horizon case +explicitly: both stacks terminate at exactly the native ``task_horizon`` step, via the wire ``done``, with +``success=False``. The two rollouts run in separate MuJoCo processes, so equality also confirms the render path is +deterministic across processes. + +This asserts fidelity against the pinned ``_MOLMO_COMMIT`` (``launcher.py``); re-run it on any bump of that pin +before merge — a sim version change can silently shift the horizon or the rollout. + +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``), and a benchmark whose task spec carries +``task_horizon_sec`` (the horizon the sim owns). 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.tests.parity --benchmark_dir +""" + +import argparse +import hashlib +import os +import subprocess +import tempfile +from pathlib import Path + +import numpy as np + +from positronic.simulator.env_server import protocol +from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.molmo_spaces import launcher, mapping +from positronic.simulator.molmo_spaces.tests import parity_record + +# parity_native.py runs only in MolmoSpaces' venv (it imports the flat, positronic-free ``env``), so reference it +# by path — importing it into positronic's interpreter would fail on that import. +_PARITY_NATIVE = Path(__file__).parent / 'parity_native.py' +_HOLD = {protocol.ACTION_COMMAND: {protocol.COMMAND_TYPE: protocol.HOLD}, protocol.ACTION_GRIP: 0.0} +_ARRAY_FIELDS = (mapping.OBS_JOINT_POS, mapping.OBS_JOINT_VEL, mapping.OBS_EEF_POS, mapping.OBS_EEF_QUAT) + + +def _drive_positronic(benchmark_dir: Path, episode_index: int, seed: int, max_steps: int) -> dict: + """Drive one episode through launcher -> env server -> wire, holding the arm to the sim's own ``done``.""" + fields: dict[str, list] = {k: [] for k in (*_ARRAY_FIELDS, mapping.OBS_GRIP)} + camera_names: list[str] = [] + cam_hashes: dict[str, list[str]] = {} + + def record(obs: dict) -> None: + for key in fields: + fields[key].append(obs[key]) + for name in camera_names: + cam_hashes[name].append(hashlib.sha256(np.ascontiguousarray(obs[name]).tobytes()).hexdigest()) + + with launcher.serve_molmo_spaces(benchmark_dir) as (host, port): + conn = EnvConnection(host, port) + try: + frame = conn.reset({mapping.TOKEN_EPISODE_INDEX: episode_index, mapping.TOKEN_SEED: seed}) + reported_horizon = frame[protocol.FRAME_HORIZON] + camera_names = [k for k, v in frame[protocol.FRAME_OBS].items() if mapping.is_rgb_frame(v)] + cam_hashes = {name: [] for name in camera_names} + record(frame[protocol.FRAME_OBS]) + out = {protocol.FRAME_DONE: False, protocol.FRAME_SUCCESS: False} + step = 0 + while not out[protocol.FRAME_DONE] and step < max_steps: + out = conn.step(_HOLD) + step += 1 + record(out[protocol.FRAME_OBS]) + finally: + conn.close() + return { + **{key: np.stack(fields[key]) for key in _ARRAY_FIELDS}, + mapping.OBS_GRIP: np.array(fields[mapping.OBS_GRIP], dtype=np.float32), + parity_record.CAMERA_NAMES: camera_names, + 'cam_hashes': cam_hashes, + 'reported_horizon': reported_horizon, + parity_record.TERMINATION_STEP: step, + parity_record.FINAL_SUCCESS: bool(out[protocol.FRAME_SUCCESS]), + } + + +def _native_env() -> dict[str, str]: + """The molmo-venv environment plus this directory, so the reference resolves ``parity_record`` flat. + + The launcher's PYTHONPATH carries what the *server* needs; ``parity_record`` is the comparison's own, so the + comparison adds it rather than the launcher knowing about a check. + """ + env = launcher.molmo_subprocess_env() + return {**env, 'PYTHONPATH': os.pathsep.join([env['PYTHONPATH'], str(Path(__file__).parent)])} + + +def _run_native(benchmark_dir: Path, episode_index: int, seed: int, max_steps: int, out_path: Path) -> dict: + """Drive the native reference (``parity_native.py``) in MolmoSpaces' venv and load its recorded rollout.""" + python = launcher.ensure_molmo_venv() + subprocess.run( + [ + str(python), + str(_PARITY_NATIVE), + parity_record.OPT_BENCHMARK_DIR, + str(benchmark_dir), + parity_record.OPT_EPISODE_INDEX, + str(episode_index), + parity_record.OPT_SEED, + str(seed), + parity_record.OPT_MAX_STEPS, + str(max_steps), + parity_record.OPT_OUT, + str(out_path), + ], + env=_native_env(), + check=True, + ) + return dict(np.load(out_path, allow_pickle=False)) + + +def _assert_parity(native: dict, positronic: dict, max_steps: int) -> None: + horizon = int(native[parity_record.HORIZON_STEPS]) + n_term = int(native[parity_record.TERMINATION_STEP]) + p_term = positronic[parity_record.TERMINATION_STEP] + assert n_term < max_steps, f'native never terminated in {max_steps} steps — raise --max_steps above the horizon' + assert p_term < max_steps, f'positronic never terminated in {max_steps} steps — raise --max_steps above the horizon' + # The horizon case: holding the arm never succeeds, so both stacks run out the native horizon and stop there. + assert n_term == horizon == p_term, ( + f'terminating step differs: native {n_term}, horizon {horizon}, positronic {p_term}' + ) + assert not bool(native[parity_record.FINAL_SUCCESS]) and not positronic[parity_record.FINAL_SUCCESS], ( + 'a held arm must not score success' + ) + # The env reports its horizon at reset (in sim-seconds); it must match native's and equal timeout's yardstick. + n_horizon = float(native[parity_record.HORIZON_SEC]) + assert n_horizon == positronic['reported_horizon'], ( + f'reported horizon differs: native {n_horizon}s, positronic {positronic["reported_horizon"]}s' + ) + + assert list(native[parity_record.CAMERA_NAMES]) == positronic[parity_record.CAMERA_NAMES], ( + 'camera sets differ between the stacks' + ) + for field in (*_ARRAY_FIELDS, mapping.OBS_GRIP): + n, p = native[field], positronic[field] + assert n.shape == p.shape, f'{field} shape differs: native {n.shape}, positronic {p.shape}' + assert np.array_equal(n, p), f'{field} differs between native and positronic rollouts' + for name in positronic[parity_record.CAMERA_NAMES]: + n_hashes, p_hashes = list(native[f'{parity_record.CAM_HASH_PREFIX}{name}']), positronic['cam_hashes'][name] + assert n_hashes == p_hashes, f'camera {name} frames differ between native and positronic rollouts' + + +def run(benchmark_dir: Path, *, episode_index: int = 0, seed: int = 0, max_steps: int = 1200) -> None: + """Run the same episode natively and through positronic and assert byte-identical parity.""" + with tempfile.TemporaryDirectory() as tmp: + native = _run_native(benchmark_dir, episode_index, seed, max_steps, Path(tmp) / 'native.npz') + positronic = _drive_positronic(benchmark_dir, episode_index, seed, max_steps) + _assert_parity(native, positronic, max_steps) + frames = positronic[parity_record.TERMINATION_STEP] + 1 + horizon = native[parity_record.HORIZON_STEPS] + print(f'PARITY PASSED — episode {episode_index}: {frames} frames, terminated at horizon {horizon}') + + +def main() -> None: + parser = argparse.ArgumentParser(description='Native-vs-positronic parity check for MolmoSpaces.') + parser.add_argument( + '--benchmark_dir', required=True, help='dir containing benchmark.json (task_horizon_sec required)' + ) + parser.add_argument('--episode_index', type=int, default=0) + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--max_steps', type=int, default=1200, help='safety cap; must exceed the benchmark horizon') + args = parser.parse_args() + run(Path(args.benchmark_dir), episode_index=args.episode_index, seed=args.seed, max_steps=args.max_steps) + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/tests/parity_native.py b/positronic/simulator/molmo_spaces/tests/parity_native.py new file mode 100644 index 000000000..05c4b7e7e --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/parity_native.py @@ -0,0 +1,144 @@ +"""Native-drive reference for the MolmoSpaces parity test — MolmoSpaces' own stack, no positronic. + +Runs in MolmoSpaces' venv, flat off ``PYTHONPATH`` like ``env.py`` (positronic-free: ``molmo_spaces`` + this +package's ``mapping``/``env`` modules). It drives one benchmark episode through MolmoSpaces' native rollout — +``JsonEvalTaskSampler`` -> ``reset``/``step``/``is_done``/``judge_success``, the sequence ``JsonEvalRunner`` runs — +holding the arm every step, and records the per-step raw sim state, per-camera frame hashes, and where the native +horizon terminates the episode. ``parity.py`` drives the *same* episode through the positronic env-server path and +asserts byte-identical outcomes against this reference. + +The reference derives what it compares from MolmoSpaces, not from the integration: the horizon from upstream's +own ``determine_task_horizon``, and each observation field read off the robot view here — the gripper closure +through upstream's own normalisation (``policy/learned_policy/pi_policy.py``). So an integration that resolves +the horizon or maps an observation differently from MolmoSpaces shows up as a parity failure rather than being +reproduced on both sides. + +Both rollouts build the task from ``env._DroidPickEvalConfig``: the comparison needs one task definition, +and MolmoSpaces has no per-benchmark config to derive a second from. + +Needs ``MLSPACES_ASSETS_DIR`` + a GL backend, like ``e2e.py``. Invoked by ``parity.py``; not run by hand. +""" + +# The ``molmo_spaces`` stack resolves only inside MolmoSpaces' own venv, where this reference runs; pyright +# checks it against positronic's deps, which cannot see it. Each such import carries its own +# ``reportMissingImports`` suppression, so one that should resolve here still fails the check. + +import argparse +import hashlib +from pathlib import Path + +# env.py (imported flat off PYTHONPATH, like mapping/server) sets MUJOCO_GL and installs the CGL stub at import, +# GL-safely pulling in the molmo_spaces stack — so import it before any other molmo_spaces import. +import env # noqa: E402 +import mapping # noqa: E402 -- positronic-free wire mappings, on PYTHONPATH +import mujoco # noqa: E402 +import numpy as np # noqa: E402 +import parity_record # noqa: E402 -- the record's field names, on PYTHONPATH beside this file + +from molmo_spaces.evaluation.benchmark_schema import ( # noqa: E402 # pyright: ignore[reportMissingImports] + load_all_episodes, +) +from molmo_spaces.evaluation.eval_main import ( # noqa: E402 # pyright: ignore[reportMissingImports] + determine_task_horizon, +) +from molmo_spaces.tasks.json_eval_task_sampler import ( # noqa: E402 # pyright: ignore[reportMissingImports] + JsonEvalTaskSampler, +) + +# The Robotiq finger qpos the DROID observation's closure is normalised against, as MolmoSpaces' own policies +# read it (``np.clip(obs["qpos"]["gripper"][0] / 0.824033, 0, 1)``, pi_policy.py:126). Transcribed from upstream +# rather than read from ``mapping``, so a wrong value there is what this reference catches instead of sharing. +_GRIPPER_QPOS_CLOSED = 0.824033 + + +def _observe(robot_view, env_obs: dict, camera_names: list[str]) -> dict: + """One frame's compared values, read off MolmoSpaces directly: measured joints, the grasp-site world pose, + the gripper closure, and each camera's frame.""" + arm = robot_view.get_move_group(mapping.MOLMO_ARM_GROUP) + eef_world = np.asarray(arm.leaf_frame_to_world, dtype=np.float64) + quat = np.zeros(4) # wxyz + mujoco.mju_mat2Quat(quat, np.ascontiguousarray(eef_world[:3, :3].reshape(9))) # pyright: ignore[reportAttributeAccessIssue] + qpos = env_obs[mapping.MOLMO_OBS_QPOS][mapping.MOLMO_GRIPPER_GROUP] + grip = np.clip(qpos[0] / _GRIPPER_QPOS_CLOSED, 0.0, 1.0) + return { + mapping.OBS_JOINT_POS: np.asarray(arm.joint_pos, dtype=np.float32), + mapping.OBS_JOINT_VEL: np.asarray(arm.joint_vel, dtype=np.float32), + mapping.OBS_EEF_POS: eef_world[:3, 3].astype(np.float32), + mapping.OBS_EEF_QUAT: quat.astype(np.float32), + mapping.OBS_GRIP: np.float32(grip), + **{name: np.ascontiguousarray(env_obs[name]) for name in camera_names}, + } + + +def _run(benchmark_dir: Path, episode_index: int, seed: int, max_steps: int, out_path: Path) -> None: + episodes = load_all_episodes(benchmark_dir) + episode = episodes[episode_index] + cfg = env._DroidPickEvalConfig() + cfg.seed = seed + native_horizon = determine_task_horizon([episode], None, cfg.policy_dt_ms) + cfg.task_horizon = native_horizon + sampler = JsonEvalTaskSampler(cfg, episode) + task = sampler.sample_task(house_index=episode.house_index) + robot_view = task.env.current_robot.robot_view + + obs, _info = task.reset() + camera_names = [k for k, v in obs[0].items() if mapping.is_rgb_frame(v)] + fields: dict[str, list] = { + k: [] + for k in ( + mapping.OBS_JOINT_POS, + mapping.OBS_JOINT_VEL, + mapping.OBS_EEF_POS, + mapping.OBS_EEF_QUAT, + mapping.OBS_GRIP, + ) + } + cam_hashes: dict[str, list[str]] = {name: [] for name in camera_names} + + def record(env_obs: dict) -> None: + payload = _observe(robot_view, env_obs, camera_names) + for key in fields: + fields[key].append(payload[key]) + for name in camera_names: + cam_hashes[name].append(hashlib.sha256(payload[name].tobytes()).hexdigest()) + + record(obs[0]) + step, success = 0, False + # The native rollout: hold the measured joints (the gripper open) and let the sim run until its own is_done — + # is_terminal or horizon expiry. A hold never succeeds, so this drives the horizon case. max_steps bounds a sim + # that never terminates (a wrong horizon); the caller asserts termination lands below it. + while not bool(task.is_done()) and step < max_steps: + measured_q = np.asarray(robot_view.get_move_group(mapping.MOLMO_ARM_GROUP).joint_pos, dtype=np.float32) + action = {mapping.MOLMO_ARM_GROUP: measured_q, mapping.MOLMO_GRIPPER_GROUP: np.array([0.0], dtype=np.float32)} + obs, _reward, _term, _trunc, _infos = task.step(action) + step += 1 + record(obs[0]) + success = bool(task.judge_success()) + if success: # end-on-success, matching env.py's step (a hold never reaches it) + break + sampler.close() + + recorded: dict = {key: np.stack(values) for key, values in fields.items()} + recorded.update({f'{parity_record.CAM_HASH_PREFIX}{name}': np.array(cam_hashes[name]) for name in camera_names}) + recorded[parity_record.CAMERA_NAMES] = np.array(camera_names) + recorded[parity_record.HORIZON_STEPS] = native_horizon + recorded[parity_record.HORIZON_SEC] = native_horizon * (cfg.policy_dt_ms / 1000.0) # env.py reports this at reset + recorded[parity_record.TERMINATION_STEP] = step + recorded[parity_record.FINAL_SUCCESS] = success + # numpy's savez **kwds stub reads a dict-unpack as possibly supplying ``allow_pickle`` (as in make_fixture.py). + np.savez(out_path, **recorded) # pyright: ignore[reportArgumentType] + + +def main() -> None: + parser = argparse.ArgumentParser(description='Native-drive MolmoSpaces reference for the parity test.') + parser.add_argument(parity_record.OPT_BENCHMARK_DIR, required=True) + parser.add_argument(parity_record.OPT_EPISODE_INDEX, type=int, default=0) + parser.add_argument(parity_record.OPT_SEED, type=int, required=True) + parser.add_argument(parity_record.OPT_MAX_STEPS, type=int, required=True) + parser.add_argument(parity_record.OPT_OUT, required=True, help='npz path for the recorded native rollout') + args = parser.parse_args() + _run(Path(args.benchmark_dir), args.episode_index, args.seed, args.max_steps, Path(args.out)) + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/tests/parity_record.py b/positronic/simulator/molmo_spaces/tests/parity_record.py new file mode 100644 index 000000000..cad4e5570 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/parity_record.py @@ -0,0 +1,23 @@ +"""The interface between ``parity_native.py`` and ``parity.py``: the options one is spawned with, and the npz +field names it writes about a rollout for the other to read back. + +Imported from two interpreters, like ``mapping``: as a package module by the comparison, and flat off +``PYTHONPATH`` by the native reference inside MolmoSpaces' venv. It holds names only — no imports at all — so +both shapes resolve without a fallback. + +The per-camera frame hashes are one field per camera name, under ``CAM_HASH_PREFIX``. +""" + +# The native reference's CLI: ``parity.py`` builds the command, ``parity_native.py``'s parser declares it. +OPT_BENCHMARK_DIR = '--benchmark_dir' +OPT_EPISODE_INDEX = '--episode_index' +OPT_SEED = '--seed' +OPT_MAX_STEPS = '--max_steps' +OPT_OUT = '--out' + +CAM_HASH_PREFIX = 'cam_hash__' +CAMERA_NAMES = 'camera_names' +HORIZON_STEPS = 'native_horizon' +HORIZON_SEC = 'horizon_sec' +TERMINATION_STEP = 'termination_step' +FINAL_SUCCESS = 'final_success' diff --git a/positronic/simulator/molmo_spaces/tests/replay_ep03.npz b/positronic/simulator/molmo_spaces/tests/replay_ep03.npz new file mode 100644 index 000000000..3b0629cfe Binary files /dev/null and b/positronic/simulator/molmo_spaces/tests/replay_ep03.npz differ diff --git a/positronic/simulator/molmo_spaces/tests/replay_ep06.npz b/positronic/simulator/molmo_spaces/tests/replay_ep06.npz new file mode 100644 index 000000000..174827ed5 Binary files /dev/null and b/positronic/simulator/molmo_spaces/tests/replay_ep06.npz differ 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..685723a36 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_adapter.py @@ -0,0 +1,97 @@ +"""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 import keys +from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol +from positronic.simulator.molmo_spaces import keys as molmo_keys +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.molmo_spaces.adapter import DEFAULT_CAMERA_DICT as CAMERA_DICT +from positronic.simulator.molmo_spaces.adapter import MolmoAdapter + +FIXTURE = Path(__file__).parent / 'droid_obs.npz' + + +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[keys.ROBOT_STATE] + assert np.allclose(state.q, payload[mapping.OBS_JOINT_POS]) + assert np.allclose(state.dq, payload[mapping.OBS_JOINT_VEL]) + assert np.allclose(state.ee_pose.translation, payload[mapping.OBS_EEF_POS]) + assert np.allclose(state.ee_pose.rotation.as_quat, payload[mapping.OBS_EEF_QUAT]) # wxyz round-trips + assert obs[keys.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[keys.WRIST_IMAGE].array, payload[mapping.MOLMO_WRIST_CAMERA]) + assert np.array_equal(obs[keys.EXTERIOR_IMAGE].array, payload[mapping.MOLMO_EXTERIOR_CAMERA]) + # Fixture marks wrist reddish, exterior greenish; a swap would flip the dominant channel. + wrist_mean = obs[keys.WRIST_IMAGE].array.reshape(-1, 3).mean(axis=0) + exterior_mean = obs[keys.EXTERIOR_IMAGE].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. + payload = _payload() + payload[mapping.MOLMO_WRIST_CAMERA_VARIANTS[0]] = payload.pop(mapping.MOLMO_WRIST_CAMERA) + obs = MolmoAdapter(CAMERA_DICT).observations(payload) + wrist_mean = obs[keys.WRIST_IMAGE].array.reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] + + +def test_privileged_forwards_sim_state(): + # The full MuJoCo state is recorded as privileged ground truth (never fed to the policy), so success can be + # recomputed offline. + state = np.arange(10, dtype=np.float64) + out = MolmoAdapter(CAMERA_DICT).privileged({mapping.OBS_SIM_STATE: state}) + assert list(out) == [mapping.OBS_SIM_STATE] and out[mapping.OBS_SIM_STATE] is state + + +def test_terminal_reports_success_only_when_done(): + adapter = MolmoAdapter(CAMERA_DICT) + done_ok = {protocol.FRAME_DONE: True, protocol.FRAME_SUCCESS: True} + done_fail = {protocol.FRAME_DONE: True, protocol.FRAME_SUCCESS: False} + running = {protocol.FRAME_DONE: False, protocol.FRAME_SUCCESS: False} + assert adapter.terminal(done_ok) == {eval_keys.SUCCESS: True} + assert adapter.terminal(done_fail) == {eval_keys.SUCCESS: False} + assert adapter.terminal(running) is None + + +def test_task_params_name_an_episode_the_way_the_reset_token_reads_it(): + adapter = MolmoAdapter(CAMERA_DICT) + params = adapter.task_params([{'name': 'put the banana in the bowl', 'episode_index': 3, 'task_horizon_sec': 30.0}]) + assert params == [ + {eval_keys.TASK: 'put the banana in the bowl', molmo_keys.EPISODE_INDEX: 3, molmo_keys.TASK_HORIZON: 30.0} + ] + + +def test_reset_token_carries_episode_and_seed(): + adapter = MolmoAdapter(CAMERA_DICT) + expected = {mapping.TOKEN_EPISODE_INDEX: 3, mapping.TOKEN_SEED: 7} + assert adapter.reset_token({molmo_keys.EPISODE_INDEX: 3, eval_keys.SEED: 7}) == expected + # An absent seed falls back to the spec's own (None here). + assert adapter.reset_token({molmo_keys.EPISODE_INDEX: 2}) == { + mapping.TOKEN_EPISODE_INDEX: 2, + mapping.TOKEN_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..6d133a030 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_mapping.py @@ -0,0 +1,258 @@ +"""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 types + +import numpy as np +import pytest + +from positronic.simulator.env_server import protocol +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( + {protocol.COMMAND_TYPE: protocol.JOINT_POS, protocol.COMMAND_JOINT_POS: 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( + {protocol.COMMAND_TYPE: protocol.JOINT_VEL, protocol.COMMAND_JOINT_VEL: 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({protocol.COMMAND_TYPE: protocol.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( + {protocol.COMMAND_TYPE: protocol.JOINT_VEL, protocol.COMMAND_JOINT_VEL: 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( + {protocol.COMMAND_TYPE: protocol.CARTESIAN, protocol.COMMAND_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, 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, variants) == variants[0] + # Variant only (default absent) -> the variant. + assert mapping.resolve_camera_key({variants[0]: 1}, 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') == 'my_cam' + + +def test_camera_key_miss_raises(): + with pytest.raises(KeyError): + mapping.resolve_camera_key({'other': 1}, mapping.MOLMO_WRIST_CAMERA) + + +def _episodes(*horizons_sec): + return [ + types.SimpleNamespace(task={} if sec is None else {mapping.MOLMO_TASK_HORIZON_SEC: sec}) for sec in horizons_sec + ] + + +def test_task_horizon_reads_the_benchmark_task_dict(): + # Where MolmoSpaces' benchmark generator writes it, and where determine_task_horizon reads it. + assert mapping.resolve_task_horizon_steps(_episodes(30, 30), 66.0) == 455 # round(30 * 1000 / 66) + + +def test_task_horizon_missing_raises(): + # The raise mirrors upstream's resolver. + with pytest.raises(ValueError): + mapping.resolve_task_horizon_steps(_episodes(30, None), 66.0) + + +def test_task_horizon_disagreeing_across_episodes_raises(): + # The horizon belongs to the benchmark, so one run has one; upstream refuses the same manifest. + with pytest.raises(ValueError, match='inconsistent'): + mapping.resolve_task_horizon_steps(_episodes(20, 30), 66.0) + + +def test_task_horizon_non_positive_raises(): + # A zero or negative span is no horizon at all; the override path already refuses one below a step. + with pytest.raises(ValueError, match='non-positive'): + mapping.resolve_task_horizon_steps(_episodes(0), 66.0) + with pytest.raises(ValueError, match='non-positive'): + mapping.resolve_task_horizon_steps(_episodes(-5), 66.0) + + +def test_task_horizon_rounding_below_one_step_raises(): + # 0.03s of a 66ms period rounds to 0 steps, which would expire the episode before its first action. + with pytest.raises(ValueError, match='rounds to 0 steps'): + mapping.resolve_task_horizon_steps(_episodes(0.03), 66.0) + + +def test_task_horizon_override_wins(): + # An explicit override pins the horizon, beating the benchmark field (mirrors --task_horizon_steps), and lets + # a benchmark that declares none, or disagrees, still resolve. + assert mapping.resolve_task_horizon_steps(_episodes(20), 66.0, override_steps=500) == 500 + assert mapping.resolve_task_horizon_steps(_episodes(20, 30, None), 66.0, override_steps=455) == 455 + + +def test_exterior_camera_variants_cover_light_randomization_and_randcam(): + # The default exterior mapping must resolve both benchmark exterior names: RandCam records + # randomized_zed2_analogue_1, not exo_camera_1. + 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, variants) == name + + +def test_unpack_wire_pose_round_trips_translation_and_rotation(): + # The client encodes a pose as Transform3D.as_vector(ROTATION_MATRIX): translation, then R row-major. + rot = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) # +90 deg about z + pos, out = mapping.unpack_wire_pose(np.concatenate([[1.0, 2.0, 3.0], rot.reshape(-1)])) + assert np.array_equal(pos, [1.0, 2.0, 3.0]) + assert np.array_equal(out, rot) + + +def test_unpack_wire_pose_rejects_wrong_width(): + with pytest.raises(ValueError): + mapping.unpack_wire_pose(np.zeros(7)) # a quaternion-encoded pose is not the wire form + + +def test_compose_world_delta_adds_translation_and_left_multiplies_rotation(): + # World-frame convention: goal_pos = ee_pos + dpos, goal_ori = R(delta) @ ee_ori. Left-multiplication is + # what keeps the delta world-framed; composing in the body frame would rotate the translation too. + cur_rot = np.eye(3) + delta_rot = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + pos, rot = mapping.compose_world_delta([1.0, 0.0, 0.0], cur_rot, [0.0, 2.0, 0.0], delta_rot) + assert np.allclose(pos, [1.0, 2.0, 0.0]) + assert np.allclose(rot, delta_rot) + + +def test_cartesian_command_resolves_through_the_supplied_ik(): + # env.py owns the solver (it needs the live model); mapping only routes the target into it. + solved = np.arange(mapping.NUM_ARM_JOINTS, dtype=np.float64) + seen = {} + + def ik(pos, rot): + seen['pos'], seen['rot'] = pos, rot + return solved + + rot = np.eye(3) + cmd = { + protocol.COMMAND_TYPE: protocol.CARTESIAN, + protocol.COMMAND_POSE: np.concatenate([[0.4, 0.1, 0.3], rot.reshape(-1)]), + } + out = mapping.wire_command_to_arm_action(cmd, np.zeros(mapping.NUM_ARM_JOINTS), ik=ik) + assert out.dtype == np.float32 and np.allclose(out, solved) + assert np.allclose(seen['pos'], [0.4, 0.1, 0.3]) and np.allclose(seen['rot'], rot) + + +def test_cartesian_delta_composes_onto_the_measured_eef_before_solving(): + # The delta is relative to the *measured* pose, so the solver must see the composed absolute target. + seen = {} + + def ik(pos, rot): + seen['pos'], seen['rot'] = pos, rot + return np.zeros(mapping.NUM_ARM_JOINTS) + + cmd = { + protocol.COMMAND_TYPE: protocol.CARTESIAN_DELTA, + protocol.COMMAND_DELTA: np.concatenate([[0.0, 0.1, 0.0], np.eye(3).reshape(-1)]), + } + mapping.wire_command_to_arm_action( + cmd, np.zeros(mapping.NUM_ARM_JOINTS), ik=ik, current_eef=(np.array([0.5, 0.0, 0.2]), np.eye(3)) + ) + assert np.allclose(seen['pos'], [0.5, 0.1, 0.2]) + + +def test_cartesian_without_an_ik_solver_raises(): + # A caller that holds no model cannot resolve a Cartesian target — fail loud rather than silently holding. + cmd = { + protocol.COMMAND_TYPE: protocol.CARTESIAN, + protocol.COMMAND_POSE: np.concatenate([np.zeros(3), np.eye(3).reshape(-1)]), + } + with pytest.raises(ValueError, match='ik solver'): + mapping.wire_command_to_arm_action(cmd, np.zeros(mapping.NUM_ARM_JOINTS)) + + +def test_unknown_command_names_the_canonical_contract(): + with pytest.raises(ValueError, match='cartesian'): # the message lists the contract the tag is not part of + mapping.wire_command_to_arm_action({protocol.COMMAND_TYPE: 'wrench'}, np.zeros(mapping.NUM_ARM_JOINTS)) + + +@pytest.mark.parametrize('command_type', protocol.CANONICAL_COMMAND_TYPES) +def test_every_canonical_command_type_converts_to_joint_targets(command_type): + """The contract is total, so every canonical type converts to the joint targets MolmoSpaces natively steps. + This is the model-free half of that property — the routing, through a stub solver; ``validate.py`` drives + the same types through the real IK against a live scene.""" + pose = np.concatenate([np.zeros(3), np.eye(3).reshape(-1)]) + payload = { + protocol.JOINT_POS: {protocol.COMMAND_JOINT_POS: np.zeros(mapping.NUM_ARM_JOINTS)}, + protocol.JOINT_VEL: {protocol.COMMAND_JOINT_VEL: np.zeros(mapping.NUM_ARM_JOINTS)}, + protocol.HOLD: {}, + protocol.CARTESIAN: {protocol.COMMAND_POSE: pose}, + protocol.CARTESIAN_DELTA: {protocol.COMMAND_DELTA: pose}, + }[command_type] + + target = mapping.wire_command_to_arm_action( + {protocol.COMMAND_TYPE: command_type, **payload}, + np.zeros(mapping.NUM_ARM_JOINTS), + ik=lambda _pos, _rot: np.zeros(mapping.NUM_ARM_JOINTS), + current_eef=(np.zeros(3), np.eye(3)), + ) + + assert target.shape == (mapping.NUM_ARM_JOINTS,) + + +def test_episode_seed_prefers_the_override_then_the_spec_then_the_index(): + spec = types.SimpleNamespace(seed=7) + assert mapping.resolve_episode_seed(spec, 3, 99) == 99 + assert mapping.resolve_episode_seed(spec, 3) == 7 + assert mapping.resolve_episode_seed(types.SimpleNamespace(seed=None), 3) == 3 + assert mapping.resolve_episode_seed(types.SimpleNamespace(), 5) == 5 diff --git a/positronic/simulator/molmo_spaces/tests/test_replay.py b/positronic/simulator/molmo_spaces/tests/test_replay.py new file mode 100644 index 000000000..a7f849a43 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_replay.py @@ -0,0 +1,91 @@ +"""Deterministic replay of recorded pi05 rollouts against the MolmoSpaces integration. + +The parity check (``parity.py``) pins that positronic drives the sim exactly as MolmoSpaces' own stack does, +but it drives a scripted hold — no real policy behaviour ever reaches the sim. This replays the commands a +real pi05 rollout emitted, open-loop from the benchmark's own seed, and asserts the sim retraces the pinned +trajectory step for step. + +It reproduces because the recorded commands are absolute joint targets, so no observation feeds back into the +stream: the replay depends only on the sim rollout and the env-server path between positronic and it. MuJoCo +CPU physics is deterministic for a pinned build, and the render nondeterminism ``exp-004`` found never reaches +an open-loop replay, which reads no images. Same-host replay of the pinned ``_MOLMO_COMMIT`` reproduces the +fixture's ``sim_state`` *exactly* — zero deviation across every replayed step of both fixtures — so a +divergence here is a real change in the integration (command mapping, horizon, the wire, scene selection), +not sim noise. + +The rollout stops where the recording stops pinning it: an eval episode's command signals end before its +observations do, so the last few steps of every recording apply commands that were never written down +(internal#130). Those steps are excluded, which is why this asserts trajectory fidelity and not the recorded +``eval.success`` — success lands inside that unwritten tail. Closing internal#130 is what would let the +verdict itself be replayed. + +Fixtures (``replay_ep*.npz``, from ``make_replay_fixture.py``) hold only the commands, the grip and +``sim_state`` checkpoints — never the videos. The commands come from the recording; the checkpoints are taken +by replaying them, so what this pins is the integration's own trajectory under a real policy's commands, and +they are regenerated together. The benchmark is a multi-hundred-MB asset pack that cannot be committed, so +the test skips unless this box has it. + +Run on a box with the asset packs (a GPU-less one uses mesa software EGL):: + + MLSPACES_ASSETS_DIR=... MUJOCO_GL=egl EGL_PLATFORM=surfaceless LIBGL_ALWAYS_SOFTWARE=1 \ + uv run --locked pytest positronic/simulator/molmo_spaces/tests/test_replay.py --no-cov +""" + +import os +from pathlib import Path + +import numpy as np +import pytest + +from positronic.simulator.molmo_spaces import mapping +from positronic.simulator.molmo_spaces.tests import make_replay_fixture as fixture_fields +from positronic.simulator.molmo_spaces.tests.make_replay_fixture import replay_commands + +FIXTURES = sorted(Path(__file__).parent.glob('replay_ep*.npz')) + +# Replay on the recording's own host reproduces every checkpoint bit-for-bit, so this budget is for the float +# drift a different CPU can introduce. It sits orders of magnitude below the ~centimetre scale that decides a +# pick, so it cannot mask a real regression; a cross-host run that exceeds it is worth investigating rather +# than widening. +SIM_STATE_TOL = 1e-6 + + +def _benchmark_dir(benchmark_path: str) -> Path: + """The benchmark the fixture was recorded against, resolved in this box's asset packs. + + The fixture pins the path from ``benchmarks/`` down, so the lookup is exact rather than a name search: + the same benchmark name sits under every scene dataset with different episodes, and resolving to the + wrong one would replay these commands against a different scene. + """ + assets = os.environ.get(mapping.ASSETS_DIR_ENV) + if not assets: + pytest.skip(f'{mapping.ASSETS_DIR_ENV} is unset — MolmoSpaces asset packs are needed to replay') + benchmark_dir = Path(assets) / mapping.ASSETS_BENCHMARKS_DIR / benchmark_path + if not (benchmark_dir / mapping.MOLMO_BENCHMARK_MANIFEST).is_file(): + pytest.skip(f'{benchmark_dir} is absent — this asset pack cannot replay the fixture') + return benchmark_dir + + +@pytest.mark.parametrize('fixture_path', FIXTURES, ids=lambda path: path.stem) +def test_recorded_rollout_replays_the_pinned_trajectory(fixture_path: Path): + fixture = np.load(fixture_path, allow_pickle=False) + commands, grips = fixture[fixture_fields.FIELD_COMMANDS], fixture[fixture_fields.FIELD_GRIPS] + episode_index = int(fixture[fixture_fields.FIELD_EPISODE_INDEX]) + benchmark_dir = _benchmark_dir(str(fixture[fixture_fields.FIELD_BENCHMARK_PATH])) + states = replay_commands(benchmark_dir, episode_index, commands, grips) + + # The sim must not end the trial inside the replayed prefix: the recording ran every one of these steps, + # so an early terminal means the integration now scores or expires the episode differently. + assert len(states) == len(commands), ( + f'replay of episode {episode_index} terminated after {len(states)} of {len(commands)} recorded steps ' + f'(its {int(fixture[fixture_fields.FIELD_UNREPLAYABLE_TAIL_STEPS])} unrecorded tail steps are excluded)' + ) + + # Checkpoints along the way, not just the end state: drift shows up long before it would flip a verdict. + checkpoint_steps = fixture[fixture_fields.FIELD_CHECKPOINT_STEPS] + for step, pinned in zip(checkpoint_steps, fixture[fixture_fields.FIELD_CHECKPOINT_SIM_STATE], strict=True): + replayed = states[int(step) - 1] # states[i] holds step i + 1; the fixture indexes steps from 1 + deviation = float(np.max(np.abs(replayed - pinned))) + assert deviation <= SIM_STATE_TOL, ( + f'sim_state diverged by {deviation:.3e} at step {step} of episode {episode_index}' + ) diff --git a/positronic/simulator/molmo_spaces/tests/test_tasks.py b/positronic/simulator/molmo_spaces/tests/test_tasks.py new file mode 100644 index 000000000..94e3a2926 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_tasks.py @@ -0,0 +1,95 @@ +"""Which episodes a MolmoSpaces eval runs: the adapter maps the env's records, the config sweeps them. + +``RemoteEnvControlSystem.tasks`` is stubbed, since a real episode list needs the MolmoSpaces venv; +``tests/e2e.py`` runs the same command against the real benchmark. +""" + +import logging +from typing import Any + +import pytest + +from positronic.cfg.eval.sim.molmo import _TIMEOUT_MARGIN_SEC, benchmark +from positronic.eval import keys as eval_keys +from positronic.simulator.env_server.proxy import RemoteEnvControlSystem +from positronic.simulator.molmo_spaces import keys as molmo_keys + +_HORIZON_SEC = 30.0 + + +@pytest.fixture +def asked(monkeypatch) -> list[Any]: + """The specs the eval sends the proxy; every spec is answered with two episodes of one horizon.""" + specs: list[Any] = [] + + def tasks(self, selection: Any) -> list[dict[str, Any]]: + specs.append(selection) + return [ + { + eval_keys.TASK: 'put the banana in the bowl', + molmo_keys.EPISODE_INDEX: 0, + molmo_keys.TASK_HORIZON: _HORIZON_SEC, + }, + { + eval_keys.TASK: 'put the mug on the shelf', + molmo_keys.EPISODE_INDEX: 1, + molmo_keys.TASK_HORIZON: _HORIZON_SEC, + }, + ] + + monkeypatch.setattr(RemoteEnvControlSystem, 'tasks', tasks) + return specs + + +def test_the_env_answers_which_episodes_the_sweep_runs(asked): + """The sweep is asked for when the run starts; an unset seed leaves the episode's own seed in force.""" + ev = benchmark.override(benchmark_dir='unused', trial_count=2).instantiate() + assert asked == [] + + trials = list(ev.tasks()) + + assert asked == [{}] + scenes = [trial.prepare_args[eval_keys.SCENE] for trial in trials] + assert [scene[molmo_keys.EPISODE_INDEX] for scene in scenes] == [0, 0, 1, 1] + assert all(eval_keys.SEED not in scene for scene in scenes) + assert [trial.meta[eval_keys.TRIAL_INDEX] for trial in trials] == [0, 1, 2, 3] + + +def test_an_episode_selection_rides_the_spec(asked): + ev = benchmark.override(benchmark_dir='unused', episodes=[0, 1]).instantiate() + ev.tasks() + assert asked == [{'episodes': [0, 1]}] + + +def test_an_explicit_seed_sweeps_each_episode(asked): + ev = benchmark.override(benchmark_dir='unused', seed=5, trial_count=2).instantiate() + scenes = [trial.prepare_args[eval_keys.SCENE] for trial in ev.tasks()] + assert [scene[eval_keys.SEED] for scene in scenes] == [5, 6, 5, 6] + + +def test_a_non_positive_trial_count_is_refused(): + with pytest.raises(ValueError, match='trial_count'): + benchmark.override(benchmark_dir='unused', trial_count=0).instantiate() + + +def test_timeout_defaults_to_the_benchmark_horizon_plus_a_margin(asked): + ev = benchmark.override(benchmark_dir='unused').instantiate() + trials = list(ev.tasks()) + assert [trial.timeout_sec for trial in trials] == [_HORIZON_SEC + _TIMEOUT_MARGIN_SEC] * 2 + + +def test_explicit_timeout_can_only_lower_the_deadline(asked, caplog): + with caplog.at_level(logging.WARNING): + short = benchmark.override(benchmark_dir='unused', timeout=20.0).instantiate().tasks() + long = benchmark.override(benchmark_dir='unused', timeout=999.0).instantiate().tasks() + assert short[0].timeout_sec == 20.0 + assert long[0].timeout_sec == _HORIZON_SEC + _TIMEOUT_MARGIN_SEC + assert len(caplog.records) == 2, 'both directions differ from the backstop, so both warn' + + +def test_timeout_matching_the_backstop_is_silent(asked, caplog): + backstop = _HORIZON_SEC + _TIMEOUT_MARGIN_SEC + with caplog.at_level(logging.WARNING): + trials = benchmark.override(benchmark_dir='unused', timeout=backstop).instantiate().tasks() + assert trials[0].timeout_sec == backstop + assert not caplog.records diff --git a/positronic/simulator/molmo_spaces/tests/validate.py b/positronic/simulator/molmo_spaces/tests/validate.py new file mode 100644 index 000000000..6cd430d25 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/validate.py @@ -0,0 +1,160 @@ +"""Validate the MolmoSpaces rig's command transforms against the live sim. + +MolmoSpaces' Franka runs a joint-position controller, so every other command reaches it only through the +conversions in ``mapping``, the Cartesian pair among them through the differential IK in ``env.py``. That +solver is arithmetic over the live MuJoCo model, which no unit test can reach (``mapping``'s tests cover the +routing with a stub solver, not the kinematics), so it is checked here against a real benchmark scene — the +same shape of check ``simulator/libero/validate.py`` runs for the LIBERO rig. + +Four properties. The kinematic three read the arm's grasp site (the frame the env observes in, so command and +observation share a frame); the fourth is the adoption's coverage of the command contract: + +- **FK identity** — the scratch-``MjData`` recompute of the measured joints reproduces the live grasp-site read, + confirming the scratch evaluation is seeded correctly and reads the same frame. +- **IK round-trip** — for reachable targets sampled by perturbing the measured joints, ``_fk(_ik(pose))`` + recovers the pose. This is the property a Cartesian policy depends on; the sampling stays near the measured + configuration so every target is reachable and the check tests the solver, not the workspace. +- **Cartesian hold** — commanding the pose the arm already holds resolves to the joints it already holds, which + is what makes an absolute Cartesian setpoint stable when a policy re-sends it. +- **Command contract** — every canonical command type converts to joint targets through the live IK. The + contract is total, so this is where the adoption's coverage of it is verified rather than asserted. + +Runs in MolmoSpaces' venv, flat off ``PYTHONPATH`` like ``parity_native.py`` (positronic-free: ``molmo_spaces`` +plus this package's ``mapping``/``env``), so positronic's interpreter cannot import it. Needs the asset packs +(``MLSPACES_ASSETS_DIR``) and a GL backend (``MUJOCO_GL``; a GPU-less box uses mesa software EGL). Launch it the +way ``parity.py`` launches the native reference — the venv python under ``launcher.molmo_subprocess_env()``:: + + uv run --locked python -c " + import subprocess + from positronic.simulator.molmo_spaces import launcher + subprocess.run([str(launcher.ensure_molmo_venv()), + 'positronic/simulator/molmo_spaces/tests/validate.py', '--benchmark_dir', ''], + env=launcher.molmo_subprocess_env(), check=True)" +""" + +# The flat ``protocol`` module resolves only inside MolmoSpaces' own venv, where this validation runs; pyright +# checks it against positronic's deps, which cannot see it. That import carries its own +# ``reportMissingImports`` suppression, so one that should resolve here still fails the check. + +import argparse +from pathlib import Path + +# env.py sets MUJOCO_GL and installs the CGL stub at import, GL-safely pulling in the molmo_spaces stack — so +# import it before any other molmo_spaces import. Reaching into its private ``_fk``/``_ik`` is the point: this +# validates that exact solver, not a re-derivation of it. +import env # noqa: E402 +import mapping # noqa: E402 -- positronic-free wire mappings, on PYTHONPATH +import numpy as np +import protocol # pyright: ignore[reportMissingImports] -- flat on PYTHONPATH beside ``server``, see ``launcher`` + +# Sampled targets perturb each measured joint by up to this much (radians): far enough that the solver has real +# work to do, near enough that every target stays reachable and away from the limits. +_JOINT_JITTER = 0.1 +_IK_SAMPLES = 16 +# The solver iterates to _IK_TOL on the 6-vector error; these are the per-component budgets that implies. +_POS_ATOL = 1e-3 # metres +_ORI_ATOL = 1e-2 # radians +# The live site is read after the sim has stepped, so it carries residual motion the scratch recompute of the +# same joints cannot reproduce exactly; float precision, not float64, is the right bar for the identity. +_FK_ATOL = 1e-5 +# The step the relative commands carry: small enough that the target stays reachable from the measured +# configuration, large enough that the conversion is not the identity. +_DELTA_POS = 0.01 # metres +_DELTA_Q = 0.01 # radians + + +def _ori_error(target_rot: np.ndarray, rot: np.ndarray) -> float: + return float(np.linalg.norm(env._pose_error(np.zeros(3), target_rot, np.zeros(3), rot)[3:])) + + +def _check_fk_identity(sim_env) -> None: + pos_fk, rot_fk = sim_env._fk(sim_env._measured_arm_q()) + pos_live, rot_live = sim_env._measured_eef_pose() + assert np.allclose(pos_fk, pos_live, atol=_FK_ATOL), f'fk pos {pos_fk} vs live {pos_live}' + assert np.allclose(rot_fk, rot_live, atol=_FK_ATOL), f'fk rot {rot_fk} vs live {rot_live}' + print(f' fk identity: OK (matches the grasp-site read, atol {_FK_ATOL})') + + +def _check_ik_roundtrip(sim_env) -> None: + measured = np.asarray(sim_env._measured_arm_q(), dtype=np.float64) + for _ in range(_IK_SAMPLES): + target_pos, target_rot = sim_env._fk(measured + np.random.uniform(-_JOINT_JITTER, _JOINT_JITTER, measured.size)) + pos, rot = sim_env._fk(sim_env._ik(target_pos, target_rot)) + ang = _ori_error(target_rot, rot) + assert np.allclose(pos, target_pos, atol=_POS_ATOL), f'ik pos off by {pos - target_pos}' + assert ang < _ORI_ATOL, f'ik orientation off by {ang} rad' + print(f' ik round-trip: OK ({_IK_SAMPLES} reachable targets, pos<{_POS_ATOL} m, ori<{_ORI_ATOL} rad)') + + +def _check_cartesian_command_is_a_noop_at_the_measured_pose(sim_env) -> None: + # Commanding the pose the arm already holds must resolve to (essentially) the joints it already holds — + # the property that makes an absolute Cartesian setpoint stable when a policy re-sends it. + pos, rot = sim_env._measured_eef_pose() + command = {protocol.COMMAND_TYPE: protocol.CARTESIAN, protocol.COMMAND_POSE: np.concatenate([pos, rot.reshape(-1)])} + target = env.mapping.wire_command_to_arm_action(command, sim_env._measured_arm_q(), ik=sim_env._ik) + drift = np.abs(np.asarray(target, dtype=np.float64) - np.asarray(sim_env._measured_arm_q(), dtype=np.float64)) + assert drift.max() < 1e-3, f'holding the measured pose moved the joints by {drift.max()} rad' + print(f' cartesian hold: OK (max joint drift {drift.max():.2e} rad)') + + +def _check_every_canonical_command_converts(sim_env) -> None: + """Drive every canonical command type through the real conversion — the adoption's whole obligation. + + The command contract is total: MolmoSpaces' Franka natively takes joint-position targets alone, so each + canonical type has to reach it as one. ``mapping``'s unit tests pin that routing against a stub solver; + here each type runs through the live IK and the measured pose, so a type the rig cannot actually resolve + fails. The iteration is over ``protocol.CANONICAL_COMMAND_TYPES`` rather than a list written here, so a + type added to the wire fails this check until the rig converts it. + """ + measured = np.asarray(sim_env._measured_arm_q(), dtype=np.float64) + pos, rot = sim_env._measured_eef_pose() + identity_rot = np.eye(3).reshape(-1) + payloads = { + protocol.CARTESIAN: {protocol.COMMAND_POSE: np.concatenate([pos, rot.reshape(-1)])}, + protocol.CARTESIAN_DELTA: {protocol.COMMAND_DELTA: np.concatenate([np.full(3, _DELTA_POS), identity_rot])}, + protocol.JOINT_POS: {protocol.COMMAND_JOINT_POS: measured}, + protocol.JOINT_VEL: {protocol.COMMAND_JOINT_VEL: np.full(measured.size, _DELTA_Q)}, + protocol.HOLD: {}, + } + unmapped = [kind for kind in protocol.CANONICAL_COMMAND_TYPES if kind not in payloads] + assert not unmapped, f'the rig has no wire payload for canonical command types {unmapped}' + + for kind in protocol.CANONICAL_COMMAND_TYPES: + command = {protocol.COMMAND_TYPE: kind, **payloads[kind]} + target = env.mapping.wire_command_to_arm_action( + command, measured, ik=sim_env._ik, current_eef=sim_env._measured_eef_pose() + ) + target = np.asarray(target, dtype=np.float64) + assert target.shape == measured.shape, f'{kind}: joint targets {target.shape} vs measured {measured.shape}' + assert np.all(np.isfinite(target)), f'{kind}: non-finite joint targets {target}' + covered = ', '.join(protocol.CANONICAL_COMMAND_TYPES) + print(f' command contract: OK ({covered} -> {measured.size} joint targets)') + + +def main() -> None: + parser = argparse.ArgumentParser(description="Validate the MolmoSpaces rig's Cartesian command transform.") + parser.add_argument('--benchmark_dir', required=True, help='dir containing benchmark.json') + parser.add_argument('--episode_index', type=int, default=0) + parser.add_argument('--seed', type=int, default=0) + # These checks are pure kinematics — they never step the sim, so the episode horizon is irrelevant to them. + # The override keeps a benchmark that declares no ``task_horizon_sec`` (the checked-in test benchmark is one) + # usable as a scene here, instead of failing the build over a field this run never reads. + parser.add_argument('--task_horizon_steps', type=int, default=1) + args = parser.parse_args() + np.random.seed(0) + + sim_env = env.MolmoSpacesEnv(Path(args.benchmark_dir), args.task_horizon_steps) + sim_env.reset({mapping.TOKEN_EPISODE_INDEX: args.episode_index, mapping.TOKEN_SEED: args.seed}) + print(f'molmo_spaces episode {args.episode_index} (seed {args.seed})') + try: + _check_fk_identity(sim_env) + _check_ik_roundtrip(sim_env) + _check_cartesian_command_is_a_noop_at_the_measured_pose(sim_env) + _check_every_canonical_command_converts(sim_env) + finally: + sim_env.close() + print('all checks passed') + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/robolab/adapter.py b/positronic/simulator/robolab/adapter.py index 73a6f7999..312bb7273 100644 --- a/positronic/simulator/robolab/adapter.py +++ b/positronic/simulator/robolab/adapter.py @@ -14,6 +14,7 @@ from positronic import geom, keys from positronic.drivers.roboarm.models import DROID_EE_FRAME from positronic.eval import keys as eval_keys +from positronic.simulator.env_server import protocol from positronic.simulator.env_server.adapter import WireCommandAdapter from positronic.simulator.mujoco.sim import MujocoFrankaState from positronic.simulator.robolab import keys as robolab_keys @@ -62,4 +63,4 @@ def privileged(self, raw_obs: dict[str, Any]) -> dict[str, Any]: def terminal(self, result: dict[str, Any]) -> dict[str, Any] | None: # ``done`` covers termination and truncation, so the trial ends either way; ``success`` is True only # when the task's success condition fired, keeping timeouts honest. - return {eval_keys.SUCCESS: bool(result['success'])} if result['done'] else None + return {eval_keys.SUCCESS: bool(result[protocol.FRAME_SUCCESS])} if result[protocol.FRAME_DONE] else None diff --git a/positronic/tests/test_keys.py b/positronic/tests/test_keys.py index 10a32366f..053247f4b 100644 --- a/positronic/tests/test_keys.py +++ b/positronic/tests/test_keys.py @@ -4,6 +4,7 @@ from positronic import keys from positronic.eval import keys as eval_keys from positronic.simulator.libero import keys as libero_keys +from positronic.simulator.molmo_spaces import keys as molmo_keys from positronic.simulator.robolab import keys as robolab_keys # Namespaced raw wire keys that denote an observation signal, and the keys a trial records — the params it @@ -33,6 +34,8 @@ libero_keys.SETTLE_STEPS, eval_keys.TASK, robolab_keys.INSTRUCTION_TYPE, + molmo_keys.EPISODE_INDEX, + molmo_keys.TASK_HORIZON, keys.OBS_TIME_NS, keys.WALL_TIME_NS, } diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 5979b4f60..83ca90dbf 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -53,7 +53,7 @@ def __init__( 'observation.state': self._derive_state, 'observation.images.left': partial(self._derive_image, wrist_camera), 'observation.images.side': partial(self._derive_image, exterior_camera), - 'task': Get(keys.TASK, ''), + keys.TASK: Get(keys.TASK, ''), } state_dim = sum(state_features.values()) @@ -128,7 +128,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={openpi.JOINT_POSITION: {keys.JOINTS: 7}, openpi.GRIPPER_POSITION: {keys.GRIP: 1}}, @@ -137,6 +138,7 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came openpi.EXTERIOR_IMAGE_LEFT: (keys.EXTERIOR_IMAGE, (224, 224)), }, task_field=openpi.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 281f7640e..ab071c96a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -267,14 +267,16 @@ include = ["client", "pimm", "positronic"] # mujoco ships no `.pyi` and re-exports from binary extension modules, so its whole API reads as # attribute errors without these. Regenerate with utilities/generate_mujoco_stubs.py. stubPath = "stubs" -# 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 stays +# checked even though it runs in MolmoSpaces' own venv: its foreign imports are suppressed in-file by a +# scoped `reportMissingImports` plus targeted per-line ignores, so every other diagnostic still fires. 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"] # ``robolab/env.py`` runs in RoboLab's isolated interpreter with the env-server modules (``server``, @@ -301,4 +303,7 @@ positronic = [ "drivers/**/*.urdf", "server/static/**", "server/templates/**", + # The launcher feeds this to `uv pip install -c` when it builds the MolmoSpaces venv, so a packaged + # install needs it on disk just as much as a checkout does. + "simulator/molmo_spaces/molmo_constraints.txt", ]