From 2a9bcd4d5132bc7dd35da9f58d0db154cb57865b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 19 Jul 2026 20:44:36 +0300 Subject: [PATCH] Add `droid_jointpos` OpenPI codec --- .../env_server/tests/test_remote_env.py | 65 +++++++++++++++++++ positronic/vendors/openpi/README.md | 11 +++- positronic/vendors/openpi/codecs.py | 12 ++++ positronic/vendors/openpi/server.py | 16 ++++- 4 files changed, 102 insertions(+), 2 deletions(-) diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index aa306875e..37d3b148e 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -11,6 +11,8 @@ from positronic.drivers.roboarm import command as roboarm_command from positronic.eval import Task from positronic.inference import main +from positronic.policy import Policy, Session +from positronic.policy.codec import ActionTimestamp from positronic.policy.tests.test_harness import StubPolicy from positronic.policy.wrappers import ChunkedSchedule from positronic.simulator.env_server.adapter import EnvAdapter @@ -242,6 +244,69 @@ def test_remote_eval_runs_to_timeout_without_done(env_server, tmp_path): assert 'sim_state.mjSTATE_INTEGRATION' in signals +class _JointposChunks(Policy): + """Chunks exactly as long as the intended open-loop cadence; ``target_grip`` encodes + ``chunk * 100 + step`` so the recorded wire signals show which actions executed, and when.""" + + def __init__(self, command: roboarm_command.CommandType, chunk_len: int): + self.command = command + self.chunk_len = chunk_len + self.chunks = 0 + + def new_session(self, context=None): + return _JointposChunkSession(self) + + +class _JointposChunkSession(Session): + def __init__(self, policy: _JointposChunks): + self._policy = policy + + def __call__(self, obs): + self._policy.chunks += 1 + return [ + {'robot_command': self._policy.command, 'target_grip': self._policy.chunks * 100.0 + i} + for i in range(self._policy.chunk_len) + ] + + +@pytest.mark.timeout(60.0) +def test_full_chunk_executes_between_replans(env_server, tmp_path): + """The recording proves the contract the DROID jointpos codec makes with RoboLab's client: every action + of every chunk lands on the wire — including the final one, which ``ActionTimestamp``'s validity + sentinel gives a full period before ``ChunkedSchedule`` re-infers — and replans arrive exactly + ``chunk_len`` control periods apart.""" + host, port = env_server + probe = make_mujoco_env([]) + control_dt = probe.reset(0)['control_dt'] + probe.close() + + chunk_len = 5 + raw = _JointposChunks(roboarm_command.JointPosition(np.zeros(7)), chunk_len) + policy = ActionTimestamp(fps=1.0 / control_dt).wrap(raw) + with pos3.mirror(): + ev = remote_stack_cubes_eval(host, port, camera_dict=CAMERAS) + ev.task.timeout = 20 * control_dt + main( + policy=policy, + evals=[replace(ev, trials=[{'eval.trial_index': 0, 'eval.seed': 100}])], + output_dir=str(tmp_path), + wrap=ChunkedSchedule(), + ) + + grip = LocalDataset(tmp_path)[0].signals['target_grip'] + executed = [(float(v), int(ts)) for v, ts in (grip[i] for i in range(len(grip)))] + values = [v for v, _ in executed if v >= 100.0] # the inter-episode home command emits 0.0 + complete_chunks = raw.chunks - 1 # the deadline cuts the last chunk short + assert complete_chunks >= 2 + expected = [c * 100.0 + i for c in range(1, complete_chunks + 1) for i in range(chunk_len)] + assert values[: len(expected)] == expected + + starts = [ts for v, ts in executed if v >= 100.0 and v % 100 == 0] + period_ns = chunk_len * control_dt * 1e9 + for earlier, later in zip(starts, starts[1:], strict=False): + assert later - earlier == pytest.approx(period_ns, abs=period_ns / (2 * chunk_len)) + + @pytest.mark.timeout(60.0) def test_server_failure_crosses_as_error_frame(env_server): """A command the env rejects comes back as an error the client re-raises — the connection survives diff --git a/positronic/vendors/openpi/README.md b/positronic/vendors/openpi/README.md index 38dca23df..02ee2ea73 100644 --- a/positronic/vendors/openpi/README.md +++ b/positronic/vendors/openpi/README.md @@ -16,12 +16,14 @@ OpenPI supports multiple codecs for different use cases: | `ee_joints_traj` | EE pose + grip + joints | Absolute EE trajectory (binarized grip) | Trajectory training with joint feedback | | `joints_traj` | Joints + grip (no EE pose) | Absolute joint trajectory (binarized grip) | Pure joint-space trajectory training | | `droid` | Joint positions + grip | Per-step `JointDelta` | Inference with pretrained DROID models | +| `droid_jointpos` | Joint positions + grip | Absolute `JointPosition` (binarized grip) | Inference with DROID jointpos models (RoboLab leaderboard) | **Key notes:** - **`ee`**: The primary codec. Handles both training data generation (LeRobot format) and inference (OpenPI format) automatically. - **`ee_joints`**: Same as `ee` but includes joint positions in the observation for richer state feedback. - **`_traj` variants**: Train on actual robot trajectory instead of commanded targets, with binarized grip signals. - **`droid`**: Inference-only codec for pretrained DROID checkpoints. The model predicts per-step joint velocities; the codec scales each into a `JointDelta` command (grip binarized) and truncates each chunk to DROID's 8-step open-loop horizon. The driver applies each delta to the live measured joints (`set_target_joints(st.q + delta)`), reproducing the DROID controller with no special client wrap. Serve with `droid` and run inference normally. +- **`droid_jointpos`**: Inference-only codec for openpi's `*_droid_jointpos` checkpoints — the policies RoboLab's leaderboard evaluates. The model emits absolute joint-position chunks; the codec decodes each step into a `JointPosition` command (grip binarized at 0.5) and executes the whole chunk before replanning, matching RoboLab's client cadence (`open_loop_horizon` = the model's `action_horizon`). Serve with `droid_jointpos`. ## 1. Prepare Data @@ -110,6 +112,9 @@ docker compose run --rm --service-ports -v ~/checkpoints:/checkpoints openpi-ser # Pretrained DROID model (pi05_droid) — preset codec, config, and public checkpoint docker compose run --rm --service-ports openpi-server droid + +# DROID jointpos model (pi05_droid_jointpos) — the RoboLab leaderboard policy +docker compose run --rm --service-ports openpi-server droid_jointpos ``` The `droid` config serves the public `pi05_droid` checkpoint from @@ -117,9 +122,13 @@ The `droid` config serves the public `pi05_droid` checkpoint from no local checkpoint mount is needed. The server emits per-step `JointDelta` commands (grip binarized); the driver applies each delta to the live joints, so no special client wrap is needed. +The `droid_jointpos` config serves openpi's `pi05_droid_jointpos` checkpoint from +`gs://openpi-assets-simeval/pi05_droid_jointpos` (openpi fetches it itself on first request). The server +emits absolute `JointPosition` chunks executed at RoboLab's leaderboard cadence — see the codec note above. + **Parameters:** - `--codec`: Codec for observation/action encoding (default: `@positronic.vendors.openpi.codecs.ee`). - Available: `ee`, `ee_joints`, `ee_traj`, `ee_joints_traj`, `joints_traj`, `droid` + Available: `ee`, `ee_joints`, `ee_traj`, `ee_joints_traj`, `joints_traj`, `droid`, `droid_jointpos` - `--checkpoints_dir`: Full path to the experiment directory containing checkpoints - `--checkpoint`: (Optional) Specific checkpoint step to load. If omitted, loads the latest checkpoint - `--config_name`: (Optional) OpenPI config name (default: `pi05_positronic_lowmem`) diff --git a/positronic/vendors/openpi/codecs.py b/positronic/vendors/openpi/codecs.py index 3163ec870..9ebf7957a 100644 --- a/positronic/vendors/openpi/codecs.py +++ b/positronic/vendors/openpi/codecs.py @@ -171,6 +171,18 @@ def observation(state_features: dict[str, int], exterior_camera: str, wrist_came # at 15 fps) so the client re-queries every 8 steps instead of playing the full chunk open-loop. droid = codecs.compose.override(obs=droid_obs, action=codecs.joint_delta_action, horizon=8 / 15) +# The DROID jointpos models (openpi `*_droid_jointpos` configs — the RoboLab leaderboard policies): the +# server returns absolute joint-position chunks ``(action_horizon, 8)`` and RoboLab's client +# (``policies/pi0_family/client.py``) executes the whole chunk before re-querying, gripper binarized at +# 0.5 — its ``open_loop_horizon`` defaults equal each variant's ``action_horizon`` (pi05 = 15, pi0 = 10). +# No ``horizon`` here: the timestamp codec's validity sentinel closes the chunk, so re-inference lands +# after the full chunk executes, whatever each variant's length. +droid_jointpos = codecs.compose.override( + obs=droid_obs, + action=codecs.absolute_joints_action.override(tgt_joints_key='robot_state.q', tgt_grip_key='grip'), + binarize_grip=('grip',), +) + class PoseDeltaAction(Codec): """Decodes pi05_libero's OSC pose-delta chunk into per-step end-effector ``CartesianDelta`` (inference only). diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 37dc7c72c..f25d09e01 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -364,6 +364,13 @@ def server( config_name='pi05_droid', checkpoints_dir='s3://PUBLIC@positronic-public/checkpoints/openpi/pi05_droid/', ) +# The RoboLab leaderboard policy: openpi's DROID jointpos model, served from the checkpoint their +# ``policies/pi0_family/README.md`` recipe pins (pass-through mode — openpi fetches gs:// itself). +droid_jointpos = server.override( + codec=codecs.droid_jointpos, + config_name='pi05_droid_jointpos', + checkpoints_dir='gs://openpi-assets-simeval/pi05_droid_jointpos', +) libero = server.override( codec=codecs.libero, config_name='pi05_libero', checkpoints_dir='gs://openpi-assets/checkpoints/pi05_libero' ) @@ -373,4 +380,11 @@ def server( init_logging() ensure_paligemma_tokenizer() with pos3.mirror(): - cfn.cli({'serve': server, 'phail': phail, 'sim_stack': sim_stack, 'droid': droid, 'libero': libero}) + cfn.cli({ + 'serve': server, + 'phail': phail, + 'sim_stack': sim_stack, + 'droid': droid, + 'droid_jointpos': droid_jointpos, + 'libero': libero, + })