Declare MolmoSpaces' grasp frame as the eval control frame (molmo_grasp) - #507
Closed
vertix wants to merge 24 commits into
Closed
Declare MolmoSpaces' grasp frame as the eval control frame (molmo_grasp)#507vertix wants to merge 24 commits into
molmo_grasp)#507vertix wants to merge 24 commits into
Conversation
… frames Convert poses between the robot's canonical `control_frame` and a policy's own EE frame (DROID's `droid_eef` = gripper base composed with `EEF_OFFSET_ROT`). The codec runs client-side: the harness injects the robot model into the local obs and `RemoteSession` keeps it off the wire. Declare RoboLab's `robot_meta` `control_frame` as `droid_eef` truthfully and drop the `#469` HACK.
Our `gripper_base` (MuJoCo Menagerie) does not coincide with RoboLab's `Robotiq_2F_85/base_link`, so routing `droid_eef` through it was wrong. Define it directly on the flange from RoboLab's DROID USD: 18.17mm along Z, +90 deg Z. Verified against the USD at 0.000mm / 0.00000 quaternion difference.
The training dual moves both poses into `to`; relabel `control_frame` too so a downstream `IKJointsAction` solves the command pose against the matching site instead of the canonical one. Document that context-reading action decoders (`RelativePositionAction`) are not yet frame-aware.
Joint-only datasets carry no `robot_command.pose`; registering a derived signal for it made the LeRobot conversion dereference a missing key. Derive a pose transform only when the episode has that signal, so `ee_frame` composed with a joint action still converts the observation pose.
Serve AllenAI's MolmoSpaces MuJoCo benchmark behind the env-server wire so positronic owns the control loop: `env.py` drives a single `BaseMujocoTask` per episode in MolmoSpaces' own venv (pinned checkout, `uv run --project`), the client-side `MolmoAdapter` maps the raw payload into the canonical contract, and `cfg/eval/sim/molmo.py` sweeps a benchmark's episodes into one CLI. Pure wire mappings live in `mapping.py` (framework-free, fixture-tested). Also lowercase the DROID prompt in the observation codec (the checkpoints train on lowercased language) and document the adapter/codec/wire-client separation of responsibilities. Ticket: Positronic-Robotics/internal#91 #refs
`build_trials` injected a random `eval.seed` per trial when the catalog seed was unset, clobbering each benchmark episode's own spec seed in the env (the `episode.seed` fallback never fired) and making runs non-reproducible. Build the trials directly: omit `eval.seed` unless the catalog seed is set, so the env uses the episode's spec seed; an explicit seed still overrides. Ticket: Positronic-Robotics/internal#91 #refs
Address Codex on #504: - `_episode_count` mirrors `load_all_episodes`' two layouts (benchmark.json, else legacy `house_*/episode_*.json`), so a whole-benchmark run no longer fails in the client for legacy dirs. - `ObservationCodec` lowercases the derived training task when `lowercase_task` is set, matching the served prompt — so a DROID codec trains and infers on one text distribution (the same-keys contract). Ticket: Positronic-Robotics/internal#91 #refs
Two problems the CPU smoke surfaced: - `uv sync --extra mujoco` also resolves the `curobo` extra (a CUDA build needing a GPU toolchain, not on the eval path), which fails on a GPU-less box. Install the `mujoco` extra into a venv with `uv pip install -e .[mujoco]` the way MolmoSpaces' own image does, and run env.py with that venv's python. - Importing `json_eval_task_sampler` directly hits a circular import; import `json_eval_runner` first to break it. Ticket: Positronic-Robotics/internal#91 #refs
An empty or mislaid `--eval.benchmark_dir` (no benchmark.json and no legacy `house_*/episode_*.json`) produced an empty episode sweep and a silent zero-trial eval. Raise a clear error instead. Ticket: Positronic-Robotics/internal#91 #refs
The default camera_dict KeyErrored on RandCam benchmarks, which record the exterior as `randomized_zed2_analogue_1` (verified: camera_configs.py, and the `--camera_names` exterior in mb-bench.md), not `exo_camera_1` or the light-randomization variant. Add it to the exterior fallback so the default mapping covers the light-randomization and RandCam suites alike. Ticket: Positronic-Robotics/internal#91 #refs
MolmoSpaces' renderer hardcodes a macOS CGL context on the CPU render path and dlopens Apple's OpenGL.framework, crashing at renderer init on Linux — so the launcher's advertised CPU rendering (MUJOCO_GL=osmesa / mesa software EGL) died before the first observation. Stub the no-op CGL module so it resolves; untouched on a GPU box. Validated end-to-end: a FrankaPickDroidMiniBench episode drives reset + steps on CPU (mesa software EGL), obs payload + camera frames intact. Ticket: Positronic-Robotics/internal#91 #refs
- Disable MolmoSpaces' internal step horizon (Codex): `JsonBenchmarkEvalConfig` defaults it to 500 steps (~33s at the 66ms policy period), which would self-terminate the task before the harness timeout and truncate the score — the env stealing the loop's authority over episode length. `task_horizon=None` runs to an infinite horizon; `is_done` then reports only the task's own success, and positronic's harness stops the trial at its `--eval.timeout`. - Add `e2e.py`: the launcher spawns the env server and a client resets + steps it over the real socket, mapping each frame through `MolmoAdapter`. Validated end-to-end on CPU (mesa software EGL) against FrankaPickDroidMiniBench. Ticket: Positronic-Robotics/internal#91 #refs
Disabling MolmoSpaces' step horizon (prev commit) combined with the default `terminate_upon_success=False` meant a successful rollout that keeps sending joint commands never reported `done`, so the harness scored it as a timeout with no `eval.success` (Codex P1). `env.py` now ends the episode on the task's own `judge_success()` (end-on-success, the benchmark's scoring) or any molmo terminal; positronic's harness timeout stays the outer deadline. Re-validated end-to-end over the socket. Ticket: Positronic-Robotics/internal#91 #refs
- Annotate `_molmo_eval`'s params + return and add `-> None` to the env/adapter constructors — the config function was untyped (basedpyright's standard mode doesn't flag missing annotations, so the ratchet didn't catch it). - Grandfather env.py + make_fixture.py's type errors in the basedpyright baseline, matching robolab/libero: env.py runs in molmo's own venv, so its foreign imports and dynamic mujoco attrs can't resolve against positronic's deps, and make_fixture hits numpy stub quirks. The ratchet is clean (0 new errors); mapping/adapter/e2e/catalog stay fully typed (not baselined). Ticket: Positronic-Robotics/internal#91 #refs
Reverse the baseline entries the previous commit added — new code is never baseline-eligible. env.py runs in MolmoSpaces' own venv (its molmo_spaces + PYTHONPATH-relative server/mapping imports don't resolve against positronic's deps), so exclude it like vendors/ — checked in its own interpreter, not grandfathered here. make_fixture's numpy typing is fixed properly (return type + a targeted savez ignore). basedpyright: 0 new errors, baseline pristine. Ticket: Positronic-Robotics/internal#91 #refs
`env.py` reports `robot_state.ee_pose` at MolmoSpaces' grasp site, but the supplied `bundled_franka_model` declares the physical FR3 flange as its control frame — so recorded episodes mislabel the pose frame and offline IK over them would solve the wrong target, the same known issue as robolab (#483). The eval is unaffected (the pi05 DROID policy is joint-driven, not eef-pose driven). A grasp-site DROID model would fix both. Ticket: Positronic-Robotics/internal#91 #refs
`task.get_task_description()` reconstructs the instruction per task type upstream (OpeningTask emits "Open the ..." even for a close episode), so it can diverge from the benchmark's authoritative goal. Read the prompt straight from `episode.language.task_description` — the same field JsonEvalTaskSampler itself uses — so recorded episodes and the served prompt always carry the benchmark's own language. Ticket: Positronic-Robotics/internal#91 #refs
env.py was excluded from basedpyright wholesale — a blanket file grandfather that also silenced the missing-annotation and wrong-type errors that are ours to fix. Check it: type the foreign-venv attributes (`_task`, `_robot_view`, ...) as `Any` so their member access is honest, scope-suppress just `reportMissingImports` file-wide (the only category that can fire in this positronic-free module — its imports are exclusively the molmo venv's), and carry targeted per-line `reportAttributeAccessIssue` ignores for the server symbols, the mujoco C-binding stub gap, and the dynamic CGL module stub. basedpyright is clean with the file included; the ratchet baseline is unchanged. Ticket: Positronic-Robotics/internal#91 #refs
#485 (`droid_eef` + `ChangeEEFrame`) predates this repo's basedpyright ratchet and carries no baseline, so stacking it here surfaces 13 unbaselined errors in its frame code: 11 are mujoco C-binding attribute accesses (`MjData`, `mj_forward`, `mj_name2id`, `mjtObj`, `mju_mat2Quat`) in `frame_transform` and the `_fk_site` test helper — the same foreign-binding class `env.py` already suppresses per line; 2 are loose types in `encode`'s wire signature and the `Derive(**derived)` unpacking. Each gets a targeted per-line `pyright: ignore`. These belong to #485's own review; this commit only keeps the stacked tree ratchet-green and drops out once #485 rebases onto main and resolves them. Ticket: #483 #refs
MolmoSpaces' env reports and accepts Cartesian poses at its arm move group's grasp site (`gripper/grasp_site`, the 2F-85 TCP), not the flange. Add it as a pure-frame site on the flange, measured off MolmoSpaces' own DROID model (`robots/franka_droid/model.xml` @ allenai/molmospaces c2f1b58): 155mm along the flange Z, no rotation — a deeper, unrotated frame than `droid_eef` (18.17mm, +90deg Z). The validation test pins the site to reproduce that measured frame. Ticket: #483 #refs
`env.py` reports `robot_state.ee_pose` at the arm move group's grasp site (`arm.leaf_frame_to_world` = `gripper/grasp_site`), so declare the served `static_meta`'s `control_frame` as `molmo_grasp` instead of the FR3 flange. Recorded episodes and offline IK then live in the frame the env actually uses, not mislabelled by the 155mm flange->grasp offset. Mirrors #485's robolab `droid_eef` flip. Removes the HACK; leaves a TODO(#483) for the per-policy frame-ownership design. Ticket: #483 #refs
6 tasks
Contributor
Author
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e0e822890
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Both surfaced by #507 routing molmo's frame through #485's `ChangeEEFrame`. 1. Read the EE frame from `static_meta` too (harness.py). `_build_obs` read the frame metadata only from the live `robot_meta` channel, but an env may ship its model client-side via `Embodiment.static_meta` while emitting `robot_meta {}` — exactly molmo (`control_frame='molmo_grasp'` in static_meta). So a molmo run with `ChangeEEFrame` got no `urdf`/`control_frame` and KeyError'd on the first obs. Merge static_meta under the live `robot_meta` (which wins), mirroring `_build_episode_meta`'s own merge order. 2. Skip null pose aliases in `ChangeEEFrame.training_encoder` (codec.py). `_RENAME_ROBOT_COMMAND` aliases a joint-only dataset's absent legacy pose to a PRESENT `robot_command.pose` whose value is `None` (`Get(..., None)`), so the `key in episode` guard wrapped `None` in a derived pose signal and joint-only training failed on dereference. Guard on the value being a real signal too. Regression tests fail on the pre-fix code, pass after: a static-meta-frame obs test and a live-robot_meta-wins test (test_harness.py); a null-alias skip test (test_change_ee_frame.py). Ticket: #483 #refs
Contributor
Author
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
v-positronic
force-pushed
the
molmo-env-server
branch
from
August 4, 2026 18:30
b389088 to
2c6b9db
Compare
v-positronic
force-pushed
the
molmo-env-server
branch
from
August 4, 2026 19:56
bd98eb6 to
56bf282
Compare
vertix
force-pushed
the
molmo-env-server
branch
from
August 10, 2026 08:28
76d863b to
dd9536e
Compare
Collaborator
|
Closing as absorbed by #504. While this sat,
So a rebase drops to an empty diff, which is why this closes rather than retargets. One residual is carried out rather than closed with it. That fix is tracked as |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The MolmoSpaces instance of #483 — Per-policy EE control frames.
env.pyreportsrobot_state.ee_poseat MolmoSpaces' arm move-group grasp site (arm.leaf_frame_to_world=gripper/grasp_site, the 2F-85 TCP), but #504 currently serves arobot_metadeclaring the FR3 flange — so recorded episodes mislabel the pose frame by the flange-to-grasp offset and offline IK over them solves the wrong target. This adds the grasp frame as amolmo_graspsite on the bundled franka model and declares it as the servedcontrol_frame, mirroring the robolabdroid_eefflip in #485. It does not affect the joint-space pi05 DROID eval (joint commands + privileged success), only the recorded pose frame / offline IK / 3D viewer.Stacked on TWO open PRs — depends on both
molmo-env-server).ChangeEEFrame(to=)codec anddroid_eefsite for per-policy EE frames — merged into this branch (its frame machinery + pure-frame-site support are prerequisites).The base is
molmo-env-server; #485 is merged in so the tree builds and tests against both. #485's lines (codec.py,ik.py,robolab/launcher.py,test_change_ee_frame.py, ...) appear in this diff until #485 lands, then drop out, leaving onlymolmo.py+models.py+ the molmo validation test.The frame — measured against MolmoSpaces' own model, not assumed
Traced
arm.leaf_frame_to_worldfor the DROID rig: thearmmove group (FrankaFR3ArmGroup) is built withgrasp_site_name="gripper/grasp_site", so the reported pose is thegripper/grasp_sitesite of the attached Robotiq 2F-85. Loadedrobots/franka_droid/model.xml@allenai/molmospaces@c2f1b58in MuJoCo and measured that site relative to the flange (positroniclink8= molmo'sattachment_site, bothlink7 + (0 0 0.107)identity — as #485 established forlink8):molmo_grasp= (0, 0, 0.155) m, identity rotation — constant across joint configs (rigid).droid_eef(18.17 mm Z, +90 deg Z): delta 136.8 mm translation, 90.0 deg rotation. So a new site, not a reuse.Validation: the new
molmo_graspsite reproduces that measured frame at 0.000 mm / 0 quat (test_molmo_grasp_matches_molmospaces_grasp_site, viaframe_transform(bundled_franka_model()['urdf'], 'link8', 'molmo_grasp')).Changes
drivers/roboarm/models.py— addmolmo_graspas a pure-frame site on the flange (mirrorsdroid_eef).cfg/eval/sim/molmo.py—static_metadeclarescontrol_frame='molmo_grasp'; the# HACKbecomes a# TODO(#483)for the per-policy frame-ownership design (still restated in the serving config here).drivers/roboarm/tests/test_ik.py— themolmo_graspvalidation test.Note: the basedpyright ratchet + #485
#485 predates this repo's type ratchet (it carries no
baseline.json), so stacking it here surfaces 13 unbaselined errors in #485's frame code — 11 mujoco C-binding attribute accesses (the same foreign-binding classenv.pyalready suppresses per line) and 2 loose types (encode's wire signature,Derive(**derived)unpacking). The first commit,Type-clean #485's frame code against the type ratchet, suppresses them with targeted per-linepyright: ignores so the stacked tree is ratchet-green. This is inherited #485 debt that belongs to #485's own review and drops out when #485 rebases onto main. My new code (molmo_graspsite, molmo config, the test) adds zero baseline entries.Test plan
uv run --locked pytest --no-cov positronic/drivers/roboarm/tests positronic/policy/tests positronic/simulator/molmo_spaces-> 159 passed. ruff + format + basedpyright ratchet clean.Do not merge — pending review, stacked on #504 and #485.