Span OpenPI LR schedule across the full run via `--lr-schedule.decay-steps` by vertix · Pull Request #391 · Positronic-Robotics/positronic

vertix · 2026-05-17T07:37:51Z

Summary

OpenPI's default CosineDecaySchedule hardcodes decay_steps=30_000 independent of --num-train-steps. For runs longer than 30k steps the LR collapses to decay_lr (~2.5e-6) and stays flat for the rest of training; shorter runs never reach the floor. OpenPI also never logged the learning rate, so the problem was invisible.

This PR is now a minimal 3-line change to positronic/vendors/openpi/train.py: when --num_train_steps is set, it also passes --lr-schedule.decay-steps=<num_train_steps> to scripts/train.py so the cosine schedule spans the whole run. lr_schedule is already CLI-exposed by OpenPI, so no wrapper is needed.

LR logging is handled upstream by merged openpi PR Positronic-Robotics/openpi#6 (logs learning_rate into the same wandb.log payload natively). This supersedes the earlier approach in this branch's history (a _launch.py launcher that monkeypatched wandb and dataclasses.replaced the schedule) — that file has been removed; net diff vs main is just the 3 lines in train.py.

Depends on: the deployed positro/openpi image / openpi-base containing openpi#6 (merged) for learning_rate to appear in wandb.

Test plan

py_compile + ruff clean.
Decay-span mechanism empirically confirmed earlier: a Nebius run with decay_steps tied to num_train_steps (=2000) produced the expected warmup→cosine curve — peak 2.50e-5 at step 1000, decaying to ~3e-6 by step 1900 (→ decay_lr at 2000), vs. the old bug where step-1900 LR would still be ~2.4e-5. --lr-schedule.decay-steps=<n> drives the identical optax schedule.
Wrapper-free revalidation passed on Nebius (run cxtxi02s, image off new openpi-base with openpi#6, no _launch.py): learning_rate logged natively; curve warms to peak 2.50e-5 at step 1000 then cosine-decays to ~3e-6 by step 1900 (→ decay_lr at 2000), matching the earlier wrapper-based run.

Known limitation

With --lr-schedule.decay-steps=num_train_steps, optax's warmup_cosine_decay_schedule computes the cosine span as decay_steps - warmup_steps (default warmup 1000). If num_train_steps <= warmup_steps this is non-positive and optax raises ValueError: cosine_decay_schedule requires positive decay_steps. Real runs are far longer than 1000 steps so this is not blocking, but very short smoke runs will crash.

openpi PR Positronic-Robotics#6 (merged) logs `learning_rate` natively, so the wandb monkeypatch is unnecessary. The decay-steps fix needs no wrapper either: `lr_schedule` is CLI-exposed, so pass `--lr-schedule.decay-steps` directly.

The squash rebase restored stage-2 file versions wholesale from the backup tag, overwriting upstream/main commits that landed between the stage-2 branch point and merge: - Re-apply Positronic-Robotics#392 idle_timeout_min auto-shutdown: - vendor_server.py: ctor param, _active_sessions/_last_activity tracking, _idle_watchdog, serve() spawns watchdog - dreamzero/gr00t/openpi server: passthrough on InferenceServer and cfn.config wrappers - Restore Positronic-Robotics#397 keyboard p -> Directive.FINISH() (was reverted to STOP) - Reset to upstream/main (stage-2 never touched these): - vendors/openpi/train.py (Positronic-Robotics#391 lr-schedule.decay-steps) - vendors/lerobot/server.py and lerobot_0_3_3/server.py (Positronic-Robotics#392) - utilities/fake_dataset_generator.py (Positronic-Robotics#400 simulate_timeout rename)

…/Session split (#381) * Stage 2 WIP: Policy/Session split + PolicyWrapper composition (squashed) Squashed rebase of worktree-inference-api onto upstream/main (post-stack1 merge in #393). Original 15-commit history preserved in tag backup-worktree-inference-api-pre-rebase. Includes: - Policy -> Policy + Session split (new_session, Session ABC) - PolicyWrapper with | composition, absolute timestamps - Refactor, defensive-coding reductions, vendor adapters - Stack1 reversion of golden_pipeline NOT applied (keep upstream's) Does NOT include the latest stack1 cancellation/Timestamped-gating fixes to harness.py and ds_writer_agent.py - those need to be re-applied. * Adapt golden pipeline test to Stage 2 Policy/Session API - ScriptedProportionalPolicy now creates a Session via new_session() - Drop simulate_inference (no longer a Harness concern; ChunkedSchedule anchors trajectories to clock.now() after inference returns) - Drop to_wire(): codecs now pass robot_command objects through directly - Regenerate golden_pipeline.json.gz against the new pipeline * Address Codex review: DHGripper trajectory + Timestamped guard P1 (positronic/drivers/gripper/dh.py): Harness now emits target_grip as a trajectory list [(ts_ns, grip), ...]. DHGripper was reading .value as a scalar, which would raise TypeError once a trajectory arrived. Switch to the TrajectoryPlayer pattern (matching Robotiq2F). P2 (positronic/dataset/ds_writer_agent.py): Restore the Timestamped guard on the list dispatch. Without it, a serializer returning a plain list (legitimate list-valued payload) would crash with AttributeError on sample.value. Empty list still matches (cancel signal). Add regression test test_serializer_plain_list_value. * Restore upstream changes lost during stage-2 squash The squash rebase restored stage-2 file versions wholesale from the backup tag, overwriting upstream/main commits that landed between the stage-2 branch point and merge: - Re-apply #392 idle_timeout_min auto-shutdown: - vendor_server.py: ctor param, _active_sessions/_last_activity tracking, _idle_watchdog, serve() spawns watchdog - dreamzero/gr00t/openpi server: passthrough on InferenceServer and cfn.config wrappers - Restore #397 keyboard p -> Directive.FINISH() (was reverted to STOP) - Reset to upstream/main (stage-2 never touched these): - vendors/openpi/train.py (#391 lr-schedule.decay-steps) - vendors/lerobot/server.py and lerobot_0_3_3/server.py (#392) - utilities/fake_dataset_generator.py (#400 simulate_timeout rename) * Address second Codex review round - P1: TrajectoryOverrideSerializer([]) now clears _buffer (was returning [] without dropping the buffered tail, so STOP/FINISH would commit canceled waypoints). - P1: RemoteSession normalizes single-dict server responses into a 1-element list to honor the Session list[dict]|None contract. - P2: Add Session.cancel() with DelegatingSession passthrough; wire ChunkedSchedule.cancel() to clear _trajectory_end; have ErrorRecovery call inner.cancel() on error entry. Prevents post-recovery stall on stale trajectory_end with the default ErrorRecovery|ChunkedSchedule. Regenerate golden_pipeline.json.gz (cancel-on-error changes the recorded trajectory across the error/recovery boundary, by design). Adds regression tests for all three. * Address third Codex review round - P1: STOP directive now emits [] on robot_commands/target_grip to clear driver TrajectoryPlayer buffers (was leaving drivers to keep playing preloaded chunks — safety regression for long trajectories). - P1: Extract _cancel_trajectories() helper and call it before STOP_EPISODE in FINISH (and on RUN-restart), so the recording's TrajectoryOverrideSerializer drops its tail before flush() commits canceled waypoints into the episode. - P2: ActionHorizon.decode uses d.get('timestamp', 0.0) so untimestamped chunk shapes (offboard tests, servers that truncate before stamping) do not raise KeyError. _cancel_trajectories also calls self._session.cancel() so the wrapper chain (ChunkedSchedule) drops stale trajectory_end too. Adds focused regression tests for STOP and FINISH ordering. * Restore server metadata in RemotePolicy.meta SampledPolicy._get_keys and inference._seed_sampler read keys from Policy.meta (default key_field='server.checkpoint_path'). With RemotePolicy.meta returning only {'type': 'remote'}, sampled remote policies fell back to index keys ('0'/'1') and the seeding/balancing counts recorded under checkpoint-path keys could not be matched. Lazily open + immediately close a ws session on first .meta access to fetch and cache server metadata (and reuse the cache when an actual new_session is later opened). Adds regression test test_remote_policy_meta_exposes_server_fields. * Revert accidental diffs back to upstream/main Three files differed from upstream/main only by re-implementations or comment rewording — not by Stage 2 design changes: - positronic/dataset/ds_writer_agent.py: pure comment rewording of the same Timestamped guard / empty-cancel fix that #393 already landed. - positronic/drivers/gripper/dh.py: Stage 2 never touched this file; the restored version reimplemented upstream's trajectory handling in a slightly different style. Upstream's #393 version is canonical. - positronic/dataset/tests/test_ds_writer_agent.py: my added tests duplicate upstream's existing test_serializer_plain_list_value and test_trajectory_override_serializer_empty_cancels_buffer. * Adopt upstream's STOP/FINISH cancel tests verbatim My versions of test_stop_cancels_in_flight_trajectory and test_finish_cancels_buffered_trajectory_before_stop_episode were re-implementations of the same tests already in upstream/main #393. Replace them with upstream's verbatim versions to remove the duplication; both work unchanged against the Stage 2 harness. * Restore simulate_inference and preserve upstream golden byte-for-byte Bringing back the inference-latency simulation the squash dropped, and fixing two follow-on issues so the upstream golden file passes unchanged: - Re-add Harness.simulate_inference: bool | float = False. Sleep is yielded after session() returns; harness post-shifts the chunks absolute timestamps and bumps the inner ChunkedSchedule._Session _trajectory_end so the scheduling gate stays consistent with the shifted trajectory. - Skip the latency Sleep when the only action is a Recover command (emergency emit from ErrorRecovery, never a model-driven chunk). - Guard target_grip.emit with if grip_traj so emitting a Recover chunk (which carries no grip) does not clear the grippers TrajectoryPlayer - matches stack1s harness behavior. - Drop to_wire(...) wrapping in DreamZeroActionCodec so action decoders return command.JointPosition objects directly (addresses Codex P1: stage-2 drivers match on Command classes, not wire dicts). - inference.py: pass simulate_inference through main_sim again; remove stale DsWriterCommandType import. Upstream's golden_pipeline.json.gz now passes byte-for-byte with only the mandatory Policy/Session API adaptation in test_golden_pipeline.py. * Address fourth Codex review round - `vendor_server.warmup`: close warmup session in `finally`, not after the awaited `infer` call. Backends with per-session sockets (DreamZero, RemotePolicy) were leaking on warmup failure since the broad `except` skipped `session.close()`. - `Harness._build_episode_meta`: restore upstream's `inference.simulate_inference` entry that the stage-2 squash dropped; recordings made with simulated latency need this in static meta to stay distinguishable from real-time runs. * Cancel grip buffer on empty action chunk `Harness._step` emitted `[]` only on `robot_commands` when the session returned an empty chunk; `target_grip` was skipped because the existing guard treated any empty `grip_traj` as "no grip targets in this chunk". For session-level cancel (empty `actions`), that left the gripper's `TrajectoryPlayer` draining stale waypoints — a partial cancel. Now emit `[]` on `target_grip` too when `actions` is empty, while still skipping the grip emit for non-empty chunks that simply don't carry a `target_grip` key. * Decode pre-PR server commands (temporary compat shim) Pre-PR vendor codecs put commands on the wire as unwrapped `to_wire()` dicts (no `__cmd__` envelope). The new client used to bridge these via `from_wire()` in the harness, but that step is gone — so a plain dict reaches the driver and the command `match` falls through. Recognize the legacy shape in `_unpack` by sniffing `type` against the known command TYPE constants and reconstructing via `from_wire(...)`. Marked with `TODO(remove-pre-PR-server-compat)` for easy grep when all deployed inference servers have been rebuilt. The new `__cmd__` envelope branch was made idempotent: the legacy shim fires bottom-up (msgpack `object_hook` order) and decodes the inner `to_wire` payload before the outer hook runs, so the outer branch now accepts either a Command or a dict. * Record inference with composable `Recorder` taps Replace the per-session `RecordingWrapper` with a `Recorder` that hands out `tap(name)` wrappers writing one rerun `.rrd` per episode. Two taps (`raw` | codec | `server`) capture both the robot and server boundaries in one correlated recording; a single `raw` tap when there is no codec. Move the rerun helpers from `codec.py` into a new `policy/recording.py`. Log action chunks structure-of-arrays as `rr.Tensor`s at one timestamp, with per-action `timestamp` stored as int64 nanoseconds. * Cancel gripper on recovery and trajectories before shutdown STOP Entering recovery emits a Recover-only chunk with no `target_grip`, so the gripper kept draining the interrupted chunk's waypoints while the arm recovered. Emit an empty grip trajectory on the recovery transition so all in-flight motion stops, not just the arm. Shutdown while recording emitted `STOP` before `_cancel_trajectories()`, so `STOP_EPISODE` flushed the unexecuted chunk tail into the recording. Cancel first, matching the FINISH/RUN paths. Regenerate the golden: the recovery fix holds the gripper instead of advancing one buffered waypoint (single `grip` sample differs; all other signals identical). * Isolate server recordings per websocket session A single shared `Recorder` wrapped every websocket session, so a second client connecting before the first disconnected logged into the first client's `.rrd` and shared recorder state. Create a fresh `Recorder` per session instead. Add a process-wide episode counter so concurrent recorders never collide on filenames. * Merge `policy.meta` into episode and handshake metadata The harness episode meta and the server ready-handshake now layer `policy.meta` (the static baseline the wrapped policy aggregates) beneath `session.meta`, which overlays per-episode specifics and wins on conflict. A policy that exposes static fields only via `Policy.meta` no longer loses them once wrapped, without adding a per-session meta-merging layer. * Drop two low-value recording tests (trivial delegation + internal step counter) * Merge `policy.meta` into basic-server ready handshake * Drop un-executed trajectory tail at episode end via flush cutoff `TrajectoryOverrideSerializer.flush(now_ns)` commits only points already due (`ts <= now_ns`) and drops the future-scheduled tail; `DsWriterAgent` passes the episode-end time on `STOP_EPISODE`. Previously the writer flushed on STOP before consuming the harness cancel, so an episode finished mid-chunk committed the un-executed tail. This makes the drop deterministic and consistent with how mid-episode overrides already truncate chunks. * Extract `EpisodeCounter` and drop `Session.on_episode_complete` Split completion-counting out of `SampledPolicy`: a focused `EpisodeCounter` owns the per-key tally (record/counts/seed-from-disk), samplers become stateless strategies (`sample(keys, context, counts)`, no `Sampler.count`), and the harness fires a generic `on_episode_complete(session, context)` callback instead of a `Session` lifecycle method. Removes `_SampledSession`; moves `group_fields` to `SampledPolicy`/`production`. * Normalize single-action chunks and tag sampled sessions with their key Address two codex review findings: - `ChunkedSchedule` now wraps a bare-dict result into a one-item list and defaults a missing `timestamp` to 0 (matching `ActionHorizon`'s tolerance), so single-action and no-codec paths execute an immediate action instead of raising `TypeError`/`KeyError`. - `SampledPolicy` wraps each session in `_KeyedSession` so the sampler's chosen key reaches `EpisodeCounter.record` even when the sub-policy doesn't expose `key_field` (fallback keys), restoring per-key completion counting. * Record both codec boundaries server-side (`raw` + `inference`) The server recorder tapped a single `inference` point outside the codec, so a recording only showed one boundary. Tap both sides instead: `raw` (obs as received over the wire + final decoded action) and `inference` (the encoded obs the model sees + its raw output), captured in one correlated `.rrd`. Falls back to a single `inference` tap when the server has no codec. * Give `production` policy slots independent `weighted_remote` copies All four slots (`groot/openpi/act/smolvla`) defaulted to the same `weighted_remote` Config object. Overriding them through a parent config (the CLI path, e.g. `--policy=@...production --policy.smolvla.port=8000 --policy.act.port=8001`) mutated that shared object in place, so every slot collapsed to one host:port and the sampler saw N identical keys. `.copy()` per slot keeps them independent; direct `.override()` users (`phail_multiple`, `spoons_ablation`) are unaffected. * Drop canceled trajectory tail on episode SUSPEND A STOP mid-chunk followed by a later FINISH could record waypoints the robot never executed: the harness emitted `SUSPEND` before the `[]` cancel, and the suspended writer skips all inputs, so the cancel never reached the serializers and the buffered tail survived to the STOP_EPISODE flush. - The writer drops each serializer's buffered tail when it handles `SUSPEND` (new `Serializer.cancel()`; `TrajectoryOverrideSerializer` clears its buffer), independent of the racy input path. - `STOP_EPISODE` flushes at the command's emit time (`cmd_msg.ts`) instead of `clock.now_ns()`, so a queued stop doesn't commit points scheduled after it. - The harness emits the cancel before `SUSPEND`, matching RUN/FINISH ordering. * Preserve already-due samples when suspending mid-trajectory Refine the SUSPEND fix: clearing the whole serializer buffer also dropped samples that had already executed before the suspend, not just the future tail. Flush at the suspend time instead (`flush(now_ns)`) — commits due samples, drops only the un-executed tail — and drop the now-unused `Serializer.cancel()`. The harness emits SUSPEND before the `[]` cancel so the writer flushes the due prefix before it starts skipping inputs. * Record the sampled key in `_KeyedSession`, not a late session key `_KeyedSession.meta` merged the chosen key before the inner session meta, so a sub-policy that exposes `key_field` only per session could overwrite the fallback key that was actually sampled. `EpisodeCounter.record()` would then count a key absent from the sampler's `keys`, leaving balanced/resume counts stuck. Merge the sampled key last so it always wins. * Carry selected sub-policy meta in `_KeyedSession` `SampledPolicy.meta` is `{}`, so a sampled sub-policy that exposes static fields only via `Policy.meta` (with an empty `Session.meta`) lost its type / checkpoint / codec fields from episode and handshake metadata. `_KeyedSession` now merges the selected policy's static `meta` under the session meta (sampled key still merged last so completion counting stays correct). * Update docs for the Policy/Session inference API Bring the docs in line with the refactor: - `connect-your-model.md`: rewrite the policy example as `Policy.new_session()` + `Session.__call__` (drop `select_action`/`reset`); fix the serialization snippet to use a session. - `codecs.md`: `BinarizeGrip` -> `BinarizeGripTraining`/`BinarizeGripInference`; document the timing codecs (`ActionTimestamp`/`ActionHorizon`/`ActionTiming`). - `offboard/README.md`, `openpi/README.md`: ready handshake is `{status:'ready', meta}`; action response is a list; openpi client example handles the status handshake; subprocess timeout 120s -> 300s. - `inference.md`: per-subcommand default driver; note `recording_dir` and sampling configs. - vendor READMEs: add the required `serve` subcommand for lerobot/lerobot_0_3_3; fix broken `groot` -> `gr00t` links and a `cGR00T` typo; fix non-existent `openpi_absolute` README example. * Send `robot_state.status` as a plain int for old-server compatibility `RobotStatus` was wired as a `{b'__robotstatus__': ...}` map. A pre-PR server's `object_hook` leaves maps it doesn't recognize undecoded, handing a bytes-keyed dict to server-side codecs/recording (old `RecordingCodec` can crash on it). `robot_state.status` is only read client-side (harness ErrorRecovery, before serialization) and nothing reconstructs it server-side, so encode it as a plain int — old and new servers both tolerate it. * Wire `RobotStatus` as a str-keyed envelope, not bytes Keep the proper enum round-trip for the new client/server (don't degrade it to a bare int forever), while staying safe against a pre-PR server: the envelope key is a plain `str`, so an old server that doesn't decode it leaves a str-keyed dict in the observation. Its recorder runs `key.endswith(...)` over dict keys, which TypeErrors on a bytes key (`b'__robotstatus__'`) but is harmless on a str one. New servers reconstruct the `RobotStatus` enum in `_unpack`. * Accept both str and bytes `__robotstatus__` keys on decode The deserializer now reconstructs `RobotStatus` from either the str-keyed envelope (current wire form, safe for pre-PR server recorders) or the bytes-keyed form. This lets a future change flip the wire to the bytes key — consistent with the other envelopes — without breaking any server already running this version; the str-key path (and the pre-PR-server compat) can then be dropped. * Document the inference execution model in `connect-your-model.md` Make it the canonical API doc: lead with the predictor/real-time-client split, relative→absolute timestamps, trajectory-as-overwritable-command, and client-side scheduling strategies (chunked, temporal ensembling, RTC) before the wire mechanics. Link the timing section in `codecs.md` to it and fix the `timestamp` row to describe chunk collapse, not failure. * Close the per-episode session on `FINISH` `FINISH` finalized recording but left `self._session` open until the next RUN or shutdown, keeping a `RemoteSession`'s websocket alive and delaying the offboard server's per-session cleanup. Close and clear it here, mirroring the RUN/shutdown paths. * Rewrite `connect-your-model.md` for a zero-context reader Lead with why inference is split into a control client and a model server and why the latency logic lives on the client, in plain language. Add the Policy/Session, Codec, and Wrapper concepts, two mermaid diagrams, and a "Debugging with recordings" section; cut duplicated protocol/codec prose; fix Docker-as-prerequisite (it is optional). Stop recommending SmolVLA as a starting point in `model-selection.md` and `README.md` (ACT instead). * Close aborted sessions on `HOME` `HOME` ends the episode (aborts recording, stops running) but left `self._session` open, keeping a `RemoteSession`'s websocket alive until the next RUN or shutdown. Release it on the abort path, mirroring `FINISH`. * Sharpen `connect-your-model.md` and restructure `codecs.md` connect-your-model: drop reader-irrelevant details (credentials, MuJoCo), stop surfacing "chunk" as a concept (server returns actions / a trajectory), rename the diagram node to Robot, frame the action list as the normal-not-only response, and generalize the recordings section around the `Recorder` wrapper. codecs: lead with the goal, then the train/inference dual that makes codecs composable, then observation/action encoding with worked examples (incl. EE→joints via IK); move vendor tables to the end.

vertix mentioned this pull request May 18, 2026

Serverless endpoint auto-shutdown + remote-training skill for Nebius Serverless #392

Merged

4 tasks

vertix added 2 commits May 18, 2026 16:27

Fix OpenPI LR schedule to span full run and log LR to wandb

4174103

Patch Run.log so injected LR survives openpi's wandb.init()

bf5fa6a

vertix mentioned this pull request May 18, 2026

Log learning_rate to wandb during training Positronic-Robotics/openpi#6

Merged

2 tasks

vertix force-pushed the openpi-training-schedule branch from 62a6fac to bb84667 Compare May 18, 2026 17:34

vertix changed the title ~~Fix OpenPI LR schedule to span full run and log LR to wandb~~ Span OpenPI LR schedule across the full run via --lr-schedule.decay-steps May 18, 2026

vertix merged commit 733bd36 into Positronic-Robotics:main May 18, 2026
7 checks passed

vertix deleted the openpi-training-schedule branch May 19, 2026 16:20

Navigation Menu

Uh oh!

Span OpenPI LR schedule across the full run via `--lr-schedule.decay-steps` - #391

Span OpenPI LR schedule across the full run via `--lr-schedule.decay-steps`#391
vertix merged 3 commits into
Positronic-Robotics:mainfrom
vertix:openpi-training-schedule

vertix commented May 17, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

vertix commented May 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Known limitation

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

vertix commented May 17, 2026 •

edited

Loading