diff --git a/.github/workflows/unit-test.yaml b/.github/workflows/unit-test.yaml index fa5a545c5..57a3b9040 100644 --- a/.github/workflows/unit-test.yaml +++ b/.github/workflows/unit-test.yaml @@ -36,7 +36,7 @@ jobs: PYTHONPATH: ${{ env.PYTHONPATH }}:$PWD run: >- uv run pytest --cov-report=html - --override-ini='testpaths=pimm/tests positronic/cfg/tests positronic/dataset/tests positronic/geom/tests positronic/offboard/tests positronic/policy/tests positronic/tests positronic/utils/tests' + --override-ini='testpaths=pimm/tests positronic/cfg/tests positronic/dataset/tests positronic/geom/tests positronic/offboard/tests positronic/policy/tests positronic/simulator/molmo_spaces/tests positronic/tests positronic/utils/tests' - name: Coverage summary (job summary) if: always() diff --git a/CLAUDE.md b/CLAUDE.md index 4a77dd8df..0f7798ca9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -75,4 +75,6 @@ # Infrastructure - Machines, Docker contexts and images: `docker/CONTEXTS.md` - Model-specific workflows: `positronic/vendors/{lerobot,gr00t,openpi}/README.md` +- Inference serving, and the adapter/codec/wire-client separation of responsibilities (read BEFORE + writing a sim/rig adapter): `positronic/offboard/README.md` - Reconstructing previous runs: read `run_metadata_*.yaml` and episode `static.json` from output directory diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ce1cbbedc..074ab813a 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -2,6 +2,34 @@ This package implements the protocol and utilities for offboard policy inference, allowing robots or simulators to stream observations to a remote server and receive actions. +## Separation of responsibilities: adapter vs codec vs wire client + +Three layers touch an observation on its way to a model, and each owns exactly one concern. +When writing a new sim/rig adapter, check this table before adding any transform to it: + +| Layer | Owns | Examples | +|---|---|---| +| **Adapter** (per sim/rig, e.g. `simulator/molmo_spaces/adapter.py`) | Rig semantics ONLY: mapping the rig's observation/action vocabulary onto positronic's raw keys | Camera-key mapping, gripper qpos → `[0, 1]` closure, decoded commands → the rig's action format | +| **Codec** (per model family, `policy/codec.py` subclasses) | Model preprocessing: everything the checkpoint's input distribution requires | Resize-with-pad to model resolution, prompt normalization (e.g. DROID lowercasing), state assembly | +| **Wire client** (`InferenceClient` / `RemotePolicy`) | Transport optimization, negotiated — never semantics | Downscaling frames to the server-advertised `image_sizes` (aspect-preserving, never upscaling), optional JPEG compression | + +Consequences: + +- **An adapter never resizes, pads, normalizes prompts, or otherwise preprocesses for the model.** + It passes frames and text through at native fidelity. If the same transform appears in an adapter + and a codec, the adapter's copy is the bug: a drifted duplicate silently changes eval inputs. +- **Bandwidth is not the adapter's problem.** The client already downsizes to what the server says + it needs: every `Codec` advertises its expected input sizes via the reserved `image_sizes` meta + key (see `Codec.meta`), the server returns it in the session handshake, and the client fits + frames to it before sending. This is default-on — an adapter that resizes "to keep the wire + payload small" is duplicating it. +- **Codecs run on either side of the wire.** positronic-native evals compose the codec around + `RemotePolicy` on the client (`cfg/policy.py` — the wire then carries model-sized encoded inputs, + and the client-side resize is disabled since `codec.meta` already reports `image_sizes`). + Thin-client deployments (a sim adapter in a foreign venv talking to a serverless endpoint) host + the codec on the server — the wire carries raw positronic keys, downsized by the negotiation + above. Both placements are supported; pick by where the dependencies can live. + ## Protocol v1 The unified WebSocket protocol is built to enable ANY hardware to connect to ANY model. All Positronic inference servers (LeRobot, GR00T, OpenPI) implement this protocol, allowing a single `.remote` policy client to work across all vendors. diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index 9d424b7a0..e7709fd27 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -18,14 +18,21 @@ class ObservationCodec(Codec): state: mapping from output state key to an ordered dict of {episode_key: dim} to concatenate. images: mapping from output image name to tuple (input_key, (width, height)). task_field: output key carrying the language prompt at inference; LeRobot training always uses ``task``. + lowercase_task: lowercase the task text at inference, for checkpoints trained on lowercased language + (the pretrained DROID models; MolmoSpaces' Pi baseline applies the same normalization). """ def __init__( - self, state: dict[str, dict[str, int]], images: dict[str, tuple[str, tuple[int, int]]], task_field: str = 'task' + self, + state: dict[str, dict[str, int]], + images: dict[str, tuple[str, tuple[int, int]]], + task_field: str = 'task', + lowercase_task: bool = False, ): self._state = state self._image_configs = images self._task_field = task_field + self._lowercase_task = lowercase_task self._derive_transforms = {k: partial(self._derive_state, k) for k in state.keys()} self._derive_transforms.update({k: partial(self._derive_image, k) for k in images.keys()}) @@ -54,7 +61,8 @@ def encode(self, inputs: dict[str, Any]) -> dict[str, Any]: obs: dict[str, Any] = {} if 'task' in inputs: - obs[self._task_field] = inputs['task'] + task = inputs['task'] + obs[self._task_field] = task.lower() if self._lowercase_task else task for out_name, (input_key, (width, height)) in self._image_configs.items(): if input_key not in inputs: diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index 37470d9bb..0f6d4eee8 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -59,12 +59,17 @@ def test_observation_encode_missing_state_inputs_raise(): def test_observation_encode_task(): enc = ObservationCodec(state={'observation.state': ['a']}, images={}) - obs = enc.encode({'a': 1.0, 'task': 'test_task'}) - assert obs['task'] == 'test_task' + obs = enc.encode({'a': 1.0, 'task': 'Test_Task'}) + assert obs['task'] == 'Test_Task' # untouched by default obs_no_task = enc.encode({'a': 1.0}) assert 'task' not in obs_no_task + # DROID-style configs lowercase the prompt: those checkpoints were trained on lowercased language, + # so capitalized benchmark task text must not reach them mixed-case. + lower = ObservationCodec(state={'observation.state': ['a']}, images={}, lowercase_task=True) + assert lower.encode({'a': 1.0, 'task': 'Pick up the Cube'})['task'] == 'pick up the cube' + def test_absolute_position_action_encode_decode_quat(): # Identity rotation, known translation/grip diff --git a/positronic/simulator/libero/make_fixture.py b/positronic/simulator/libero/tests/make_fixture.py similarity index 95% rename from positronic/simulator/libero/make_fixture.py rename to positronic/simulator/libero/tests/make_fixture.py index b88b47969..ce98e4fc3 100644 --- a/positronic/simulator/libero/make_fixture.py +++ b/positronic/simulator/libero/tests/make_fixture.py @@ -10,7 +10,7 @@ The e2e replay needs only each demo's action sequence and its initial full state — a few KB per episode, not the multi-GB benchmark. Run once on a box that has the demos, then commit the ``.npz`` next to the test:: - uv run --no-project positronic/simulator/libero/make_fixture.py \ + uv run --no-project positronic/simulator/libero/tests/make_fixture.py \ --demo-path "$LIBERO_DATASETS/libero_spatial/_demo.hdf5" \ --out positronic/simulator/libero/tests/libero_spatial_task0.npz """ diff --git a/positronic/simulator/libero/tests/test_e2e.py b/positronic/simulator/libero/tests/test_e2e.py index e639d9d0a..bd121fcf9 100644 --- a/positronic/simulator/libero/tests/test_e2e.py +++ b/positronic/simulator/libero/tests/test_e2e.py @@ -7,7 +7,7 @@ The fixture is generated once on a LIBERO box:: - uv run --no-project positronic/simulator/libero/make_fixture.py \ + uv run --no-project positronic/simulator/libero/tests/make_fixture.py \ --demo-path "$LIBERO_DATASETS/libero_spatial/_demo.hdf5" \ --out positronic/simulator/libero/tests/libero_spatial_task0.npz 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..56614319a --- /dev/null +++ b/positronic/simulator/molmo_spaces/adapter.py @@ -0,0 +1,423 @@ +"""Bridge a pi05_droid policy served over positronic's inference protocol into a MolmoSpaces eval. + +MolmoSpaces drives a DROID FR3 + Robotiq rig and expects a ``BasePolicy`` whose ``get_action`` maps its +observation dict to a per-move-group action dict. positronic serves pi05_droid behind an ``InferenceServer`` +(FastAPI on port 8000) whose DROID codec consumes the raw positronic observation keys +(``robot_state.q``, ``grip``, ``image.wrist``, ``image.exterior``, ``task``) and returns a chunk of decoded +per-step ``JointDelta`` commands (7 joint velocities already scaled by ``MAX_JOINT_DELTA``) plus a binarized grip. + +This module holds: + - two pure mapping functions — ``molmo_obs_to_positronic`` and ``positronic_action_to_molmo`` — that carry the + whole translation and are import-free of both frameworks so they run under a bare pytest; + - ``ChunkBuffer``, which plays one buffered chunk step per policy tick and re-queries when the chunk drains, + matching DROID's open-loop horizon; + - ``FakePolicy``, a server-free stand-in honoring the positronic client contract, for smoke runs; + - ``MolmoSpacesPolicy``, the ``InferencePolicy`` subclass wiring the above into MolmoSpaces (only usable where + molmo_spaces is installed; the mapping logic above is not). + +MolmoSpaces and positronic are imported behind guards so the pure logic and its tests need neither installed. +""" + +import collections.abc as cabc +from dataclasses import dataclass +from typing import Any + +import numpy as np + +try: + from molmo_spaces.policy.base_policy import InferencePolicy as _InferencePolicyBase + + HAS_MOLMO_SPACES = True +except ImportError: + _InferencePolicyBase = object + HAS_MOLMO_SPACES = False + + +# MolmoSpaces DROID rig observation keys (FrankaDroidCameraSystem + RobotJointPositionSensor). +MOLMO_WRIST_CAMERA = 'wrist_camera' +MOLMO_EXTERIOR_CAMERA = 'exo_camera_1' +# MolmoSpaces' Zed-wrist / light-randomized benchmark variants emit these instead of the defaults, and its own +# Pi policy prefers them when present (molmo_spaces/policy/learned_policy/pi_policy.py) — mirrored here. +MOLMO_WRIST_CAMERA_VARIANTS = ('wrist_camera_zed_mini',) +MOLMO_EXTERIOR_CAMERA_VARIANTS = ('droid_shoulder_light_randomization',) +MOLMO_ARM_GROUP = 'arm' +MOLMO_GRIPPER_GROUP = 'gripper' + +# positronic raw observation keys the DROID codec (positronic/vendors/openpi/codecs.py:droid_obs) reads. +POS_JOINTS = 'robot_state.q' +POS_GRIP = 'grip' +POS_WRIST_IMAGE = 'image.wrist' +POS_EXTERIOR_IMAGE = 'image.exterior' +POS_TASK = 'task' + +NUM_ARM_JOINTS = 7 + +# Robotiq closure at which the FR3 gripper qpos saturates; the pi baseline normalizes proprio grip by it +# (molmospaces pi_policy.py:126) so the served model sees the [0, 1] closure it was trained on. +GRIPPER_QPOS_CLOSED = 0.824033 + +# Robotiq 2F-85 control convention: 0 fully open, 255 fully closed (franka_droid_view.py:43). +ROBOTIQ_OPEN = 0.0 +ROBOTIQ_CLOSED = 255.0 + +# DROID re-queries after an 8-step open-loop horizon (codecs.py:172); mirror it in the fake client. +DROID_CHUNK_STEPS = 8 + +# JointDelta scales the clipped [-1, 1] velocities by this (positronic policy/action.py:194). +MAX_JOINT_DELTA = 0.2 + + +def _single_env(observation: Any) -> dict: + """MolmoSpaces yields ``observation`` as a list (one dict per batch env); single-env eval uses index 0.""" + if isinstance(observation, cabc.Mapping): + return observation + if isinstance(observation, cabc.Sequence): + return observation[0] + raise TypeError(f'Unexpected observation container: {type(observation)}') + + +def molmo_obs_to_positronic( + observation: Any, + task: str, + *, + wrist_key: str = MOLMO_WRIST_CAMERA, + exterior_key: str = MOLMO_EXTERIOR_CAMERA, + gripper_qpos_closed: float = GRIPPER_QPOS_CLOSED, +) -> dict[str, Any]: + """Map a MolmoSpaces DROID observation to the raw positronic keys the DROID codec consumes. + + ``observation`` may be the batch list or a single env dict. The arm joints pass through as 7 absolute radians; + the Robotiq finger qpos is normalized to the ``[0, 1]`` closure the model was trained on. Everything else is + the mapping only: camera frames and task text pass through untouched — model preprocessing (resize-with-pad, + prompt normalization) is the server codec's job, and wire downsizing is negotiated by the inference client + from the server's advertised ``image_sizes``. + """ + env = _single_env(observation) + qpos = env['qpos'] + arm = np.asarray(qpos['arm'], dtype=np.float32).reshape(-1) + grip_qpos = float(np.asarray(qpos['gripper']).reshape(-1)[0]) + grip = float(np.clip(grip_qpos / gripper_qpos_closed, 0.0, 1.0)) + + return { + POS_JOINTS: arm, + POS_GRIP: np.array([grip], dtype=np.float32), + POS_WRIST_IMAGE: _camera_image(env, wrist_key, MOLMO_WRIST_CAMERA, MOLMO_WRIST_CAMERA_VARIANTS), + POS_EXTERIOR_IMAGE: _camera_image(env, exterior_key, MOLMO_EXTERIOR_CAMERA, MOLMO_EXTERIOR_CAMERA_VARIANTS), + POS_TASK: task, + } + + +def _camera_image(env: cabc.Mapping, key: str, default: str, variants: tuple[str, ...]) -> Any: + """Resolve a camera image the way MolmoSpaces' own policies do: a benchmark-variant key takes precedence + over the default when present; an explicitly configured non-default key is read as-is.""" + if key != default: + return env[key] + for candidate in (*variants, key): + if candidate in env: + return env[candidate] + raise KeyError(f'observation has none of {(*variants, key)}; available keys: {sorted(env)}') + + +def _wire_get(mapping: cabc.Mapping, name: str, default: Any = None) -> Any: + """Read a key that may arrive str- or bytes-keyed: msgpack deserialisation yields bytes keys on some + positronic/msgpack version combinations, so both spellings must be tried.""" + for key in (name, name.encode()): + if key in mapping: + return mapping[key] + return default + + +def _joint_delta_velocities(robot_command: Any) -> np.ndarray: + """Read the 7 joint velocities from a decoded positronic ``JointDelta`` (object) or its wire dict.""" + if hasattr(robot_command, 'velocities'): + return np.asarray(robot_command.velocities, dtype=np.float32).reshape(-1) + if isinstance(robot_command, cabc.Mapping): + velocities = _wire_get(robot_command, 'velocities') + if velocities is not None: + return np.asarray(velocities, dtype=np.float32).reshape(-1) + raise TypeError(f'Cannot read joint velocities from {type(robot_command)}') + + +def positronic_action_to_molmo( + action: cabc.Mapping, + current_arm_qpos: np.ndarray, + *, + arm_group: str = MOLMO_ARM_GROUP, + gripper_group: str = MOLMO_GRIPPER_GROUP, +) -> dict[str, np.ndarray]: + """Turn one decoded positronic DROID action into a MolmoSpaces per-move-group action. + + The server returns a per-step ``JointDelta`` (velocities relative to the live measured joints, DROID's control + convention) plus a binarized grip. Integrating each delta onto the joints measured this tick reproduces the + positronic driver's ``set_target_joints(q + delta)``; the grip maps to the Robotiq open/closed control values. + """ + robot_command = _wire_get(action, 'robot_command') + if robot_command is None: + raise KeyError(f'no robot_command in action; keys={list(action.keys())}') + velocities = _joint_delta_velocities(robot_command) + current = np.asarray(current_arm_qpos, dtype=np.float32).reshape(-1) + if velocities.shape[0] != current.shape[0]: + raise ValueError(f'Joint count mismatch: delta {velocities.shape[0]} vs qpos {current.shape[0]}') + + grip = float(_wire_get(action, 'target_grip', 0.0)) + gripper = ROBOTIQ_CLOSED if grip > 0.5 else ROBOTIQ_OPEN + return {arm_group: (current + velocities).astype(np.float32), gripper_group: np.array([gripper], dtype=np.float32)} + + +class ChunkBuffer: + """Plays one action per tick from a buffered inference chunk, re-querying the session when it drains. + + The positronic session returns a whole action chunk per call; MolmoSpaces asks for one action per policy step. + Buffering here re-queries every ``len(chunk)`` steps, matching DROID's open-loop horizon. + """ + + def __init__(self, session: Any): + self._session = session + self._pending: list[Any] = [] + + def next(self, obs: dict[str, Any]) -> Any: + if not self._pending: + chunk = self._session(obs) + # The serving layer ends each chunk with a horizon marker carrying only `timestamp` + # (droid's action window is horizon=8/15: 8 actions + the window-end stamp). It is not + # an action — playing it as one KeyErrors on `robot_command`, so keep only action-bearing entries. + self._pending = [c for c in (chunk or []) if _wire_get(c, 'robot_command') is not None] + if not self._pending: + raise RuntimeError('Inference session returned an empty action chunk') + return self._pending.pop(0) + + def reset(self) -> None: + self._pending = [] + + def close(self) -> None: + close = getattr(self._session, 'close', None) + if close is not None: + close() + + +@dataclass +class _FakeJointDelta: + """Duck-typed stand-in for positronic ``command.JointDelta`` (read via ``.velocities``).""" + + TYPE = 'joint_delta' + velocities: np.ndarray + + +class FakeSession: + """Server-free session emitting action chunks in the shape a positronic ``RemoteSession`` returns.""" + + def __init__( + self, *, chunk_size: int, num_joints: int, max_joint_delta: float, mode: str, rng: np.random.Generator + ): + self._chunk_size = chunk_size + self._num_joints = num_joints + self._max_joint_delta = max_joint_delta + self._mode = mode + self._rng = rng + + def _velocities(self) -> np.ndarray: + if self._mode == 'zero': + return np.zeros(self._num_joints, dtype=np.float32) + return (self._rng.uniform(-1.0, 1.0, self._num_joints) * self._max_joint_delta).astype(np.float32) + + def _grip(self) -> float: + if self._mode == 'zero': + return 0.0 + return 1.0 if self._rng.random() > 0.5 else 0.0 + + def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]]: + return [ + {'robot_command': _FakeJointDelta(self._velocities()), 'target_grip': self._grip()} + for _ in range(self._chunk_size) + ] + + @property + def meta(self) -> dict[str, Any]: + return {'type': 'fake'} + + def close(self) -> None: + return None + + +class FakePolicy: + """Drop-in for positronic ``RemotePolicy`` that needs no server — random or zero DROID action chunks. + + ``mode='random'`` draws velocities uniformly in ``[-max_joint_delta, max_joint_delta]`` and a random binary + grip; ``mode='zero'`` holds the arm and keeps the gripper open. Sessions are deterministic given ``seed``. + """ + + def __init__( + self, + *, + chunk_size: int = DROID_CHUNK_STEPS, + num_joints: int = NUM_ARM_JOINTS, + max_joint_delta: float = MAX_JOINT_DELTA, + mode: str = 'random', + seed: int = 0, + ): + if mode not in ('random', 'zero'): + raise ValueError(f"mode must be 'random' or 'zero', got {mode!r}") + self._chunk_size = chunk_size + self._num_joints = num_joints + self._max_joint_delta = max_joint_delta + self._mode = mode + self._seed = seed + self._session_count = 0 + + def new_session(self, context: dict[str, Any] | None = None) -> FakeSession: + rng = np.random.default_rng(self._seed + self._session_count) + self._session_count += 1 + return FakeSession( + chunk_size=self._chunk_size, + num_joints=self._num_joints, + max_joint_delta=self._max_joint_delta, + mode=self._mode, + rng=rng, + ) + + @property + def meta(self) -> dict[str, Any]: + return {'type': 'fake', 'mode': self._mode} + + def close(self) -> None: + return None + + +def make_policy_client( + *, + fake: bool = False, + host: str = 'localhost', + port: int = 8000, + model_id: str | None = None, + secure: bool = False, + fake_mode: str = 'random', + fake_seed: int = 0, + fake_chunk_size: int = DROID_CHUNK_STEPS, +) -> Any: + """Return the inference client the policy talks to: a ``FakePolicy`` when ``fake`` else positronic ``RemotePolicy``. + + ``RemotePolicy`` is imported lazily so this module (and its tests) load without positronic installed. + """ + if fake: + return FakePolicy(mode=fake_mode, seed=fake_seed, chunk_size=fake_chunk_size) + from positronic.policy.remote import RemotePolicy + + return RemotePolicy(host, port=port, model_id=model_id, secure=secure) + + +@dataclass +class AdapterConfig: + """Static rig/codec wiring for the adapter, independent of any framework object.""" + + wrist_key: str = MOLMO_WRIST_CAMERA + exterior_key: str = MOLMO_EXTERIOR_CAMERA + arm_group: str = MOLMO_ARM_GROUP + gripper_group: str = MOLMO_GRIPPER_GROUP + gripper_qpos_closed: float = GRIPPER_QPOS_CLOSED + + +def adapter_config_from_exp_config(config: Any) -> AdapterConfig: + """An ``AdapterConfig`` honoring the eval's ``policy_config.camera_names``, when declared. + + MolmoSpaces' documented custom-policy convention (evaluation/README.md) is that the policy config carries + ``camera_names`` and ``eval_main.py --camera_names`` overrides it; a policy reads its cameras from there. + The wrist role is the name containing ``'wrist'`` (true of every upstream wrist camera, incl. the Zed + variant), the exterior role the first name without it; a role with no match keeps its default key, so the + benchmark-variant fallback in ``_camera_image`` still applies. + """ + names = getattr(getattr(config, 'policy_config', None), 'camera_names', None) or () + wrist = next((n for n in names if 'wrist' in n), MOLMO_WRIST_CAMERA) + exterior = next((n for n in names if 'wrist' not in n), MOLMO_EXTERIOR_CAMERA) + return AdapterConfig(wrist_key=wrist, exterior_key=exterior) + + +def client_kwargs_from_exp_config(config: Any) -> dict[str, Any]: + """Client kwargs honoring the eval's ``policy_config.remote_config``, when declared. + + MolmoSpaces' learned-policy configs carry ``remote_config: dict(host, port)`` for a policy served remotely + (policy_configs_baselines.py), so the normal ``policy_factory(exp_config, task)`` path must reach the + configured endpoint rather than ``make_policy_client``'s defaults. + """ + remote = getattr(getattr(config, 'policy_config', None), 'remote_config', None) or {} + return {k: remote[k] for k in ('host', 'port', 'model_id', 'secure') if k in remote} + + +class MolmoSpacesPolicy(_InferencePolicyBase): + """MolmoSpaces ``InferencePolicy`` serving pi05_droid through a positronic inference client. + + Wires the pure mapping functions and ``ChunkBuffer`` into the ``obs_to_model_input`` → ``inference_model`` → + ``model_output_to_action`` pipeline. Instantiable only where molmo_spaces is installed; the mapping logic it + delegates to is not, and is unit-tested directly. + """ + + def __init__( + self, + config: Any = None, + task: Any = None, + *, + client: Any = None, + prompt: str = '', + adapter_config: AdapterConfig | None = None, + client_kwargs: dict[str, Any] | None = None, + ): + if HAS_MOLMO_SPACES: + super().__init__(config, task) + if client is None: + kwargs = client_kwargs if client_kwargs is not None else client_kwargs_from_exp_config(config) + client = make_policy_client(**kwargs) + self._client = client + self._prompt = prompt + self._adapter = adapter_config if adapter_config is not None else adapter_config_from_exp_config(config) + self._buffer: ChunkBuffer | None = None + self._task_text = prompt + self._current_arm_qpos = np.zeros(NUM_ARM_JOINTS, dtype=np.float32) + + def prepare_model(self) -> None: + self._open_session() + + def reset(self) -> None: + task = getattr(self, 'task', None) + self._task_text = task.get_task_description() if task is not None else self._prompt + self._open_session() + + def _open_session(self) -> None: + if self._buffer is not None: + self._buffer.close() + self._buffer = ChunkBuffer(self._client.new_session()) + + def close(self) -> None: + """Close the open inference session. MolmoSpaces' runner never calls this itself, so the caller closes + each episode's policy explicitly — otherwise every per-episode policy leaks one live server session + (each of which also resets a serverless endpoint's idle clock).""" + if self._buffer is not None: + self._buffer.close() + self._buffer = None + + def obs_to_model_input(self, observation: Any) -> dict[str, Any]: + env = _single_env(observation) + self._current_arm_qpos = np.asarray(env['qpos']['arm'], dtype=np.float32).reshape(-1) + return molmo_obs_to_positronic( + env, + self._task_text, + wrist_key=self._adapter.wrist_key, + exterior_key=self._adapter.exterior_key, + gripper_qpos_closed=self._adapter.gripper_qpos_closed, + ) + + def inference_model(self, model_input: dict[str, Any]) -> Any: + if self._buffer is None: + self._open_session() + return self._buffer.next(model_input) + + def model_output_to_action(self, model_output: cabc.Mapping) -> dict[str, np.ndarray]: + return positronic_action_to_molmo( + model_output, + self._current_arm_qpos, + arm_group=self._adapter.arm_group, + gripper_group=self._adapter.gripper_group, + ) + + def get_info(self) -> dict[str, Any]: + info = super().get_info() if HAS_MOLMO_SPACES else {} + info['client'] = getattr(self._client, 'meta', {}) + return info 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..96d249604 Binary files /dev/null and b/positronic/simulator/molmo_spaces/tests/droid_obs.npz differ diff --git a/positronic/simulator/molmo_spaces/tests/make_fixture.py b/positronic/simulator/molmo_spaces/tests/make_fixture.py new file mode 100644 index 000000000..0e876d8c3 --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/make_fixture.py @@ -0,0 +1,51 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +"""Regenerate the synthetic MolmoSpaces DROID observation fixture used by test_adapter.py. + +MolmoSpaces renders real MuJoCo scenes that need the full asset stack and a GPU, so committing a real render is +impractical; the adapter under test only touches observation *structure* (keys, shapes, dtypes, gripper qpos +scaling, image resize), which a tiny synthetic frame exercises exactly. Images are kept small (36x64, the DROID +16:9 aspect) and orientation-marked so a resize regression or a wrist/exterior swap is visible. + +Run: uv run --no-project positronic/simulator/molmo_spaces/tests/make_fixture.py +Output: droid_obs.npz next to this script (well under 100 KB) +""" + +from pathlib import Path + +import numpy as np + +RIG_HEIGHT, RIG_WIDTH = 36, 64 # (H, W); DROID exo/wrist cameras are 16:9. + + +def _marked_frame(base_rgb: tuple[int, int, int]) -> np.ndarray: + """A solid-color frame with a white top-left block and a black bottom-right block — an orientation marker. + + resize_with_pad preserves orientation, so a vertical flip or a left/right swap moves these markers detectably. + """ + frame = np.zeros((RIG_HEIGHT, RIG_WIDTH, 3), dtype=np.uint8) + frame[:] = base_rgb + frame[:8, :12] = 255 # top-left white + frame[-8:, -12:] = 0 # bottom-right black + return frame + + +def build_observation() -> dict[str, np.ndarray]: + wrist = _marked_frame((200, 40, 40)) # reddish wrist view + exterior = _marked_frame((40, 160, 40)) # greenish exterior view + qpos_arm = np.array([0.0, -0.785, 0.0, -2.356, 0.0, 1.571, 0.0], dtype=np.float32) # FR3 init_qpos + qpos_gripper = np.array([0.412016, 0.412016], dtype=np.float32) # half of GRIPPER_QPOS_CLOSED -> grip 0.5 + return {'wrist_camera': wrist, 'exo_camera_1': exterior, 'qpos_arm': qpos_arm, 'qpos_gripper': qpos_gripper} + + +def main() -> None: + out = Path(__file__).parent / 'droid_obs.npz' + obs = build_observation() + np.savez_compressed(out, **obs) + print(f'Wrote {out} ({out.stat().st_size} bytes)') + + +if __name__ == '__main__': + main() diff --git a/positronic/simulator/molmo_spaces/tests/test_adapter.py b/positronic/simulator/molmo_spaces/tests/test_adapter.py new file mode 100644 index 000000000..917d97cba --- /dev/null +++ b/positronic/simulator/molmo_spaces/tests/test_adapter.py @@ -0,0 +1,330 @@ +"""Unit tests for the pi05_droid <-> MolmoSpaces adapter mapping logic. + +Runs with NEITHER molmo_spaces nor positronic installed: the adapter import-guards both, and every test here +exercises the pure mapping functions, ``ChunkBuffer``, and ``FakePolicy`` — none of which touch either framework. + +Run: uv run --locked pytest positronic/simulator/molmo_spaces/tests/test_adapter.py --no-cov +""" + +from pathlib import Path + +import numpy as np +import pytest + +from positronic.simulator.molmo_spaces import adapter +from positronic.simulator.molmo_spaces.adapter import ( + NUM_ARM_JOINTS, + POS_EXTERIOR_IMAGE, + POS_GRIP, + POS_JOINTS, + POS_TASK, + POS_WRIST_IMAGE, + ROBOTIQ_CLOSED, + ROBOTIQ_OPEN, + ChunkBuffer, + FakePolicy, + _FakeJointDelta, + adapter_config_from_exp_config, + molmo_obs_to_positronic, + positronic_action_to_molmo, +) + +FIXTURE = Path(__file__).parent / 'droid_obs.npz' + + +def _load_env_obs() -> dict: + data = np.load(FIXTURE) + return { + 'wrist_camera': data['wrist_camera'], + 'exo_camera_1': data['exo_camera_1'], + 'qpos': {'arm': data['qpos_arm'], 'gripper': data['qpos_gripper']}, + } + + +# --- import guard ----------------------------------------------------------------------------------------------- + + +def test_module_imports_without_frameworks(): + # The pure logic must be usable even when molmo_spaces is absent (the common test/dev box). + assert isinstance(adapter.HAS_MOLMO_SPACES, bool) + assert callable(molmo_obs_to_positronic) + assert callable(positronic_action_to_molmo) + + +# --- observation mapping ---------------------------------------------------------------------------------------- + + +def test_obs_mapping_key_set(): + obs = molmo_obs_to_positronic(_load_env_obs(), 'pick up the cube') + assert set(obs) == {POS_JOINTS, POS_GRIP, POS_WRIST_IMAGE, POS_EXTERIOR_IMAGE, POS_TASK} + + +def test_obs_mapping_shapes_and_dtypes(): + env = _load_env_obs() + obs = molmo_obs_to_positronic(env, 'Pick up the cube') + assert obs[POS_JOINTS].shape == (NUM_ARM_JOINTS,) and obs[POS_JOINTS].dtype == np.float32 + assert obs[POS_GRIP].shape == (1,) and obs[POS_GRIP].dtype == np.float32 + # Frames and task text pass through untouched: model preprocessing (resize-with-pad, prompt lowercasing) + # belongs to the server codec, wire downsizing to the inference client. + assert np.array_equal(obs[POS_WRIST_IMAGE], env['wrist_camera']) + assert np.array_equal(obs[POS_EXTERIOR_IMAGE], env['exo_camera_1']) + assert obs[POS_TASK] == 'Pick up the cube' + + +def test_obs_mapping_accepts_batch_list_and_single_dict(): + env = _load_env_obs() + from_dict = molmo_obs_to_positronic(env, 't') + from_list = molmo_obs_to_positronic([env], 't') # MolmoSpaces yields a per-env list + assert np.array_equal(from_dict[POS_JOINTS], from_list[POS_JOINTS]) + assert np.array_equal(from_dict[POS_WRIST_IMAGE], from_list[POS_WRIST_IMAGE]) + + +def test_obs_mapping_does_not_swap_cameras(): + obs = molmo_obs_to_positronic(_load_env_obs(), 't') + # Fixture marks the wrist view reddish and the exterior view greenish; a swap would flip the dominant channel. + wrist_mean = obs[POS_WRIST_IMAGE].reshape(-1, 3).mean(axis=0) + exterior_mean = obs[POS_EXTERIOR_IMAGE].reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] # wrist: red > green + assert exterior_mean[1] > exterior_mean[0] # exterior: green > red + + +def test_obs_mapping_resolves_benchmark_variant_camera_keys(): + # Zed-wrist / light-randomized variants replace the default camera keys; MolmoSpaces' own Pi policy + # prefers the variant key when present, so the adapter must too (regression: hard indexing raised + # KeyError on those benchmark observations). + env = _load_env_obs() + env['wrist_camera_zed_mini'] = env.pop('wrist_camera') + env['droid_shoulder_light_randomization'] = env.pop('exo_camera_1') + obs = molmo_obs_to_positronic(env, 't') + wrist_mean = obs[POS_WRIST_IMAGE].reshape(-1, 3).mean(axis=0) + exterior_mean = obs[POS_EXTERIOR_IMAGE].reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] # the reddish wrist view still lands on the wrist key + assert exterior_mean[1] > exterior_mean[0] + + # A variant key coexisting with the default wins, matching the upstream policy's precedence. + both = _load_env_obs() + both['droid_shoulder_light_randomization'] = both['wrist_camera'] # reddish, unlike exo_camera_1 + obs = molmo_obs_to_positronic(both, 't') + exterior_mean = obs[POS_EXTERIOR_IMAGE].reshape(-1, 3).mean(axis=0) + assert exterior_mean[0] > exterior_mean[1] + + # An explicitly configured non-default key is read as-is, never shadowed by a variant. + custom = _load_env_obs() + custom['my_cam'] = custom['wrist_camera'] + custom['wrist_camera_zed_mini'] = custom['exo_camera_1'] # greenish decoy + obs = molmo_obs_to_positronic(custom, 't', wrist_key='my_cam') + wrist_mean = obs[POS_WRIST_IMAGE].reshape(-1, 3).mean(axis=0) + assert wrist_mean[0] > wrist_mean[1] + + +class _Ns: + def __init__(self, **kw): + self.__dict__.update(kw) + + +def test_adapter_config_honors_policy_config_camera_names(): + # An eval configured with non-default camera names (policy_config.camera_names, per MolmoSpaces' + # custom-policy convention) must reach the adapter's camera keys (regression: a default AdapterConfig + # ignored them and KeyErrored on the renamed observations). + cfg = _Ns(policy_config=_Ns(camera_names=['randomized_zed2_analogue_1', 'wrist_camera_zed_mini'])) + adapter_cfg = adapter_config_from_exp_config(cfg) + assert adapter_cfg.wrist_key == 'wrist_camera_zed_mini' + assert adapter_cfg.exterior_key == 'randomized_zed2_analogue_1' + + # No camera_names declared (BasePolicyConfig doesn't define the field) -> the defaults, variant fallback intact. + for cfg in (None, _Ns(policy_config=_Ns()), _Ns(policy_config=_Ns(camera_names=[]))): + adapter_cfg = adapter_config_from_exp_config(cfg) + assert adapter_cfg.wrist_key == adapter.MOLMO_WRIST_CAMERA + assert adapter_cfg.exterior_key == adapter.MOLMO_EXTERIOR_CAMERA + + # A single-role list keeps the missing role on its default. + adapter_cfg = adapter_config_from_exp_config(_Ns(policy_config=_Ns(camera_names=['exo_top_down']))) + assert adapter_cfg.wrist_key == adapter.MOLMO_WRIST_CAMERA + assert adapter_cfg.exterior_key == 'exo_top_down' + + # The policy's default path derives from the config it was constructed with. force_enable_depth is the one + # field MolmoSpaces' BasePolicy.__init__ reads, so the stub carries it and the test also passes on boxes + # where molmo_spaces IS installed and the real superclass path runs. + names = ['randomized_zed2_analogue_1', 'wrist_camera_zed_mini'] + cfg = _Ns(policy_config=_Ns(camera_names=names, force_enable_depth=False)) + policy = adapter.MolmoSpacesPolicy(cfg, client=FakePolicy()) + assert policy._adapter.wrist_key == 'wrist_camera_zed_mini' + assert policy._adapter.exterior_key == 'randomized_zed2_analogue_1' + + +def test_client_kwargs_honor_policy_config_remote_config(): + # MolmoSpaces' learned-policy configs carry remote_config=dict(host, port) for remotely served policies; + # the default policy_factory(exp_config, task) path must reach that endpoint, not the localhost defaults. + cfg = _Ns(policy_config=_Ns(remote_config={'host': 'h100.example', 'port': 9000})) + assert adapter.client_kwargs_from_exp_config(cfg) == {'host': 'h100.example', 'port': 9000} + # Keys make_policy_client doesn't take are filtered out. + cfg = _Ns(policy_config=_Ns(remote_config={'host': 'h', 'checkpoint_path': '/x'})) + assert adapter.client_kwargs_from_exp_config(cfg) == {'host': 'h'} + # Absent or None remote_config (or no config at all) -> defaults. + assert adapter.client_kwargs_from_exp_config(None) == {} + assert adapter.client_kwargs_from_exp_config(_Ns(policy_config=_Ns(remote_config=None))) == {} + + +def test_gripper_proprio_normalization(): + closed = adapter.GRIPPER_QPOS_CLOSED + + def grip_for(qpos_val: float) -> float: + env = _load_env_obs() + env['qpos']['gripper'] = np.array([qpos_val, qpos_val], dtype=np.float32) + return float(molmo_obs_to_positronic(env, 't')[POS_GRIP][0]) + + assert grip_for(0.0) == 0.0 + assert abs(grip_for(closed / 2) - 0.5) < 1e-4 + assert abs(grip_for(closed) - 1.0) < 1e-6 + assert grip_for(closed * 2) == 1.0 # saturates, never exceeds 1 + + +# --- action mapping --------------------------------------------------------------------------------------------- + + +def test_action_integrates_delta_onto_live_joints(): + current = np.arange(NUM_ARM_JOINTS, dtype=np.float32) + velocities = np.full(NUM_ARM_JOINTS, 0.1, dtype=np.float32) + action = {'robot_command': _FakeJointDelta(velocities), 'target_grip': 1.0} + out = positronic_action_to_molmo(action, current) + assert out['arm'].shape == (NUM_ARM_JOINTS,) and out['arm'].dtype == np.float32 + assert np.allclose(out['arm'], current + velocities) + assert out['gripper'].shape == (1,) + + +def test_action_gripper_convention(): + current = np.zeros(NUM_ARM_JOINTS, dtype=np.float32) + vel = _FakeJointDelta(np.zeros(NUM_ARM_JOINTS, dtype=np.float32)) + + def gripper_for(target_grip: float) -> float: + out = positronic_action_to_molmo({'robot_command': vel, 'target_grip': target_grip}, current) + return float(out['gripper'][0]) + + assert gripper_for(1.0) == ROBOTIQ_CLOSED == 255.0 + assert gripper_for(0.0) == ROBOTIQ_OPEN == 0.0 + assert gripper_for(0.9) == ROBOTIQ_CLOSED # binarized above 0.5 + assert gripper_for(0.1) == ROBOTIQ_OPEN + + +def test_action_reads_velocities_from_object_or_wire_dict(): + current = np.zeros(NUM_ARM_JOINTS, dtype=np.float32) + vel = np.linspace(-0.2, 0.2, NUM_ARM_JOINTS, dtype=np.float32) + from_obj = positronic_action_to_molmo({'robot_command': _FakeJointDelta(vel), 'target_grip': 0.0}, current) + from_dict = positronic_action_to_molmo({'robot_command': {'velocities': vel}, 'target_grip': 0.0}, current) + assert np.allclose(from_obj['arm'], from_dict['arm']) + # Bytes-keyed wire form (msgpack deserialisation on some client versions keys with bytes). + from_bytes = positronic_action_to_molmo({b'robot_command': {b'velocities': vel}, b'target_grip': 1.0}, current) + assert np.allclose(from_obj['arm'], from_bytes['arm']) + + +def test_action_joint_count_mismatch_raises(): + current = np.zeros(NUM_ARM_JOINTS, dtype=np.float32) + bad = {'robot_command': _FakeJointDelta(np.zeros(6, dtype=np.float32)), 'target_grip': 0.0} + with pytest.raises(ValueError): + positronic_action_to_molmo(bad, current) + + +# --- FakePolicy ------------------------------------------------------------------------------------------------- + + +def test_fake_policy_chunk_shape_and_range(): + policy = FakePolicy(chunk_size=8, seed=3) + chunk = policy.new_session()({POS_TASK: 't'}) + assert len(chunk) == 8 + for step in chunk: + vel = step['robot_command'].velocities + assert vel.shape == (NUM_ARM_JOINTS,) + assert np.all(np.abs(vel) <= adapter.MAX_JOINT_DELTA + 1e-6) + assert step['target_grip'] in (0.0, 1.0) + + +def test_fake_policy_is_deterministic_per_seed(): + a = FakePolicy(seed=7).new_session()({POS_TASK: 't'}) + b = FakePolicy(seed=7).new_session()({POS_TASK: 't'}) + for sa, sb in zip(a, b, strict=True): + assert np.array_equal(sa['robot_command'].velocities, sb['robot_command'].velocities) + assert sa['target_grip'] == sb['target_grip'] + + +def test_fake_policy_zero_mode_holds(): + chunk = FakePolicy(mode='zero').new_session()({POS_TASK: 't'}) + for step in chunk: + assert np.array_equal(step['robot_command'].velocities, np.zeros(NUM_ARM_JOINTS)) + assert step['target_grip'] == 0.0 + + +# --- ChunkBuffer ------------------------------------------------------------------------------------------------ + + +class _CountingSession: + def __init__(self, chunk_size: int): + self.calls = 0 + self._chunk_size = chunk_size + + def __call__(self, obs): + self.calls += 1 + return [ + { + 'robot_command': _FakeJointDelta(np.full(NUM_ARM_JOINTS, self.calls, dtype=np.float32)), + 'target_grip': 0.0, + } + for _ in range(self._chunk_size) + ] + + +def test_chunk_buffer_replays_one_per_tick_then_requeries(): + session = _CountingSession(chunk_size=3) + buf = ChunkBuffer(session) + first = [buf.next({}) for _ in range(3)] + assert session.calls == 1 # one chunk covered three ticks + fourth = buf.next({}) + assert session.calls == 2 # drained -> re-queried + assert float(first[0]['robot_command'].velocities[0]) == 1.0 + assert float(fourth['robot_command'].velocities[0]) == 2.0 + + +def test_chunk_buffer_empty_chunk_raises(): + buf = ChunkBuffer(lambda obs: []) + with pytest.raises(RuntimeError): + buf.next({}) + + +def test_chunk_buffer_drops_trailing_horizon_marker(): + # A real droid chunk is 8 actions + a window-end entry carrying only `timestamp`; consumed as an action the + # marker KeyErrors on `robot_command`, so the buffer must drop it and re-query rather than replay it. + vel = np.zeros(NUM_ARM_JOINTS, dtype=np.float32) + chunk = [{'robot_command': _FakeJointDelta(vel), 'target_grip': 0.0, 'timestamp': i / 15} for i in range(8)] + chunk.append({'timestamp': 8 / 15}) + calls = {'n': 0} + + def session(obs): + calls['n'] += 1 + return list(chunk) + + buf = ChunkBuffer(session) + for _ in range(8): + assert 'robot_command' in buf.next({}) + assert calls['n'] == 1 + buf.next({}) # 9th tick: marker dropped, buffer re-queries instead of replaying it + assert calls['n'] == 2 + + +# --- end to end (server-free) ----------------------------------------------------------------------------------- + + +def test_end_to_end_fake_pipeline(): + env = _load_env_obs() + client = FakePolicy(mode='random', seed=1, chunk_size=8) + buffer = ChunkBuffer(client.new_session()) + + pos_obs = molmo_obs_to_positronic(env, 'pick up the cube') + current_arm = np.asarray(env['qpos']['arm'], dtype=np.float32) + action = buffer.next(pos_obs) + molmo_action = positronic_action_to_molmo(action, current_arm) + + assert set(molmo_action) == {'arm', 'gripper'} + assert molmo_action['arm'].shape == (NUM_ARM_JOINTS,) + assert molmo_action['gripper'].shape == (1,) + assert molmo_action['gripper'][0] in (ROBOTIQ_OPEN, ROBOTIQ_CLOSED) + expected_arm = current_arm + action['robot_command'].velocities + assert np.allclose(molmo_action['arm'], expected_arm) diff --git a/positronic/simulator/robolab/make_fixture.py b/positronic/simulator/robolab/tests/make_fixture.py similarity index 96% rename from positronic/simulator/robolab/make_fixture.py rename to positronic/simulator/robolab/tests/make_fixture.py index 624fe89de..8811f3ade 100644 --- a/positronic/simulator/robolab/make_fixture.py +++ b/positronic/simulator/robolab/tests/make_fixture.py @@ -11,7 +11,7 @@ episode, not the multi-GB recording. Run once against a RoboLab recording (e.g. the repo's ``examples/recorded_data/RubiksCubeAndBananaTask/data.hdf5``), then commit the ``.npz`` next to the test:: - uv run --no-project positronic/simulator/robolab/make_fixture.py \ + uv run --no-project positronic/simulator/robolab/tests/make_fixture.py \ --demo-path /examples/recorded_data/RubiksCubeAndBananaTask/data.hdf5 \ --out positronic/simulator/robolab/tests/rubiks_cube_and_banana.npz """ diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 9ebf7957a..9e2cc31b6 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -138,7 +138,8 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came # Pretrained DROID models read joints and gripper as separate observation keys and the language -# prompt under `prompt` (see openpi `droid_policy.DroidInputs`). +# prompt under `prompt` (see openpi `droid_policy.DroidInputs`), lowercased — the checkpoints were +# trained on lowercased language and MolmoSpaces' Pi baseline normalizes the same way. droid_obs = cfn.Config( GenericObservationCodec, state={'observation/joint_position': {'robot_state.q': 7}, 'observation/gripper_position': {'grip': 1}}, @@ -147,6 +148,7 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came 'observation/exterior_image_1_left': ('image.exterior', (224, 224)), }, task_field='prompt', + lowercase_task=True, ) ee = codecs.compose.override(obs=ee_obs, action=codecs.absolute_pos_action) diff --git a/pyproject.toml b/pyproject.toml index 21108ce35..d9010f2ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,7 @@ testpaths = [ "positronic/offboard/tests", "positronic/policy/tests", "positronic/simulator/env_server/tests", + "positronic/simulator/molmo_spaces/tests", "positronic/tests", "positronic/utils/tests", "positronic/vendors/gr00t/tests",