diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 33274d34d..cd984ef26 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -1,23 +1,5 @@ { "files": { - "./pimm/core.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 8, - "endColumn": 17, - "lineCount": 1 - } - }, - { - "code": "reportReturnType", - "range": { - "startColumn": 15, - "endColumn": 23, - "lineCount": 1 - } - } - ], "./pimm/shared_memory.py": [ { "code": "reportIndexIssue", @@ -995,22 +977,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 62, - "endColumn": 77, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 96, - "endColumn": 112, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -1285,22 +1251,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 58, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 63, - "endColumn": 70, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1317,14 +1267,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 47, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1587,22 +1529,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 110, - "endColumn": 112, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 56, - "endColumn": 58, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1611,14 +1537,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 57, - "endColumn": 59, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1635,14 +1553,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 110, - "endColumn": 112, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1658,14 +1568,6 @@ "endColumn": 73, "lineCount": 1 } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 106, - "endColumn": 108, - "lineCount": 1 - } } ], "./positronic/dataset/edits.py": [ @@ -2357,22 +2259,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 24, - "endColumn": 28, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 24, - "endColumn": 28, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4125,22 +4011,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 28, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 40, - "endColumn": 44, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4157,14 +4027,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 33, - "endColumn": 52, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4215,22 +4077,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 27, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 39, - "endColumn": 43, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4339,22 +4185,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 38, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 43, - "endColumn": 47, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4371,14 +4201,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 30, - "endColumn": 44, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4403,14 +4225,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 77, - "endColumn": 83, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4419,14 +4233,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 29, - "endColumn": 36, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4847,22 +4653,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 27, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 39, - "endColumn": 43, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4891,38 +4681,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 35, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 24, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 41, - "endColumn": 45, - "lineCount": 1 - } - }, { "code": "reportGeneralTypeIssues", "range": { @@ -5874,91 +5632,26 @@ } } ], - "./positronic/policy/harness.py": [ + "./positronic/policy/harness.py": [], + "./positronic/policy/observation.py": [ { - "code": "reportOptionalCall", + "code": "reportArgumentType", "range": { - "startColumn": 18, - "endColumn": 56, + "startColumn": 78, + "endColumn": 88, "lineCount": 1 } }, { - "code": "reportArgumentType", + "code": "reportIncompatibleMethodOverride", "range": { - "startColumn": 39, - "endColumn": 55, + "startColumn": 8, + "endColumn": 14, "lineCount": 1 } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 20, - "endColumn": 27, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 41, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 59, - "endColumn": 61, - "lineCount": 1 - } - }, - { - "code": "reportOptionalOperand", - "range": { - "startColumn": 65, - "endColumn": 79, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 11, - "endColumn": 40, - "lineCount": 1 - } - } - ], - "./positronic/policy/observation.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 78, - "endColumn": 88, - "lineCount": 1 - } - }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 14, - "lineCount": 1 - } - } - ], - "./positronic/policy/recording.py": [ + } + ], + "./positronic/policy/recording.py": [ { "code": "reportArgumentType", "range": { @@ -6003,38 +5696,6 @@ } ], "./positronic/policy/tests/test_golden_pipeline.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 35, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 19, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 35, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -6076,64 +5737,7 @@ } } ], - "./positronic/policy/tests/test_harness.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 19, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 65, - "endColumn": 73, - "lineCount": 1 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 15, - "endColumn": 32, - "lineCount": 1 - } - }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 35, - "endColumn": 40, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 39, - "endColumn": 44, - "lineCount": 1 - } - } - ], + "./positronic/policy/tests/test_harness.py": [], "./positronic/policy/tests/test_policy_io.py": [ { "code": "reportArgumentType", @@ -6484,16 +6088,6 @@ } } ], - "./positronic/policy/wrappers.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - } - ], "./positronic/probe.py": [ { "code": "reportAttributeAccessIssue", @@ -6643,14 +6237,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 97, - "endColumn": 101, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6659,14 +6245,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 52, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6675,14 +6253,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 53, - "endColumn": 55, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6691,14 +6261,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 53, - "endColumn": 60, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6707,14 +6269,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 95, - "endColumn": 102, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6722,14 +6276,6 @@ "endColumn": 56, "lineCount": 1 } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 57, - "endColumn": 63, - "lineCount": 1 - } } ], "./positronic/server/dataset_utils.py": [ @@ -7258,14 +6804,6 @@ "endColumn": 32, "lineCount": 1 } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 54, - "endColumn": 62, - "lineCount": 1 - } } ], "./positronic/simulator/env_server/server.py": [ @@ -7287,22 +6825,6 @@ } ], "./positronic/simulator/env_server/tests/mujoco_env.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 42, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 43, - "endColumn": 58, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -7441,14 +6963,6 @@ "lineCount": 1 } }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 15, - "endColumn": 33, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -7763,38 +7277,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 45, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 24, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 47, - "endColumn": 51, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -8271,22 +7753,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 41, - "endColumn": 53, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 38, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -8327,14 +7793,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 33, - "endColumn": 48, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -8651,38 +8109,6 @@ "lineCount": 1 } }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 19, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 19, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 32, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 33, - "endColumn": 47, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -8690,14 +8116,6 @@ "endColumn": 45, "lineCount": 1 } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 46, - "endColumn": 53, - "lineCount": 1 - } } ], "./positronic/wire.py": [ diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index bd094b2a7..1d28d359a 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -111,8 +111,11 @@ For each open comment, decide and note severity if the bot tagged one (e.g. Code resolve only if you land a concrete change, else leave open. - **Decline** — wrong, not applicable, or contradicts a deliberate decision → reasoned reply, **leave open**. -- **Defer** — valid but out of scope for this PR → reply (note where it's tracked), - **leave open**. +- **Defer** — valid but out of scope for this PR → reply saying so, **leave open**, and name it in + the Step 6 report as wanting a tracking issue. **Never open an issue without asking first.** An + issue is filed under the user's own account, so it reads to their team as their judgement about + what matters; that call is theirs. Propose the title and the substance and wait for a yes. If an + existing issue already covers it, cite that instead — citing is not filing. - **Discuss** — the reviewer is asking a question or opening a design discussion, not requesting a change → answer it, **leave open** for them to respond. @@ -122,8 +125,16 @@ Present the triage as a short numbered list: comment → verdict → planned fix declines without prompting. "Unclear" is a verdict you reach after the research above, never before it — a question you could have answered by reading the code costs the user a round trip and comes back as *go and find out*, which is the same work plus a delay. Pause and confirm -only for what research cannot settle: a preference, a cost, a scope call, or a change whose -blast radius is genuinely the user's to accept. The user may drop or override any item. +only for what research cannot settle: a preference, a cost, or a change whose blast radius is +genuinely the user's to accept. The user may drop or override any item. + +**Whether a finding belongs in this PR is your call.** "Fix it here or track it for later" is +the verdict this skill exists to reach, and by the time you can pose it the research is done, +so the human has nothing to add that you don't already have. Reach it, reply, and propose the +tracking issue in the Step 6 report, where they can approve or overrule it — filing it is +theirs to authorize, never yours. Bring a deferral back for a decision only when leaving it +undone changes what the PR ships — a gap in the behaviour the user asked for, not a follow-up +the report already names. When you do bring something back, explain it in full and in plain words: what the thing is and how it works, what the reviewer asked for, what you did or did not do and why, and what each @@ -256,10 +267,62 @@ Summarize: - which threads you resolved (fixes only) vs left open (declines / defers / discussion), - any follow-ups the user should track. +Then **offer** the walk below, in one line, and go to Step 7 without waiting: + +> 3 threads are open on my call: , <title>, <title>. Want to go through them? + A bot will re-review on push and may add comments. **Don't hand the watch back to the user** — go to Step 7, which watches for that re-review in the background and loops you through another pass automatically until the reviewer converges. +## Walking your declines — only on an explicit go + +Every thread you declined or deferred is a decision the human has not seen. They are entitled to +review each one, during the cycle or after it converges. **Offer it; never start it uninvited.** A +proposal is one line (above). Presenting a thread before they say go is the failure this section +exists to prevent — as is presenting all of them at once. + +On their go, **one thread per message, no section headings**, covering four things: + +1. what the thing is, +2. what the reviewer wanted, +3. why you declined it, +4. the decision you need — bolded, on its own line, with the options named. + +**Simple beats short.** The failure mode is density, not length: four sentences carrying three ideas +each are worse than twelve carrying one each. Unpack — one idea per sentence, plain words, and +describe the thing before arguing about it. Length is the budget for being understood; spend it on +unpacking and on nothing else. + +**A short code snippet is welcome** where code says it faster than prose — the branch that decides, +the line that was removed. A few lines, and say what it shows. Long listings, file tours and full +diffs are not. + +Then **stop**. Do not append the next thread, a summary, or what you plan to do after. + +Withheld until asked: the measurements, the alternative designs, the options table, the wider diff. +Fine to *have* — `git show`, instrumentation, a longer argument — none of it goes in unasked. + +- **"too long" means cut it in half** — not rewrite it at the same length. +- **"too complicated" / "I don't understand" means unpack it** — more sentences, fewer ideas per + sentence, and a snippet if one exists. Cutting further is the wrong reflex and makes it worse. +- **A follow-up question gets that answer alone.** "why declined?" is not an invitation to re-present + the thread; "which three?" wants three names and nothing else. +- Assume they have not read the code and will not. No "see the call site". + +Their verdicts are terse. What each one means: + +| They say | You do | +|---|---| +| "next" / "resolve, next" | reply, resolve, present the next thread — same shape, no re-asking | +| "leave it" | reply that it is accepted as a known limit, resolve | +| "fix it" | fix, run the gates, commit, push, reply with the SHA, resolve | +| "file it" / "open a bug" | file the issue, reply linking it, resolve — this verdict is the only thing that authorizes filing one | +| "why declined?" | the reasoning only | + +This is the one path on which a declined thread gets resolved: the human closed it, thread by thread. +Absent that, Step 5 stands and it stays open. + ## Step 7: Watch for convergence in the background (so the user doesn't have to) Two signals decide whether a push actually lands the PR, and **both are asynchronous**: CI diff --git a/.claude/skills/check-rules/SKILL.md b/.claude/skills/check-rules/SKILL.md index 4505a9992..be2b7bc9f 100644 --- a/.claude/skills/check-rules/SKILL.md +++ b/.claude/skills/check-rules/SKILL.md @@ -39,6 +39,14 @@ of one cleanly checked. Each needs its rule and a way to tell which code changed not know how that code came to be: a cold-headed reviewer answering one question, not a colleague who sat through the reasoning. +Every prompt is the same four parts: the rule pasted whole and verbatim, the scope Step 1 settled, the +return format below, and the instructions that follow it. Nothing else goes in. + +Whatever you would add comes from knowing the change, and that knowledge is what kept you from seeing +the violation yourself. Emphasis, an example, a "pay attention to", a scope trimmed to the lines you +touched — each hands the agent your reading of the rule in place of the rule, and the reading you did +not have is the one you spawned it for. + **Spawn them on Sonnet** — `model: 'sonnet'`, on every agent. Each one answers a narrow question: one rule against one diff. A fan-out of that width spends most of what it costs on being that wide rather than on the judging. The parameter is not optional: omitted, an agent inherits whatever model the diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 168521c79..a7cea9daa 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -36,9 +36,9 @@ library must supply the tools that make late binding possible — codecs, per-tr projections over raw recordings. **Every decision lives with the party that has the information.** Only a driver knows its motion -capabilities, so drivers plan through waypoints. Only a sensor knows its own cadence, so sensors -run at their own rate instead of a rate the loop imposes. Only a policy knows what its model was -trained on, so translation to model I/O ships with the policy. +capabilities, so drivers own how they reach a commanded setpoint. Only a sensor knows its own +cadence, so sensors run at their own rate instead of a rate the loop imposes. Only a policy knows +what its model was trained on, so translation to model I/O ships with the policy. **Components are functions over flowing data.** A component sees nothing but its inputs and touches nothing but its outputs. Whatever varies enters as data — time is an observation field, hardware @@ -116,21 +116,30 @@ control system its clock, and no component reads time at point of use. Trajector the same time frame the observations carry, so a virtual clock, a slowed sim, or a replayed episode changes nothing downstream. -**Trajectory is the command.** Ownership puts execution with the driver: the policy emits a -trajectory of waypoints with absolute timestamps, and the driver plays it at its own control rate, -planning through the waypoints as well as it knows how. Signals are last-value-wins, so a new -trajectory overwrites the current one — the previous command is merely context for the next. -Continuous-update schemes (RTC, temporal ensembling) therefore need no special mechanism: they are -wrappers that rewrite the command more often. An empty trajectory cancels the channel and the -device holds. +**The wrapper owns the plan, the harness plays it, the driver executes.** A policy speaks in +trajectories — waypoints with absolute timestamps — because a model predicts a horizon, not an +instant. But a trajectory on the wire makes every driver buffer the future, and makes the recording +guess which prefix of that buffer actually ran. So the plan stops at the harness: a command channel +carries the single command due at the moment it is emitted, the driver executes the latest one and +holds otherwise, and emission time *is* execution time. Continuous-update schemes (RTC, temporal +ensembling) therefore need no special mechanism: they are wrappers that hand back a new trajectory +more often, and the harness keeps playing the old one until they do. **The harness stays thin.** It is the one layer standing between any policy and any embodiment, so anything it encodes about either side breaks the any-to-any goal. It assembles the observation -dict, calls the session, demuxes the returned waypoints per command channel, and runs episode -lifecycle — nothing else. Scheduling, blending, history stacking and error recovery live in the -wrapper stack around the policy; a session returning `None` means "keep executing the current +dict, calls the session, plays the returned trajectory one command per channel per round, and runs +episode lifecycle — nothing else. Scheduling, blending, history stacking and error recovery live in +the wrapper stack around the policy; a session returning `None` means "keep executing the current trajectory". +**Inference cost is a fact of the trial, owned by the harness.** A call costs the trial either the +wall time it took or nothing: the trial context's `inference_latency` flag asks a sim for the former, +and a real rig pays it regardless. Only the harness reads the flag. Paying nothing means holding the +world for the call, which holds a virtual clock still; paying wall time means letting the world run, +though no further ahead of the call's start than wall time has. The clock the harness hands the +policy stack (`now`) reads the instant the in-flight call's output takes effect, so a scheduling +wrapper stamps its chunk at `now()` and never learns the mode. + **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every model-facing view — action space, control frame, vendor format — is a codec's projection. diff --git a/CODE_RULES.md b/CODE_RULES.md index 88c5c0179..e7b2f7d11 100644 --- a/CODE_RULES.md +++ b/CODE_RULES.md @@ -190,8 +190,9 @@ inside the single entity that uses it. Distance costs the reader a search in bot use, to find out what it does; from the definition, to find out why it exists. A private name has every user in the file, so its place is determined. With a single user, first ask -whether the name is worth keeping: a body that says as much as its name is better inlined than moved. -Not touching `self` is no reason to stay at module level — that is what `@staticmethod` is for. +whether the name is worth keeping at all (`earn-its-place`); this rule only decides where a name that +is worth keeping goes. Not touching `self` is no reason to stay at module level — that is what +`@staticmethod` is for. A public name is looser: most of its users are elsewhere, and a module may order its surface deliberately. Group it with its in-file callers where that ordering does not say otherwise. Module @@ -285,14 +286,20 @@ except json.JSONDecodeError: ### earn-its-place -Don't add a class, file, field or parameter for a distinction the code already encodes. Check each of -the places one can already live — the dict an entry sits in, an enum member, the calling context, a -module that owns the subject — and use that instead. +Don't add a class, file, function, field or parameter for a distinction the code already encodes. Check +each of the places one can already live — the dict an entry sits in, an enum member, the calling +context, a module that owns the subject — and use that instead. The cost is not the lines. Every new type is another thing a reader holds, another place a value can live, and another edge to keep in sync; a field duplicating what its caller already knows goes stale the first time only one of the two is set. +With a single user, ask what the name adds over what already stands at that one use. A function whose +body is its name spelled out, a class whose fields the caller holds as locals anyway, a parameter only +ever passed one value — in each the use site already carries the distinction, and the name is a second +place to look. A name earns its place by holding what the use site cannot say: a diagnostic it raises, +an invariant it keeps, a domain predicate that takes more than one clause. + A new distinction lands beside its siblings, in the module that owns the subject. Put it there without asking. Raise it only when nothing suggests a home: that is a finding about the layout, and a human decides it. This rule does not apply to a repository still laying out its structure. diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 5954d942e..122cad7f1 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -101,6 +101,7 @@ The client sends the full raw robot state as a dict. Keys are flat strings (the | `robot_state.ee_pose` | float32 | (7,) | End-effector pose: `x, y, z, qw, qx, qy, qz` (quaternion is **wxyz**, scalar first) | | `robot_state.q` | float32 | (7,) | Joint positions (radians) | | `robot_state.dq` | float32 | (7,) | Joint velocities (radians/s) | +| `robot_state.fault` | bool | scalar | Whether the arm is faulted. A faulted arm sends no `robot_state.ee_pose`/`q`/`dq` at all — it has no sample to give — so a stack that reads them needs `StopOnFault` ahead of it | | `grip` | float32 | scalar | Gripper closure in `[0, 1]`: 0 = open, 1 = closed | | `image.<name>` | uint8 | (H, W, 3) | Camera RGB. Every eval target — PhAIL and each sim — sends `image.exterior` and `image.wrist`, whatever the underlying benchmark calls those cameras, so one codec reads them all; a target with more views adds its own names beside them (the MuJoCo sim adds `image.agent_view`) | | `obs_time_ns` | int | scalar | Harness-clock timestamp of this observation (ns) | @@ -108,7 +109,7 @@ The client sends the full raw robot state as a dict. Keys are flat strings (the | `task` | str | — | Language instruction for the episode | | `descriptor` | str | — | Embodiment the observation came from (e.g. `mujoco.franka`); empty string when unset. Lets a multi-embodiment policy adapt to the current robot | -Your server receives every key each step. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. +Your server receives every key each step, except that a faulted arm omits its `robot_state.*` measurements and sends only `robot_state.fault`. The standard stack puts `StopOnFault` ahead of the model, which answers such a step itself and never forwards it; a stack without it reaches the model with those keys missing. A mid-reset arm is not that case: it produces no step at all, so the client skips the round rather than sending a partial observation. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. ### Actions (server → client) @@ -169,7 +170,7 @@ from positronic.drivers.roboarm import command from positronic.offboard import PolicyServer from positronic.policy import Policy, Session from positronic.policy.spec import PolicySource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault class MySession(Session): @@ -200,20 +201,20 @@ class MyPolicy(Policy): return {'type': 'my_model'} -pipeline = ChunkedSchedule() | remote | PolicySource(MyPolicy(load_my_model())) +pipeline = StopOnFault() | ChunkedSchedule() | remote | PolicySource(MyPolicy(load_my_model())) PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() ``` -The pipeline reads left to right: everything left of the `remote` marker is the client-side stack the server declares in its handshake (here the standard `ChunkedSchedule`); everything right of it runs on the server. `PolicySource` is the pipeline's terminal — a model source that serves one already-built policy. +The pipeline reads left to right: everything left of the `remote` marker is the client-side stack the server declares in its handshake (here the standard `StopOnFault` and `ChunkedSchedule`); everything right of it runs on the server. `PolicySource` is the pipeline's terminal — a model source that serves one already-built policy. -The left side is not optional: a pipeline with nothing there is refused when the server starts, and a rig refuses a handshake that declares nothing. It needs a scheduler in particular. Actions come back timestamped relative to their chunk, and `ChunkedSchedule` is what turns those into times on the rig's clock; a stack that leaves them relative — or anchors them twice — makes the harness reject the chunk at the first inference, since it schedules nothing more than `MAX_ACTION_SKEW_SEC` from now. +The left side is not optional: a pipeline with nothing there is refused when the server starts, and a rig refuses a handshake that declares nothing. It needs a scheduler in particular, and `StopOnFault` outside that scheduler — a faulted arm is not tracking the plan it was given, so the wrapper answers the empty trajectory and the rig stops rather than resuming a chunk stamped before the fault. It is what makes the session above safe to write: a faulted observation arrives without `robot_state.ee_pose`/`q`/`dq`, so a session reading them by name needs the wrapper ahead of it. Actions come back timestamped relative to their chunk, and `ChunkedSchedule` is what turns those into times on the rig's clock; a stack that leaves them relative — or anchors them twice — makes the harness reject the chunk at the first inference, since it schedules nothing more than `MAX_ACTION_SKEW_SEC` from now. `new_session`'s `now` argument is the runtime clock that wrappers scheduling against live time read; a policy that does no scheduling of its own just accepts and ignores it (server-side it is `None`). If you put a `Codec` right of the marker (`ChunkedSchedule() | remote | codec | PolicySource(...)`), your session works entirely in *model space* — it receives encoded observations and returns model-native actions, and the codec handles the wire format. A codec that encodes images should also bound them on the rig, so full-resolution frames never cross the wire — that is what the built-in vendor pipelines do: ```python -ChunkedSchedule() | RestrictImageSize() | remote | codec | source +StopOnFault() | ChunkedSchedule() | RestrictImageSize() | remote | codec | source ``` Give it the geometry your codec encodes to — `RestrictImageSize(224, 224)` for a 224x224 model — so a frame is shrunk once, on the rig. The default is a loose 640x640, for a codec that resizes to nothing in particular. Leaving it out costs bandwidth, not correctness. diff --git a/pimm/core.py b/pimm/core.py index 34c3df7ea..6e4e3d6ab 100644 --- a/pimm/core.py +++ b/pimm/core.py @@ -11,7 +11,11 @@ class NoValueException(Exception): pass -NODEFAULT = object() +class _NoDefault: + """No default at all, which `T | None` cannot express: `None` is itself a value a receiver may default to.""" + + +NODEFAULT = _NoDefault() @dataclass @@ -164,7 +168,7 @@ def emit(self, data: T, ts: int = -1): class ControlSystemReceiver(SignalReceiver[T]): """Receiver adaptor bound to a single upstream signal on behalf of a system.""" - def __init__(self, owner: ControlSystem, default: T | None = NODEFAULT, maxsize: int | None = None): + def __init__(self, owner: ControlSystem, default: T | None | _NoDefault = NODEFAULT, maxsize: int | None = None): self._owner = owner self._default = default self._internal: SignalReceiver[T] | None = None @@ -187,8 +191,8 @@ def read(self) -> Message[T] | None: value = self._internal.read() if value is not None: return value - if self._default is not NODEFAULT: - # Always not-updated; the check above excludes the sentinel, which `T | None` cannot express. + if not isinstance(self._default, _NoDefault): + # Always not-updated; `None` reaches here as the value it was given, not as an absent default. return Message(cast(T, self._default), -1, False) return None @@ -227,7 +231,9 @@ class ReceiverDict(dict[str, ControlSystemReceiver[U]]): Pass fake=True for all fake receivers, or fake={'key1', 'key2'} for specific keys. """ - def __init__(self, owner: ControlSystem, *, default: U | None = NODEFAULT, fake: bool | Iterable[str] = False): + def __init__( + self, owner: ControlSystem, *, default: U | None | _NoDefault = NODEFAULT, fake: bool | Iterable[str] = False + ): super().__init__() self._owner = owner self._default = default diff --git a/positronic/cfg/embodiment.py b/positronic/cfg/embodiment.py index 29e5116af..42876b3e7 100644 --- a/positronic/cfg/embodiment.py +++ b/positronic/cfg/embodiment.py @@ -25,7 +25,7 @@ def droid(robot_arm, gripper, cameras): """Real single-arm Franka (DROID) + Robotiq gripper + ZED cameras.""" observations = { - 'robot_state': Observation(robot_arm.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot_arm.state, Serializers.robot_state), keys.GRIP: Observation(gripper.grip, None), **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } @@ -48,7 +48,7 @@ def droid(robot_arm, gripper, cameras): def yam(robot_arm, cameras): """Real single-arm i2rt YAM: the arm driver carries the gripper (they share one CAN chain).""" observations = { - 'robot_state': Observation(robot_arm.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot_arm.state, Serializers.robot_state), keys.GRIP: Observation(robot_arm.grip, None), **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } @@ -96,21 +96,22 @@ def yam_bimanual(left_channel: str, right_channel: str, mounts: dict[str, list[f for side, channel in (('left', left_channel), ('right', right_channel)) } observations = { - **{f'robot_state.{s}': Observation(arm.state, Serializers.robot_state) for s, arm in arms.items()}, - **{f'grip.{s}': Observation(arm.grip, None) for s, arm in arms.items()}, + **{f'{keys.ROBOT_STATE}.{s}': Observation(arm.state, Serializers.robot_state) for s, arm in arms.items()}, + **{f'{keys.GRIP}.{s}': Observation(arm.grip, None) for s, arm in arms.items()}, **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { **{ - f'robot_command.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) + f'{keys.ROBOT_COMMAND}.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) for s, arm in arms.items() }, - **{f'target_grip.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, + **{f'{keys.TARGET_GRIP}.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, } - joint_signals = {side: f'robot_state.{side}.q' for side in arms} + joint_signals = {side: f'{keys.ROBOT_STATE}.{side}.q' for side in arms} static_meta = { keys.JOINT_SIGNALS: list(joint_signals.values()), - keys.POSE_SIGNALS: [f'robot_state.{s}.ee_pose' for s in arms] + [f'robot_command.{s}.pose' for s in arms], + keys.POSE_SIGNALS: [f'{keys.ROBOT_STATE}.{s}.ee_pose' for s in arms] + + [f'{keys.ROBOT_COMMAND}.{s}.pose' for s in arms], keys.MOUNTS: {sig: mounts[side] for side, sig in joint_signals.items()}, } return Embodiment( @@ -133,7 +134,7 @@ def mujoco_franka(sim, camera_dict): Mujoco does not render the second image when using only 2 cameras. """ observations = { - 'robot_state': Observation(sim.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(sim.state, Serializers.robot_state), keys.GRIP: Observation(sim.grip, None), **{name: Observation(sim.cameras[orig], Serializers.camera_images) for name, orig in camera_dict.items()}, } diff --git a/positronic/cfg/wrappers.py b/positronic/cfg/wrappers.py index d8c2300e4..6f2c0f6aa 100644 --- a/positronic/cfg/wrappers.py +++ b/positronic/cfg/wrappers.py @@ -1,7 +1,7 @@ import configuronic as cfn from positronic import keys as obs_keys -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack chunked_schedule = cfn.Config(ChunkedSchedule) temporal_stack = cfn.Config(TemporalStack) @@ -47,4 +47,4 @@ def video_context_wrappers(history_frames: int, stride: int, keys: tuple[str, .. stack = TemporalStack( keys=tuple(keys), offsets_sec=_frame_offsets_sec(history_frames, stride, fps), pad_start=pad_start ) - return stack | ChunkedSchedule() + return StopOnFault() | stack | ChunkedSchedule() diff --git a/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index fb81e2c56..3d6e22dca 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -12,7 +12,7 @@ import pimm import positronic.cfg.policy as policy_cfg -from positronic import telemetry, telemetry_keys, utils, wire +from positronic import keys, telemetry, telemetry_keys, utils, wire from positronic.cfg.eval import placeholder from positronic.cli.eval.submit import submit from positronic.dataset.ds_writer_agent import TimeMode @@ -194,7 +194,7 @@ def run( eval: Eval | str, policy, output_dir=None, - inference_latency=False, + inference_latency: bool = False, timing=False, policy_image: str | None = None, alias: str | None = None, @@ -222,8 +222,9 @@ def run( if not isinstance(eval, Eval): raise SystemExit(f'--eval={eval!r} is a name, not a config: pass --policy-image to run it on the platform') # The eval config owns the trial sweep (seed, task range); ``inference_latency`` is the CLI's per-run knob - # (sim inference-cost simulation). Overlay it onto every trial context, then self-drive the eval. - eval = replace(eval, trials=[{**trial, 'inference_latency': inference_latency} for trial in eval.trials]) + # (whether a sim charges each call its wall time). Overlay it onto every trial context, then self-drive + # the eval. + eval = replace(eval, trials=[{**trial, keys.INFERENCE_LATENCY: inference_latency} for trial in eval.trials]) main(policy=policy, evals=[eval], output_dir=output_dir, timing=timing) return None diff --git a/positronic/data_collection.py b/positronic/data_collection.py index fb18915d9..c4cef98cf 100644 --- a/positronic/data_collection.py +++ b/positronic/data_collection.py @@ -160,9 +160,9 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self.sound.emit(abort_wav_path) tracker.turn_off() recording = False - self.robot_commands.emit([(clock.now_ns(), roboarm.command.Reset())]) + self.robot_commands.emit(roboarm.command.Reset()) - self.target_grip.emit([(clock.now_ns(), button_handler.get_value('right_trigger'))]) + self.target_grip.emit(button_handler.get_value('right_trigger')) cp_msg = self.controller_positions.read() if cp_msg.updated: target_robot_pos = tracker.update(cp_msg.data['right']) @@ -175,7 +175,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self.sound.emit(error_wav_path) if not in_error and cp_msg.updated: cmd = roboarm.command.CartesianPosition(target_robot_pos) - self.robot_commands.emit([(clock.now_ns(), cmd)]) + self.robot_commands.emit(cmd) yield pimm.Sleep(0.001) diff --git a/positronic/dataset/README.md b/positronic/dataset/README.md index e31623873..376670a69 100644 --- a/positronic/dataset/README.md +++ b/positronic/dataset/README.md @@ -354,6 +354,7 @@ Each line of `edits.jsonl` is one JSON record carrying its op. `{"op": "set_stat Key ideas - Inputs are registered explicitly through `DsWriterAgent.add_signal(name, serializer=None)`. - The agent polls inputs at a configurable rate and appends only on updates. +- Recording is best effort, and this is a deliberate trade rather than an oversight. Each input arrives over a one-slot `pimm` signal where a new value overwrites one still unread, so a recorder that stalls for longer than the gap between two samples loses the older one — commands exactly as much as camera frames or arm state. An episode is what the recorder managed to observe, not a guaranteed-complete log of what happened; treat a missing sample as possible in any analysis that counts them. - A separate `command` channel controls episode lifecycle. - `time_mode` selects how timestamps are recorded: `CLOCK` (default) stamps samples when the agent ingests them (useful during live data collection so every signal reflects when the recorder could act on it), while `MESSAGE` preserves the timestamp attached by the emitting control system (ideal for analysing inference latency by keeping original emission times). diff --git a/positronic/dataset/ds_writer_agent.py b/positronic/dataset/ds_writer_agent.py index f6ec10428..f5d9b286f 100644 --- a/positronic/dataset/ds_writer_agent.py +++ b/positronic/dataset/ds_writer_agent.py @@ -10,7 +10,7 @@ from .dataset import DatasetWriter from .episode import EpisodeWriter -from .serializers import Serializer, StatefulSerializer, Timestamped, _PureSerializer, expand_suffixed +from .serializers import Serializer, StatefulSerializer, _PureSerializer, expand_suffixed logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) @@ -59,76 +59,6 @@ def ABORT(): return DsWriterCommand(DsWriterCommandType.ABORT_EPISODE) -class TrajectoryOverrideSerializer(StatefulSerializer): - """Flatten policy trajectories into a single monotonic per-point stream. - - A policy emits whole trajectories ``[(abs_ts_ns, value), ...]``. A newer - trajectory overrides the overlapping tail of the previous one - (last-writer-wins on the timeline): given ``1:[1..10]`` then ``2:[5..15]`` - the recorded stream is ``1@1..4`` then ``2@5..15``. A point is committed - only once a newer trajectory starting after it proves it final; the - remainder is drained by :meth:`flush` at episode end. - - HACK: lossy. Drops the notion of a *predicted* trajectory and cannot - represent overlapping schedulers (RTC/temporal ensembling) that replan into - the already-committed past — such points are dropped to keep timestamps - strictly increasing. Faithful full-command recording needs an - object-valued Signal (``Kind.OBJECT``); tracked in TODO(positronic#NNN). - """ - - def __init__(self, inner: Serializer | None): - self._inner = inner - self._buffer: list[tuple[int, Any]] = [] # latest trajectory, (abs_ts_ns, value), ts-sorted - self._last_ts: int | None = None - - def reset(self) -> None: - self._buffer = [] - self._last_ts = None - - def _encode(self, value: Any) -> Any: - return self._inner(value) if self._inner is not None else value - - def _committable(self, points: list[tuple[int, Any]]) -> list[Timestamped]: - # Guard only bites in the overlap-degrade case (RTC replanning into the - # past); under ChunkedSchedule the prefix is always already ahead. - if self._last_ts is not None: - points = [(ts, v) for ts, v in points if ts > self._last_ts] - if points: - self._last_ts = points[-1][0] - return [Timestamped(ts, self._encode(v)) for ts, v in points] - - def __call__(self, message: list[tuple[int, Any]]) -> list[Timestamped]: - if not message: - # Empty trajectory is the cancel signal (the Harness emits it at episode end): - # drop the buffered tail so flush() does not commit canceled waypoints. - self._buffer = [] - return [] - - start = message[0][0] - # Buffer is ts-sorted: everything before the new trajectory's start is - # final; the rest is overridden and dropped by the reassignment below. - cut = next((i for i, (ts, _) in enumerate(self._buffer) if ts >= start), len(self._buffer)) - committed = self._committable(self._buffer[:cut]) - self._buffer = list(message) - return committed - - def flush(self, now_ns: int | None = None) -> list[Timestamped]: - # At episode end, commit only points already due (ts <= now_ns); the - # remaining future-scheduled tail never executed, so drop it. ``now_ns`` - # is None only for callers wanting the legacy "commit everything". - points = self._buffer if now_ns is None else [(ts, v) for ts, v in self._buffer if ts <= now_ns] - out = self._committable(points) - self._buffer = [] - return out - - -def _append(ep_writer: EpisodeWriter, name: str, value: Any, ts_ns: int, extra_ts: dict[str, int] | None = None): - for full_name, v in expand_suffixed(name, value): - if v is None: - continue - ep_writer.append(full_name, v, ts_ns, extra_ts) - - class TimeMode(IntEnum): """Mode of timestamping for the dataset writer.""" @@ -194,7 +124,26 @@ def add_signal(self, name: str, serializer: Serializer | StatefulSerializer | No def inputs(self) -> dict[str, pimm.ControlSystemReceiver[Any]]: return frozen_keys_dict(self._inputs) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C901 + def _record(self, ep_writer: EpisodeWriter, name: str, msg: pimm.Message, clock: pimm.Clock) -> None: + """Append one input's sample, stamped as ``time_mode`` selects and carrying every clock beside it.""" + world_time_ns, message_time_ns = clock.now_ns(), msg.ts + primary_ts = world_time_ns if self._time_mode == TimeMode.CLOCK else message_time_ns + + extra_ts = {'message': message_time_ns, 'system': pimm.world.SystemClock().now_ns()} + # Only add 'world' if clock is not system clock + if not isinstance(clock, pimm.world.SystemClock): + extra_ts['world'] = world_time_ns + + with self._telemetry_span(): + serializer = self._serializers.get(name) + value = msg.data + if serializer is not None: + value = serializer(value) + for full_name, v in expand_suffixed(name, value): + if v is not None: + ep_writer.append(full_name, v, primary_ts, extra_ts) + + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): """Main loop: process commands and append updated inputs to the episode.""" limiter = pimm.utils.RateLimiter(clock, hz=self._poll_hz) pace = (lambda: pimm.Yield()) if self._virtual_time else limiter.wait @@ -212,12 +161,14 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 opened = False if start: was_open = ep_writer is not None - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) opened = ep_writer is not None and not was_open if ep_writer is not None: for name, reader in self._inputs.items(): msg = reader.read() + if msg is None: + continue # Scope a command turn's drain to the episode window: the open turn keeps only # samples after START (dropping the inter-episode home command and any pre-reset # frame), the closing turn keeps only samples at or before STOP (dropping post-STOP @@ -226,37 +177,17 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 # normally. after_start = not opened or msg.ts > cmd_msg.ts before_stop = not closing or msg.ts <= cmd_msg.ts - if msg is not None and msg.updated and after_start and before_stop: - world_time_ns, message_time_ns = clock.now_ns(), msg.ts - primary_ts = world_time_ns if self._time_mode == TimeMode.CLOCK else message_time_ns - - extra_ts = {'message': message_time_ns, 'system': pimm.world.SystemClock().now_ns()} - # Only add 'world' if clock is not system clock - if not isinstance(clock, pimm.world.SystemClock): - extra_ts['world'] = world_time_ns - - with self._telemetry_span(): - serializer = self._serializers.get(name) - value = msg.data - if serializer is not None: - value = serializer(value) - # Gate on `Timestamped` so plain list-valued samples - # (e.g. list-state vectors) still go through `_append`. - # Empty list matches too — used as the cancel signal. - if isinstance(value, list) and (not value or isinstance(value[0], Timestamped)): - for sample in value: - _append(ep_writer, name, sample.value, sample.ts, None) - else: - _append(ep_writer, name, value, primary_ts, extra_ts) + if msg.updated and after_start and before_stop: + self._record(ep_writer, name, msg, clock) if closing: - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) yield pace() finally: cmd_msg = self.command.read() if cmd_msg.updated: - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) if ep_writer is not None: try: @@ -265,9 +196,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 ep_writer.__exit__(None, None, None) logger.info(f'DsWriterAgent: [ABORT] Episode {ep_counter}') - def _handle_command( # noqa: C901 - self, cmd: DsWriterCommand, ep_writer: EpisodeWriter | None, ep_counter: int, now_ns: int | None = None - ): + def _handle_command(self, cmd: DsWriterCommand, ep_writer: EpisodeWriter | None, ep_counter: int): match cmd.type: case DsWriterCommandType.START_EPISODE: if ep_writer is None: @@ -282,10 +211,6 @@ def _handle_command( # noqa: C901 logger.warning('Episode already started, ignoring start command') case DsWriterCommandType.STOP_EPISODE: if ep_writer is not None: - with self._telemetry_span(): - for name, ser in self._serializers.items(): - for sample in ser.flush(now_ns): - _append(ep_writer, name, sample.value, sample.ts, None) for k, v in cmd.static_data.items(): ep_writer.set_static(k, v) ep_writer.__exit__(None, None, None) diff --git a/positronic/dataset/serializers.py b/positronic/dataset/serializers.py index 90d92fafd..d04578446 100644 --- a/positronic/dataset/serializers.py +++ b/positronic/dataset/serializers.py @@ -1,5 +1,4 @@ from collections.abc import Callable, Iterator -from dataclasses import dataclass from typing import Any import numpy as np @@ -26,20 +25,9 @@ # - use "" (empty string) to keep the base name as-is # - any dict entry with value None is skipped # * None -> the sample is dropped -# * a list[Timestamped] -> a self-timestamped stream (recording only): each item is -# recorded at its own ``ts_ns``. An empty list defers; a StatefulSerializer may emit -# the remainder later from ``flush()``. The per-item ``value`` follows the rules above. Serializer = Callable[[Any], Any | dict[str, Any]] -@dataclass -class Timestamped: - """A sample paired with its own absolute timestamp (ns).""" - - ts: int - value: Any - - class StatefulSerializer: """Base for serializers registered with ``DsWriterAgent``. @@ -51,19 +39,9 @@ class StatefulSerializer: def reset(self) -> None: pass - def __call__(self, value: Any) -> Any | dict[str, Any] | list['Timestamped']: + def __call__(self, value: Any) -> Any | dict[str, Any]: raise NotImplementedError - def flush(self, now_ns: int | None = None) -> list['Timestamped']: - """Drain any buffered samples at episode end (mirror of ``reset``). - - Called once on ``STOP_EPISODE`` before the episode is finalized. ``now_ns`` - is the episode-end time; serializers that buffer future-scheduled samples - use it to drop the un-executed tail. The default keeps stateless - serializers a no-op. - """ - return [] - class _PureSerializer(StatefulSerializer): """Wraps a plain callable so every serializer has a uniform interface.""" diff --git a/positronic/dataset/tests/test_ds_writer_agent.py b/positronic/dataset/tests/test_ds_writer_agent.py index 95a97ba5c..8f3aaf997 100644 --- a/positronic/dataset/tests/test_ds_writer_agent.py +++ b/positronic/dataset/tests/test_ds_writer_agent.py @@ -9,17 +9,12 @@ import pimm from positronic import geom, keys, telemetry, telemetry_keys from positronic.dataset import DatasetWriter, EpisodeWriter -from positronic.dataset.ds_writer_agent import ( - DsWriterAgent, - DsWriterCommand, - DsWriterCommandType, - TimeMode, - TrajectoryOverrideSerializer, -) +from positronic.dataset.ds_writer_agent import DsWriterAgent, DsWriterCommand, DsWriterCommandType, TimeMode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers -from positronic.drivers.roboarm import RobotStatus, State +from positronic.drivers.roboarm import RobotStatus from positronic.drivers.roboarm import command as rcmd +from positronic.drivers.roboarm.tests.fakes import FakeRobotState from positronic.tests.testing_coutils import run_scripted_agent @@ -356,33 +351,9 @@ def test_transform_3d_serializer(world): np.testing.assert_allclose(names_vals[0][1][3:], q.as_quat) -class _FakeState(State): - def __init__(self, q, dq, ee_pose, status): - self._q = q - self._dq = dq - self._ee = ee_pose - self._status = status - - @property - def q(self): - return self._q - - @property - def dq(self): - return self._dq - - @property - def ee_pose(self): - return self._ee - - @property - def status(self): - return self._status - - def test_robot_state_serializer_drops_reset_and_emits_components(world): ds = FakeDatasetWriter() - agent, cmd_em, emitters = build_agent_with_pipes({'robot_state': Serializers.robot_state}, ds, world) + agent, cmd_em, emitters = build_agent_with_pipes({keys.ROBOT_STATE: Serializers.robot_state}, ds, world) q = np.arange(7, dtype=np.float32) dq = np.arange(7, dtype=np.float32) + 10 @@ -391,8 +362,8 @@ def test_robot_state_serializer_drops_reset_and_emits_components(world): script = [ (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.START_EPISODE)), 0.001), - (lambda: emitters['robot_state'].emit(_FakeState(q, dq, pose, RobotStatus.RESETTING)), 0.001), - (lambda: emitters['robot_state'].emit(_FakeState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(FakeRobotState(q, dq, pose, RobotStatus.RESETTING)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(FakeRobotState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.STOP_EPISODE)), 0.001), ] @@ -482,36 +453,18 @@ def test_pickles_with_every_constructor_argument_filled(): virtual_time=True, telemetry_span=partial(telemetry.span, telemetry_keys.SPAN_RECORD_IO), ) - agent.add_signal('robot_command', TrajectoryOverrideSerializer(Serializers.robot_command)) - agent.add_signal('robot_state', Serializers.robot_state) + agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) + agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) loaded = pickle.loads(pickle.dumps(agent)) - assert set(loaded.inputs) == {'robot_command', 'robot_state'} + assert set(loaded.inputs) == {keys.ROBOT_COMMAND, keys.ROBOT_STATE} with loaded._telemetry_span(): pass -def test_trajectory_override_serializer(): - s = TrajectoryOverrideSerializer(None) - s.reset() - - # First trajectory: nothing is final yet (could be overridden). - assert s([(1, 'a'), (2, 'b'), (3, 'c')]) == [] - - # Next trajectory starts at ts=2 -> only ts<2 ('a') is final; 'b','c' overridden. - out = s([(2, 'B'), (3, 'C'), (4, 'D')]) - assert [(t.ts, t.value) for t in out] == [(1, 'a')] - - # Episode end drains the still-live buffer. - assert [(t.ts, t.value) for t in s.flush()] == [(2, 'B'), (3, 'C'), (4, 'D')] - - def test_serializer_plain_list_value(world): - """A serializer returning a plain list (non-`Timestamped`) is appended as one sample. - - The trajectory-stream dispatch must not hijack legitimate list-valued samples. - """ + """A serializer returning a plain list is appended as one sample, values and all.""" ds = FakeDatasetWriter() def to_list(_): @@ -527,52 +480,3 @@ def to_list(_): w = ds.created[-1] assert [(s, v) for (s, v, _, _) in w.appends] == [('v', [1, 2, 3])] - - -def test_trajectory_override_serializer_empty_cancels_buffer(): - """Empty trajectory is the Harness STOP cancel signal: drop the buffered tail.""" - s = TrajectoryOverrideSerializer(None) - s.reset() - - # Buffer a trajectory (nothing committed yet). - assert s([(1, 'a'), (2, 'b'), (3, 'c')]) == [] - # Empty trajectory = cancel: nothing committed AND buffer cleared. - assert s([]) == [] - # Subsequent flush must not emit the canceled waypoints. - assert s.flush() == [] - - -def test_trajectory_override_serializer_flush_cutoff(): - """flush(now_ns) commits only points already due; the future tail is dropped.""" - s = TrajectoryOverrideSerializer(None) - s.reset() - - # Buffer a chunk scheduled at ts 1..4 (nothing committed yet). - assert s([(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')]) == [] - - # Episode ends at ts=2: only the due points (ts <= 2) are committed; 'c','d' dropped. - assert [(t.ts, t.value) for t in s.flush(now_ns=2)] == [(1, 'a'), (2, 'b')] - - # No cutoff keeps the legacy "commit everything" behavior. - s.reset() - assert s([(1, 'a'), (2, 'b')]) == [] - assert [(t.ts, t.value) for t in s.flush()] == [(1, 'a'), (2, 'b')] - - -def test_stop_commits_due_drops_future_trajectory(world): - """A mid-trajectory STOP commits already-due samples and drops the un-executed tail.""" - ds = FakeDatasetWriter() - agent, cmd_em, emitters = build_agent_with_pipes({'traj': TrajectoryOverrideSerializer(None)}, ds, world) - - future = 10**18 # far beyond the test clock, so it stays an un-executed tail - script = [ - (partial(cmd_em.emit, DsWriterCommand(DsWriterCommandType.START_EPISODE)), 0.001), - (partial(emitters['traj'].emit, [(0, 'due'), (future, 'tail')]), 0.001), - (partial(cmd_em.emit, DsWriterCommand(DsWriterCommandType.STOP_EPISODE)), 0.001), - ] - run_scripted_agent(agent, script, world=world) - - w = ds.created[-1] - # 'due' (ts <= stop time) is committed; the future 'tail' is dropped. - assert [(s, v) for (s, v, _, _) in w.appends] == [('traj', 'due')] - assert w.exited is True diff --git a/positronic/drivers/gripper/dh.py b/positronic/drivers/gripper/dh.py index df62fa463..9b3a6b524 100644 --- a/positronic/drivers/gripper/dh.py +++ b/positronic/drivers/gripper/dh.py @@ -3,7 +3,6 @@ import pimm from positronic.drivers import vendor_import -from positronic.drivers.roboarm.command import Trajectory, TrajectoryPlayer with vendor_import('pymodbus', 'Gripper support'): import pymodbus.client as ModbusClient @@ -13,7 +12,7 @@ class DHGripper(pimm.ControlSystem): def __init__(self, port: str): self.port = port self.grip: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.SignalReceiver[Trajectory[float]] = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.force: pimm.SignalReceiver = pimm.ControlSystemReceiver(self, default=100) self.speed: pimm.SignalReceiver = pimm.ControlSystemReceiver(self, default=100) @@ -33,18 +32,14 @@ def _state_r(): while _state_g() != 1 and _state_r() != 1: yield pimm.Sleep(0.1) - player = TrajectoryPlayer() last_grip = 0.0 - # TODO: We must translate these to physical units (N and m/s) + # TODO: Should we translate these to physical units (N and m/s)? while not should_stop.value: try: grip_msg = self.target_grip.read() - if grip_msg.updated: - player.set(grip_msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - last_grip = grip + if grip_msg is not None and grip_msg.updated: + last_grip = grip_msg.data width = round((1 - max(0, min(last_grip, 1))) * 1000) client.write_register(0x103, c_uint16(width).value, slave=1) client.write_register(0x101, c_uint16(self.force.value).value, slave=1) @@ -78,7 +73,7 @@ def _state_r(): force.emit(100) for width in np.sin(np.linspace(0, 10 * np.pi, 60)) + 1: - target_grip.emit([(world.clock.now_ns(), width)]) + target_grip.emit(width) time.sleep(0.5) try: print(f'Real grip position: {grip.value}') diff --git a/positronic/drivers/gripper/robotiq.py b/positronic/drivers/gripper/robotiq.py index f7bdfb59b..893f3def5 100644 --- a/positronic/drivers/gripper/robotiq.py +++ b/positronic/drivers/gripper/robotiq.py @@ -4,7 +4,6 @@ import pimm from positronic.drivers import vendor_import -from positronic.drivers.roboarm.command import Trajectory, TrajectoryPlayer with vendor_import('pymodbus', 'Gripper support'): import pymodbus.client as ModbusClient @@ -22,7 +21,7 @@ class Robotiq2F(pimm.ControlSystem): def __init__(self, port: str): self._port = port self.grip = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.ControlSystemReceiver[Trajectory[float]] = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.force = pimm.ControlSystemReceiver(self, default=255) # device scale 0..255 self.speed = pimm.ControlSystemReceiver(self, default=255) # device scale 0..255 @@ -37,15 +36,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p client.write_registers(_REG_CMD, [0x0000, 0x0000, 0x0000], device_id=_SLAVE) client.write_registers(_REG_CMD, [0x0100, 0x0000, 0x0000], device_id=_SLAVE) - player = TrajectoryPlayer() - while not should_stop.value: pos_msg = self.target_grip.read() - if pos_msg.updated: - player.set(pos_msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - pos = int(max(0, min(1, grip)) * 255) + if pos_msg is not None and pos_msg.updated: + pos = int(max(0, min(1, pos_msg.data)) * 255) spd = int(max(0, min(255, self.speed.value))) frc = int(max(0, min(255, self.force.value))) @@ -81,7 +75,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p while True: if time.time() - start > i * 1.0: - tgt.emit([(world.clock.now_ns(), waypoints[i])]) + tgt.emit(waypoints[i]) i += 1 if i >= len(waypoints): break diff --git a/positronic/drivers/roboarm/command.py b/positronic/drivers/roboarm/command.py index d231efce9..f7c62e016 100644 --- a/positronic/drivers/roboarm/command.py +++ b/positronic/drivers/roboarm/command.py @@ -1,7 +1,7 @@ """Collection of commands that can be sent to the robot.""" from dataclasses import dataclass, field -from typing import Any, TypeAlias, TypeVar +from typing import Any import numpy as np @@ -80,45 +80,6 @@ def apply(self, current: geom.Transform3D) -> geom.Transform3D: CommandType = Reset | CartesianPosition | JointPosition | JointDelta | CartesianDelta -_T = TypeVar('_T') - -# The wire shape of every command channel: waypoints stamped with absolute clock ns. A single immediate -# command is the one-waypoint trajectory ``[(clock.now_ns(), value)]``; ``[]`` cancels whatever is in flight. -Trajectory: TypeAlias = list[tuple[int, _T]] - - -def _combine(acc: CommandType, cmd: CommandType) -> CommandType: - match (acc, cmd): - case (CartesianDelta(a, frame_a), CartesianDelta(b, frame_b)): - if not np.allclose(frame_a.as_matrix, frame_b.as_matrix): - raise ValueError('Cannot accumulate cartesian deltas expressed in different frames') - return CartesianDelta(_compose_delta(a, b), frame_a) - case (JointDelta(a), JointDelta(b)): - return JointDelta(a + b) - case (CartesianDelta() | JointDelta(), _) | (_, CartesianDelta() | JointDelta()): - raise ValueError(f'Cannot reduce {type(acc).__name__} then {type(cmd).__name__} in one tick') - case _: - return cmd - - -def reduce(due: Trajectory[CommandType]) -> CommandType: - """Collapse the commands due in one control tick into the single command to execute. - - Folds the batch in timestamp order. A run of same-space deltas accumulates (their motion is summed, so a - missed tick is caught up rather than dropped); a run of absolute commands keeps the last. Mixing an absolute - with a delta, or two delta spaces, has no faithful single-command form -- a delta binds to the pose measured - when it is consumed, which an absolute target or a foreign space cannot supply -- and raises. - """ - result = due[0][1] - for _, cmd in due[1:]: - result = _combine(result, cmd) - return result - - -def _reduce_last(due: Trajectory[Any]) -> Any: - """The trailing value wins -- the right collapse for absolute setpoints and gripper targets.""" - return due[-1][1] - def to_wire(command: CommandType) -> dict[str, Any]: match command: @@ -138,35 +99,6 @@ def to_wire(command: CommandType) -> dict[str, Any]: } -class TrajectoryPlayer: - """Plays back a timestamped trajectory at the driver's control rate. - - Call ``set()`` when a new trajectory arrives, then ``advance(now)`` each tick to get the single command to - apply: every waypoint whose timestamp has been reached is collapsed by ``reduce`` into one value (the last - one by default; the arm channels pass ``command.reduce`` to accumulate due deltas instead of dropping them). - """ - - def __init__(self, reduce=_reduce_last): - self._trajectory: Trajectory[Any] = [] - self._index: int = 0 - self._reduce = reduce - - def set(self, trajectory: Trajectory[Any]): - self._trajectory = trajectory - self._index = 0 - - def advance(self, current_time: int): - """Collapse every waypoint whose timestamp <= current_time into the single value to apply, or None.""" - due = [] - while self._index < len(self._trajectory): - ts, value = self._trajectory[self._index] - if ts > current_time: - break - self._index += 1 - due.append((ts, value)) - return self._reduce(due) if due else None - - def from_wire(wire: dict[str, Any]) -> CommandType: match wire['type']: case 'reset': diff --git a/positronic/drivers/roboarm/franka.py b/positronic/drivers/roboarm/franka.py index 4614cbbc6..e8c31effb 100644 --- a/positronic/drivers/roboarm/franka.py +++ b/positronic/drivers/roboarm/franka.py @@ -141,9 +141,7 @@ def __init__( self._home_joints_variation = ( home_joints_variation if home_joints_variation is not None else [0.03, 0.05, 0.08, 0.08, 0.10, 0.10, 0.10] ) - self.commands: pimm.SignalReceiver[command.Trajectory[command.CommandType]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) self.state: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) self._load = load @@ -318,7 +316,7 @@ def _park(self, robot) -> Iterator[pimm.Sleep]: except Exception: logging.exception('Parking failed, the arm stays where it stands') - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: with self._desk_session(): robot = self._ensure_robot() try: @@ -332,7 +330,6 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p yield from self._reset(robot, robot_state, rate_limiter, should_stop) in_error = False - player = command.TrajectoryPlayer(reduce=command.reduce) while not should_stop.value: st = robot.state() @@ -344,23 +341,17 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p logging.warning(f'Robot error: {st.error_message}') cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) if in_error: # The driver always clears a recoverable error itself; making it optional (hold in # ERROR for out-of-band recovery instead) is a config knob to add when an embodiment # needs it. robot.recover_from_errors() - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale - # waypoint once the error clears. - player.set([]) yield rate_limiter.wait() continue - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): yield from self._reset(robot, robot_state, rate_limiter, should_stop) case command.CartesianPosition(pose): @@ -376,8 +367,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p robot.set_target_joints(positions) case command.JointDelta(velocities=joint_delta): robot.set_target_joints(st.q + joint_delta) - case _: - raise NotImplementedError(f'Unsupported command {cmd}') + case other: + raise NotImplementedError(f'Unsupported command {other}') yield rate_limiter.wait() @@ -419,7 +410,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p if time.monotonic() > start + duration: print(f'Moving to {pos + origin.translation}') target = command.CartesianPosition(geom.Transform3D(pos + origin.translation, origin.rotation)) - commands.emit([(world.clock.now_ns(), target)]) + commands.emit(target) i += 1 else: time.sleep(0.01) diff --git a/positronic/drivers/roboarm/kinova/driver.py b/positronic/drivers/roboarm/kinova/driver.py index baf573687..3f3a11040 100644 --- a/positronic/drivers/roboarm/kinova/driver.py +++ b/positronic/drivers/roboarm/kinova/driver.py @@ -68,9 +68,7 @@ def __init__(self, ip: str, relative_dynamics_factor=0.2, home_joints: list[floa self.relative_dynamics_factor = relative_dynamics_factor self.solver = KinematicsSolver() self.home_joints = home_joints if home_joints is not None else [0.0, -0, 0.5, -1.5, 0.0, -0.5, 1.57079633] - self.commands: pimm.SignalReceiver[command.Trajectory[command.CommandType]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) self.state: pimm.SignalEmitter[KinovaState] = pimm.ControlSystemEmitter(self) def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: @@ -89,15 +87,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p joint_controller.compute_torque(q, dq, tau) current_command = np.zeros(api.actuator_count, dtype=np.float32) - player = command.TrajectoryPlayer(reduce=command.reduce) - while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): joint_controller.set_target_qpos(self.home_joints) case command.CartesianPosition(pose): @@ -110,8 +103,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p case command.JointPosition(positions): qpos = np.array(positions, dtype=np.float32) joint_controller.set_target_qpos(qpos) - case _: - print(f'Unsuported command: {cmd}') + case other: + print(f'Unsuported command: {other}') torque_command = joint_controller.compute_torque(q, dq, tau) np.divide(torque_command, torque_constant, out=current_command) diff --git a/positronic/drivers/roboarm/so101/driver.py b/positronic/drivers/roboarm/so101/driver.py index f9578630f..82a558670 100644 --- a/positronic/drivers/roboarm/so101/driver.py +++ b/positronic/drivers/roboarm/so101/driver.py @@ -59,13 +59,12 @@ def __init__(self, motor_bus: MotorBus, home_joints: list[float] | None = None): self.kinematic = Kinematics(_SO101_URDF_PATH, _SO101_EE_JOINT) self.joint_limits = self.kinematic.joint_limits self.home_joints = home_joints if home_joints is not None else [0.0, 0.0, 0.0, 0.0, 0.0] - self.commands: pimm.SignalReceiver[roboarm_command.Trajectory[roboarm_command.CommandType]] = ( - pimm.ControlSystemReceiver(self, default=[]) - ) - self.target_grip: pimm.SignalReceiver[roboarm_command.Trajectory[float]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[roboarm_command.CommandType](self) + self.target_grip = pimm.ControlSystemReceiver[float](self) self._last_grip: float = 0.0 + # The arm half of the motor setpoint, in normalized units. ``None`` until the first arm command: + # a gripper target arriving before one has no arm position to pair with. + self._last_qpos: np.ndarray | None = None self.grip: pimm.SignalEmitter[float] = pimm.ControlSystemEmitter(self) self.state: pimm.SignalEmitter[SO101State] = pimm.ControlSystemEmitter(self) @@ -88,40 +87,30 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p rate_limit = pimm.RateLimiter(hz=1000, clock=clock) state = SO101State() - player = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - grip_player = roboarm_command.TrajectoryPlayer() - while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) grip_msg = self.target_grip.read() - if grip_msg.updated: - grip_player.set(grip_msg.data) - grip = grip_player.advance(clock.now_ns()) + grip = grip_msg.data if grip_msg is not None and grip_msg.updated else None + command = cmd_msg.data if cmd_msg is not None and cmd_msg.updated else None if grip is not None: self._last_grip = grip - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if command is not None: + match command: case roboarm_command.Reset(): raise NotImplementedError('Reset not implemented') case roboarm_command.CartesianPosition(pose): - qpos = self._solve_ik(state, pose) - q_with_gripper = np.concatenate([qpos, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) + self._last_qpos = self._solve_ik(state, pose) case roboarm_command.CartesianDelta() as delta_cmd: ee_pose, _ = self._forward_kinematics(self.motor_bus.position) - target = delta_cmd.apply(ee_pose) - qpos = self._solve_ik(state, target) - q_with_gripper = np.concatenate([qpos, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) + self._last_qpos = self._solve_ik(state, delta_cmd.apply(ee_pose)) case roboarm_command.JointPosition(qpos): - q_norm = self.rad_to_norm(qpos) - q_with_gripper = np.concatenate([q_norm, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) - case _: - raise ValueError(f'Unknown command: {cmd}') + self._last_qpos = self.rad_to_norm(qpos) + case other: + raise ValueError(f'Unknown command: {other}') + # The arm and the gripper are one setpoint on a shared bus, but they arrive as two channels that + # need not carry a value in the same round, so either one changing rewrites the whole vector. + if (command is not None or grip is not None) and self._last_qpos is not None: + self.motor_bus.set_target_position(np.concatenate([self._last_qpos, [self._last_grip]])) q = self.motor_bus.position dq = self.motor_bus.velocity[:-1] diff --git a/positronic/drivers/roboarm/tests/fakes.py b/positronic/drivers/roboarm/tests/fakes.py new file mode 100644 index 000000000..9440bd4a4 --- /dev/null +++ b/positronic/drivers/roboarm/tests/fakes.py @@ -0,0 +1,41 @@ +"""Fakes for the ``roboarm`` interfaces, for tests that drive an arm without a driver.""" + +import numpy as np + +from positronic import geom +from positronic.drivers.roboarm import RobotStatus, State + + +class FakeRobotState(State): + """A ``State`` over the four values it is given, copied on read so a caller cannot reach the arrays + a producer keeps emitting from.""" + + def __init__(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, status: RobotStatus) -> None: + self._q = q + self._dq = dq + self._ee_pose = ee_pose + self._status = status + + @property + def q(self) -> np.ndarray: + return self._q.copy() + + @property + def dq(self) -> np.ndarray: + return self._dq.copy() + + @property + def ee_pose(self) -> geom.Transform3D: + return geom.Transform3D(translation=self._ee_pose.translation.copy(), rotation=self._ee_pose.rotation) + + @property + def status(self) -> RobotStatus: + return self._status + + +def make_robot_state(translation, joints, status: RobotStatus = RobotStatus.AVAILABLE) -> FakeRobotState: + """A stationary arm at ``translation`` with identity rotation: zero joint velocity throughout.""" + translation = np.asarray(translation, dtype=np.float32) + joints = np.asarray(joints, dtype=np.float32) + pose = geom.Transform3D(translation=translation, rotation=geom.Rotation.identity) + return FakeRobotState(joints, np.zeros_like(joints), pose, status) diff --git a/positronic/drivers/roboarm/yam.py b/positronic/drivers/roboarm/yam.py index 829514fc9..f4abd423e 100644 --- a/positronic/drivers/roboarm/yam.py +++ b/positronic/drivers/roboarm/yam.py @@ -184,13 +184,13 @@ def __init__( self._sim = sim self._connect = connect - self.commands = pimm.ControlSystemReceiver[command.Trajectory[command.CommandType]](self, default=[]) - self.target_grip = pimm.ControlSystemReceiver[command.Trajectory[float]](self, default=[]) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.state = pimm.ControlSystemEmitter[YamState](self) self.grip = pimm.ControlSystemEmitter[float](self) self.robot_meta = pimm.ControlSystemEmitter[dict[str, Any]](self) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: arm = self._connect(self._channel, self._sim) try: kin = _Kinematics() @@ -199,33 +199,21 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p robot_state = YamState() limiter = pimm.RateLimiter(clock, hz=100) - player = command.TrajectoryPlayer(reduce=command.reduce) - grip_player = command.TrajectoryPlayer() q_target = self._reset(arm, kin, robot_state) grip_target = 0.0 while not should_stop.value: - if (cmd_msg := self.commands.read()) is not None and cmd_msg.updated: - player.set(cmd_msg.data) + cmd_msg = self.commands.read() if (grip_msg := self.target_grip.read()) is not None and grip_msg.updated: - grip_player.set(grip_msg.data) - - grip = grip_player.advance(clock.now_ns()) - if grip is not None: - grip_target = float(grip) + grip_target = float(grip_msg.data) obs = arm.get_observations() q = obs['joint_pos'] - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): - # Drop the in-flight trajectories so the homed arm holds position rather than - # resuming stale waypoints on the next tick. - player.set([]) - grip_player.set([]) q_target = self._reset(arm, kin, robot_state) grip_target = 0.0 obs = arm.get_observations() @@ -239,8 +227,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p case command.CartesianDelta() as delta_cmd: target = delta_cmd.apply(self._base_pose * kin.fk(q)) q_target = self._ik_or_hold(kin, target, q, q_target) - case _: - raise NotImplementedError(f'Unsupported command {cmd}') + case other: + raise NotImplementedError(f'Unsupported command {other}') arm.command_joint_pos(np.append(q_target, 1.0 - grip_target)) @@ -329,8 +317,8 @@ def close(self) -> None: with pimm.World() as world: # `World.pair` cannot express that it returns the counterpart of the port it is given, so the four # payload types are named here. - commands = cast(pimm.SignalEmitter[command.Trajectory[command.CommandType]], world.pair(robot.commands)) - target_grip = cast(pimm.SignalEmitter[command.Trajectory[float]], world.pair(robot.target_grip)) + commands = cast(pimm.SignalEmitter[command.CommandType], world.pair(robot.commands)) + target_grip = cast(pimm.SignalEmitter[float], world.pair(robot.target_grip)) state = cast(pimm.SignalReceiver[YamState], world.pair(robot.state)) grip = cast(pimm.SignalReceiver[float], world.pair(robot.grip)) @@ -356,12 +344,12 @@ def pump(seconds: float): assert home_err < 1e-4, home_err # Grip round-trip: polarity inverted on the way out (command) and on the way back (observation). - target_grip.emit([(world.clock.now_ns(), 0.8)]) + target_grip.emit(0.8) pump(0.5) assert fake.last_command is not None assert abs(fake.last_command[6] - 0.2) < 1e-6, fake.last_command # positronic 0.8 closed -> chain 0.2 assert abs(grip.value - 0.8) < 0.02, grip.value - target_grip.emit([(world.clock.now_ns(), 0.0)]) + target_grip.emit(0.0) pump(0.5) assert abs(fake.last_command[6] - 1.0) < 1e-6, fake.last_command assert abs(grip.value) < 0.02, grip.value @@ -369,7 +357,7 @@ def pump(seconds: float): # Unfold toward the workspace, then drive a Cartesian square through the driver's IK. The square # sits well inside the reach envelope, at the unfolded posture's wrist orientation. reach_q = np.array([0.0, 1.2, 1.2, 0.0, 0.6, 0.0]) - commands.emit([(world.clock.now_ns(), command.JointPosition(reach_q))]) + commands.emit(command.JointPosition(reach_q)) pump(0.5) if fake is not None: assert np.allclose(state.value.q, reach_q, atol=0.02), state.value.q @@ -383,7 +371,7 @@ def pump(seconds: float): assert solution is not None, f'IK failed for {target}' ik_err = np.linalg.norm(kin.fk(solution).translation - target.translation) assert ik_err < 5e-3, ik_err # FK↔IK consistency - commands.emit([(world.clock.now_ns(), command.CartesianPosition(target))]) + commands.emit(command.CartesianPosition(target)) pump(0.7) reached = np.linalg.norm(state.value.ee_pose.translation - target.translation) print(f'Moved to {target.translation}, error {reached * 1000:.2f} mm') diff --git a/positronic/keys.py b/positronic/keys.py index 9ee511da4..587a07dd6 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -26,11 +26,20 @@ def is_robot_command(name: str) -> bool: return name == ROBOT_COMMAND or name.startswith(f'{ROBOT_COMMAND}.') -JOINTS = 'robot_state.q' -JOINT_VEL = 'robot_state.dq' -EE_POSE = 'robot_state.ee_pose' +# The arm's state channel, and the signals a recorded state unfolds into. As on the command side, the +# suffixes are ``Serializers.robot_state``'s, so the names derive from the channel rather than restating it. +ROBOT_STATE = 'robot_state' +JOINTS = f'{ROBOT_STATE}.q' +JOINT_VEL = f'{ROBOT_STATE}.dq' +EE_POSE = f'{ROBOT_STATE}.ee_pose' +# Whether the arm is faulted, in every observation the harness builds. A faulted arm has no sample to give +# and is not tracking the plan it was handed, so the policy stack — not the harness — decides what happens +# next: the rest of ``ROBOT_STATE`` is absent from that observation. +ROBOT_FAULT = f'{ROBOT_STATE}.fault' GRIP = 'grip' TASK = 'task' +# The embodiment an observation came from, so a multi-embodiment policy can tell which robot it is driving. +DESCRIPTOR = 'descriptor' # The prefix that identifies a camera on the wire: an embodiment declares its cameras by naming # them this way, and every consumer picks them out of the observations by it. IMAGE_PREFIX = 'image.' @@ -88,3 +97,8 @@ def is_robot_command(name: str) -> bool: # leaves it absent on failure — a reader defaults it rather than assuming a False. EVAL_SUCCESS = 'eval.success' EVAL_TERMINATED = 'eval.terminated' + +# Whether a model call charges the world clock its own wall duration. A flag: any other type is rejected +# when the episode starts. A sim trial without it charges nothing (the world holds still per call); +# hardware always pays wall whatever the trial asks. +INFERENCE_LATENCY = 'inference_latency' diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 6a3b2c417..596847b91 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -57,7 +57,9 @@ class TestPrepareObs: def test_images_pass_through_untouched_by_default(self): session = RemoteSession(_mock_ws_session()) obs = {'cam': _make_image(480, 640), 'state': np.array([1.0])} - assert session._prepare_obs(obs) is obs + prepared = session._prepare_obs(obs) + assert prepared.keys() == obs.keys() + assert all(prepared[key] is value for key, value in obs.items()) def test_compression_reaches_nested_images(self): session = RemoteSession(_mock_ws_session(), compress_images=True) diff --git a/positronic/policy/action.py b/positronic/policy/action.py index cb9a801d6..7af571de9 100644 --- a/positronic/policy/action.py +++ b/positronic/policy/action.py @@ -26,6 +26,8 @@ def _relative_rot_vec(q_current: np.ndarray, q_target: np.ndarray, representatio class AbsolutePositionAction(Codec): + WIRE_NAME = 'absolute_position_action' + def __init__(self, tgt_ee_pose_key: str, tgt_grip_key: str, rotation_rep: RotRep | str = RotRep.QUAT): self.rot_rep = RotRep(rotation_rep) self.tgt_ee_pose_key = tgt_ee_pose_key @@ -54,7 +56,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'absolute_position_action', + 'name': self.WIRE_NAME, 'args': { 'tgt_ee_pose_key': self.tgt_ee_pose_key, 'tgt_grip_key': self.tgt_grip_key, @@ -64,6 +66,8 @@ def to_spec(self): class AbsoluteJointsAction(Codec): + WIRE_NAME = 'absolute_joints_action' + def __init__(self, tgt_joints_key: str, tgt_grip_key: str, num_joints: int = 7): self.tgt_joints_key = tgt_joints_key self.tgt_grip_key = tgt_grip_key @@ -92,7 +96,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'absolute_joints_action', + 'name': self.WIRE_NAME, 'args': { 'tgt_joints_key': self.tgt_joints_key, 'tgt_grip_key': self.tgt_grip_key, @@ -137,6 +141,8 @@ def training_encoder(self): class RelativePositionAction(Codec): + WIRE_NAME = 'relative_position_action' + def __init__( self, rotation_rep: RotRep | str = RotRep.QUAT, @@ -194,7 +200,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'relative_position_action', + 'name': self.WIRE_NAME, 'args': { 'rotation_rep': self.rot_rep.value, 'robot_pose_key': self.robot_pose_key, @@ -211,6 +217,8 @@ class JointDeltaAction(Codec): into a ``JointDelta`` command; the driver integrates each delta onto the live measured joints. """ + WIRE_NAME = 'joint_delta_action' + # General DROID form scales each normalized velocity by its own per-joint delta limit, then # renorms the velocity vector so no joint exceeds its limit: # RELATIVE_MAX_JOIN_DELTA = np.array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]) @@ -241,4 +249,4 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return {keys.ROBOT_COMMAND: command.JointDelta(velocities=velocities), keys.TARGET_GRIP: grip} def to_spec(self): - return {'name': 'joint_delta_action', 'args': {'num_joints': self.num_joints}} + return {'name': self.WIRE_NAME, 'args': {'num_joints': self.num_joints}} diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 3bb4d66b0..ed16a804e 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,11 +1,12 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Callable -from typing import Any +from collections.abc import Callable, Mapping +from typing import Any, ClassVar Now = Callable[[], float] + # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. SEQ = 'seq' PAR = 'par' @@ -32,7 +33,7 @@ class Session(ABC): """ @abstractmethod - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: Mapping[str, Any]) -> list[dict[str, Any]] | None: """Predict actions for the given observation.""" @property @@ -147,11 +148,15 @@ def wrap_session(self, inner: Session, context: dict[str, Any] | None, now: Now """Wrap a single session. Subclasses override this for per-session wrapping.""" raise NotImplementedError('Override wrap_session or wrap') + # The name this wrapper travels under, set by every deliverable subclass. ``WIRE_WRAPPERS`` is keyed by + # it, so the name is written once and both sides of the wire read the same attribute. + WIRE_NAME: ClassVar[str] + def to_spec(self) -> dict[str, Any]: """Plain-data wire spec of this wrapper, for a server's local-stack declaration. Only wrappers registered in ``positronic.policy.spec.WIRE_WRAPPERS`` are deliverable to a rig. - The spec is ``{'name': <wire name>}`` plus ``{'args': {...}}`` when the wrapper takes any; + The spec is ``{'name': WIRE_NAME}`` plus ``{'args': {...}}`` when the wrapper takes any; ``args`` are constructor keywords, since the rig rebuilds by calling the constructor with them. """ raise NotImplementedError(f'{type(self).__name__} is not deliverable to a rig (no wire spec)') diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 16ae3a000..bba31fd04 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -240,6 +240,8 @@ class ActionTimestamp(Codec): At training time, surfaces ``action_fps`` as transform metadata. """ + WIRE_NAME = 'action_timestamp' + def __init__(self, *, fps: float): self._fps = fps self._dt = 1.0 / fps @@ -266,7 +268,7 @@ def meta(self): return {'action_fps': self._fps} def to_spec(self): - return {'name': 'action_timestamp', 'args': {'fps': self._fps}} + return {'name': self.WIRE_NAME, 'args': {'fps': self._fps}} class ActionHorizon(Codec): @@ -285,6 +287,8 @@ class ActionHorizon(Codec): At training time, surfaces ``action_horizon_sec`` as transform metadata. """ + WIRE_NAME = 'action_horizon' + def __init__(self, horizon_sec: float): self._horizon_sec = horizon_sec @@ -310,7 +314,7 @@ def meta(self): return {'action_horizon_sec': self._horizon_sec} def to_spec(self): - return {'name': 'action_horizon', 'args': {'horizon_sec': self._horizon_sec}} + return {'name': self.WIRE_NAME, 'args': {'horizon_sec': self._horizon_sec}} def ActionTiming(*, fps: float, horizon_sec: float | None = None) -> Codec: @@ -335,6 +339,8 @@ class BinarizeGripTraining(Codec): timing | BinarizeGripTraining(('grip', 'target_grip')) | BinarizeGripInference() | obs & action """ + WIRE_NAME = 'binarize_grip_training' + def __init__(self, keys: tuple[str, ...], threshold: float = 0.5): self._keys = keys self._threshold = threshold @@ -361,7 +367,7 @@ def _derive(episode): return Group(Derive(**transforms), Identity()) def to_spec(self): - return {'name': 'binarize_grip_training', 'args': {'keys': list(self._keys), 'threshold': self._threshold}} + return {'name': self.WIRE_NAME, 'args': {'keys': list(self._keys), 'threshold': self._threshold}} class BinarizeGripInference(Codec): @@ -372,6 +378,8 @@ class BinarizeGripInference(Codec): timing | BinarizeGripInference() | obs & action """ + WIRE_NAME = 'binarize_grip_inference' + def __init__(self, threshold: float = 0.5, key: str = obs_keys.TARGET_GRIP): self._threshold = threshold self._key = key @@ -389,7 +397,7 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return data def to_spec(self): - return {'name': 'binarize_grip_inference', 'args': {'threshold': self._threshold, 'key': self._key}} + return {'name': self.WIRE_NAME, 'args': {'threshold': self._threshold, 'key': self._key}} class FlipGrip(Codec): @@ -404,6 +412,8 @@ class FlipGrip(Codec): timing | FlipGrip() | obs & action """ + WIRE_NAME = 'flip_grip' + def encode(self, data): # Copy: the original dict is also the decode ``context`` and the raw recording tap's input. if obs_keys.GRIP in data: @@ -420,7 +430,7 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return data def to_spec(self): - return {'name': 'flip_grip'} + return {'name': self.WIRE_NAME} def _scaled(image: np.ndarray, width: int, height: int) -> np.ndarray: @@ -444,6 +454,8 @@ class RestrictImageSize(Codec): ChunkedSchedule() | RestrictImageSize() | remote | codec | source """ + WIRE_NAME = 'restrict_image_size' + def __init__(self, width: int = 640, height: int = 640): self._width = width self._height = height @@ -475,7 +487,7 @@ def training_encoder(self) -> EpisodeTransform: ) def to_spec(self): - return {'name': 'restrict_image_size', 'args': {'width': self._width, 'height': self._height}} + return {'name': self.WIRE_NAME, 'args': {'width': self._width, 'height': self._height}} class ChangeEEFrame(Codec): @@ -491,6 +503,8 @@ class ChangeEEFrame(Codec): of the observation/action codecs. """ + WIRE_NAME = 'change_ee_frame' + @staticmethod def _move(value: Any, transform: geom.Transform3D) -> Any: """A pose vector or an arm command, re-expressed through ``transform``.""" @@ -565,6 +579,6 @@ def meta(self): def to_spec(self): # Lists, not tuples, so the spec is identical before and after a wire round-trip. return { - 'name': 'change_ee_frame', + 'name': self.WIRE_NAME, 'args': {'transform': self._transform.as_vector(_QUAT).tolist(), 'keys': list(self._keys)}, } diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index cee52a980..fd50adf49 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,35 +1,37 @@ +import concurrent.futures +import logging import time +from collections import deque from collections.abc import Generator, Iterable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum from typing import Any +import numpy as np from opentelemetry.trace import Span import pimm from positronic import keys, telemetry, telemetry_keys from positronic.dataset.ds_writer_agent import DsWriterCommand from positronic.dataset.serializers import expand_suffixed +from positronic.drivers import roboarm from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Task -from positronic.policy.base import Policy, Session -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.base import Policy from positronic.utils import flatten_dict, frozen_view -# How far from now an action may be scheduled. A chunk spans seconds, so this is loose enough that no real -# trajectory approaches it, and tight enough to catch a rig-side stack that left timestamps relative to the -# chunk (decades behind) or anchored them twice (decades ahead). +# How far from now an action may be scheduled: past any real chunk, short of the decades a rig-side stack is +# off by when it leaves timestamps chunk-relative or anchors them twice. MAX_ACTION_SKEW_SEC = 60.0 +# How long a real-time round may last when no waypoint is due sooner. It bounds how late a directive is +# noticed, and with it the granularity every command timestamp is quantized to. +POLL_PERIOD_SEC = 0.01 -def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: - """Reject a chunk whose timestamps are not times on the harness clock.""" - skew = max((abs(action[keys.ACTION_TIMESTAMP] - now) for action in actions), default=0.0) - if skew > MAX_ACTION_SKEW_SEC: - raise ValueError( - f'Action scheduled {skew:.0f}s from now, over the {MAX_ACTION_SKEW_SEC:.0f}s bound: the rig-side ' - f'stack is not anchoring chunks to the harness clock' - ) +# How long a submitted call may take and still resolve within its round: a wrapper that skips inference +# answers in microseconds, a real model call runs far past this and is throttled across rounds. +SKIP_REPLY_SEC = 0.001 class DirectiveType(Enum): @@ -61,14 +63,106 @@ def ABORT(cls) -> 'Directive': return cls(DirectiveType.ABORT) -class _EpisodeTelemetry: - """The live rollout's wall-clock telemetry: the episode span, its index, its step count and the virtual - instant it began. Inert while telemetry is unbound, so the harness calls it unconditionally. +class _InferenceWorker: + """One episode's policy session, called one at a time on a thread of its own so the harness keeps + playing while the model runs. - The span stays anchored while open, so the rollout's phase spans (reset, env.step, policy.infer, - record.io) parent to it rather than to the pass. + ``charge_wall`` is what a call costs the trial: the wall time it took, or nothing — the loop is held + for the call, which holds a virtual clock still. """ + def __init__(self, policy: Policy, context: dict[str, Any], charge_wall: bool, clock: pimm.Clock) -> None: + self._charge_wall = charge_wall + self._clock = clock + # World clock and ``time.monotonic()`` at the in-flight call's submit, anchored at the episode's + # start so ``effect_time`` reads a trial instant from the moment the session exists. + self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() + self._session = policy.new_session(context, self.effect_time) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') + self._call: Future[list[dict[str, Any]] | None] | None = None + + @property + def meta(self) -> dict[str, Any]: + return self._session.meta + + @property + def idle(self) -> bool: + return self._call is None + + @property + def done(self) -> bool: + """Whether the call in flight has returned.""" + return self._call is not None and self._call.done() + + def effect_time(self) -> float: + """The trial instant the in-flight call's output takes effect: its submit, plus its wall duration so + far when the trial pays wall time.""" + wall = time.monotonic() - self._wall_t0 if self._charge_wall else 0.0 + return self._t0_ns / 1e9 + wall + + @staticmethod + def _owned(obs: dict[str, Any]) -> dict[str, Any]: + """The observation with its arrays copied, so nothing rewrites what the worker is still reading. + + A producer may reuse one buffer for every sample it emits — a camera renders into the array behind + the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. + Copying at dispatch pays once per call rather than per round. + """ + return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} + + def submit(self, obs: dict[str, Any]) -> None: + """Start a call on ``obs``. The moment's wait lets a wrapper that skips inference resolve in the round + it was asked.""" + self._t0_ns, self._wall_t0 = self._clock.now_ns(), time.monotonic() + self._call = self._executor.submit(self._session, frozen_view(self._owned(obs))) + concurrent.futures.wait([self._call], timeout=SKIP_REPLY_SEC) + + def throttle(self) -> None: + """Slow the loop for the call in flight as the trial's mode requires: until the call returns when the + world is held for it, else only while the world is ahead of the call's own wall clock.""" + assert self._call is not None + # Wall time cannot be held still, so the world runs no further ahead of the call's start than it has. + timeout = max(self._clock.now() - self.effect_time(), 0.0) if self._charge_wall else None + concurrent.futures.wait([self._call], timeout=timeout) + + def result(self) -> list[dict[str, Any]] | None: + """The returned call's trajectory — ``None`` when it had nothing to place — leaving the worker idle.""" + assert self._call is not None + # Read on the loop thread, so a failing call still seals the episode. + actions, self._call = self._call.result(), None + return actions + + @staticmethod + def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: + """Report the failure of a call nobody is waiting for any more.""" + # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no + # caller left to raise to; the log is the only place its failure can go. + if not future.cancelled() and (exc := future.exception()) is not None: + logging.error(f'Inference failed after the episode that asked for it ended: {exc}') + + def abandon(self) -> None: + """Let go of the call in flight: its answer lands nowhere and its failure only reaches the log.""" + if self._call is not None: + self._call.add_done_callback(self._report_abandoned) + self._call = None + self._executor.shutdown(wait=False, cancel_futures=True) + + def join(self) -> None: + """Wait out an abandoned call, then close the session it was inside. + + ``shutdown(cancel_futures=True)`` cancels only what is still queued, so until this returns the call + still holds the session's resources: a ``RemoteSession``'s websocket, or the in-process model every + session shares. + """ + self._executor.shutdown(wait=True) + self._session.close() + + +class _EpisodeTelemetry: + """The live rollout's episode span, with the index, step count and virtual start it is stamped with. + Inert while telemetry is unbound, so the harness calls it unconditionally. The span stays anchored while + open, so the rollout's phase spans parent to it rather than to the pass.""" + def __init__(self) -> None: self._span: Span | None = None self._index = -1 @@ -95,8 +189,8 @@ def step(self) -> None: self._steps += 1 def end(self, virtual_now: float) -> None: - """Close a finished rollout, stamped with its step count and its virtual duration up to - ``virtual_now`` — captured when the rollout ended, before the flush round advances the sim clock.""" + """Close the rollout, stamped with its step count and its virtual duration up to ``virtual_now`` — + captured when the rollout ended, before the flush round advances the sim clock.""" if self._span is None: return self._close(virtual_now) @@ -110,15 +204,13 @@ def abort(self) -> None: self._end_span() def seal(self, virtual_now: float) -> None: - """Close a rollout abandoned mid-flight by a raising ``reset`` / ``new_session`` / session call, stamped - like a clean end and marked ``episode.partial``. Ending it is what exports it — the batch processor drops - an unended span, orphaning the finished children and losing their phases. Partial rather than aborted so - the reduce keeps it. Inert when no span is open.""" + """Close a rollout abandoned mid-flight by a raising ``reset`` / ``new_session`` / session call, marked + ``episode.partial`` so the reduce keeps it. Ending it is what exports it: the batch processor drops an + unended span, orphaning the finished children and losing their phases.""" if self._span is None: return telemetry.set_attrs(self._span, **{telemetry_keys.ATTR_EPISODE_PARTIAL: True}) - self._close(virtual_now) - telemetry.force_flush() + self.end(virtual_now) def _close(self, virtual_now: float) -> None: # A rollout whose first observation never landed — a reset that raised, or a task already done before @@ -137,23 +229,16 @@ def _end_span(self) -> None: class Harness(pimm.ControlSystem): - """Control system that runs the episode lifecycle and forwards trajectories to drivers. - - Handles directives (RUN/FINISH/ABORT) and dataset recording. Inference intelligence — scheduling, - error recovery, blending, absolute time stamping — lives in the policy/session layer; the harness - calls the session, demuxes the action dicts into per-channel trajectories, and emits. The RUN context - is handed whole to the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). - - A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial - whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's - ``timeout`` bounds every trial, self-driven or operator-driven, so an attended episode still terminates - at the deadline if the operator never sends FINISH. A bounded trial also ends early on a truthy - privileged ``done``, recording ``eval.terminated`` True and the delivered payload in its static data; a - timeout records False. A task-less session has neither deadline nor budget and ends only on directives. - - The ``Embodiment`` supplies the observation serializers (which own the canonical key names), the command - channels and the home action. The policy owns its wrapper stack; the harness runs what it is given, - passing ``new_session`` the clock the scheduling wrapper anchors chunks to. + """Control system that runs the episode lifecycle and plays the policy's trajectory to the drivers. + + The wrapper owns the plan, the harness plays it, one command per channel per round. The session call + runs on a worker so playing continues while the model does. A call costs the trial either the wall time + it took or nothing — the world held still for it — and the ``now`` handed to ``new_session`` reads the + instant the call's output takes effect, so wrappers stamp for it without knowing the mode. + + A ``trials`` plan makes the harness self-driving: it starts the next trial whenever idle and returns + once the plan is exhausted. A task's ``timeout`` bounds every trial either way, and a truthy privileged + ``done`` ends one early — ``eval.terminated`` records which. A task-less session ends only on directives. """ def __init__( @@ -173,31 +258,32 @@ def __init__( self.policy: Policy = policy self.context: dict[str, Any] = {} self._static_meta = static_meta or {} - self._policy_session: Session | None = None - # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. - self._running = False - # Sim-only, delivered on the RUN context: ``True`` advances the sim clock by the measured wall cost - # of the inference call, a float by a fixed deterministic amount (the reproducible golden). The chunk - # is anchored before the sleep, so ``_step`` post-shifts it. - self._inference_latency: bool | float = False + # This episode's session and the thread it runs on. ``None`` while no episode is live: between RUN + # and FINISH/ABORT, stepping and recording happen together. + self._worker: _InferenceWorker | None = None + # A worker let go of mid-call, kept until the join that makes closing its session safe. + self._retiring: _InferenceWorker | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None - # Whether this episode's first observation has landed. Until it does the deadline stands where the - # reset put it, which bounds an episode whose first observation never arrives. + # False until this episode's first observation lands; until then the deadline stands where the reset + # put it, which bounds an episode that never gets one. self._rollout_started = False # Wall-clock telemetry for the live rollout, opened under ``--timing`` and inert otherwise. self._telemetry = _EpisodeTelemetry() - # Channels that have not delivered since this episode's reset. A receiver latches its last value, so - # emptying this set is what keeps the first inference off the previous episode's final frame. + # Channels that have not delivered since this episode's RUN; the first inference waits until every one + # has. A receiver latches its last value, so a producer silent between episodes would otherwise feed + # the previous episode's final frame. Delivery is judged by ``updated``, not ``ts``: some producers + # stamp ``ts`` on their own clock. self._awaiting_obs: set[str] = set() - self._descriptor = embodiment.descriptor self.observations = pimm.ReceiverDict(self) self.commands = pimm.EmitterDict(self) for name in embodiment.observations: self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] + # Each channel's waypoints not yet played, stamped with absolute clock ns and ascending. + self._schedules: dict[str, deque[tuple[int, Any]]] = {name: deque() for name in embodiment.commands} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) @@ -220,79 +306,87 @@ def _build_episode_meta(self, context: dict[str, Any]) -> dict[str, Any]: meta['eval.timeout'] = self._task.timeout # ``policy.meta`` is the static baseline; the session overlays per-episode specifics (e.g. the # sampled sub-policy) and wins on conflict. - session_meta = self.policy.meta | (self._policy_session.meta if self._policy_session else {}) + session_meta = self.policy.meta | (self._worker.meta if self._worker else {}) for k, v in flatten_dict(session_meta).items(): meta[f'{keys.POLICY_META}.{k}'] = v meta.update(context) return meta - def _emit_now(self, action: dict[str, Any], clock: pimm.Clock) -> None: - now = clock.now_ns() + def _emit(self, action: dict[str, Any]) -> None: for name, value in action.items(): - self.commands[name].emit([(now, value)]) + self.commands[name].emit(value) - def _home(self, clock: pimm.Clock) -> None: - self._emit_now(self._embodiment.home, clock) + def _home(self) -> None: + self._emit(self._embodiment.home) - def _pace(self) -> pimm.Command: + def _pace(self, clock: pimm.Clock) -> pimm.Command: """Sim: yield, so the simulator's control-period sleep is the sole time-master and the policy reads - each observation instantly. Real: sleep the poll period to hold wall-clock rate.""" - return pimm.Yield() if self._embodiment.simulated else pimm.Sleep(0.01) - - def _bump_schedule_end(self, delta_sec: float) -> None: - """Shift the active ``ChunkedSchedule._Session`` ``_trajectory_end`` by ``delta_sec``. - - ``_step`` post-shifts a chunk's timestamps to pay for ``inference_latency``; the wrapper's - end-of-chunk gate must follow, or it re-infers before the driver has played the shifted trajectory. - - TODO: reaching the wrapper by attribute-name guessing and writing its private state breaks on any rename. - """ - s = self._policy_session - while s is not None: - if isinstance(s, ChunkedSchedule._Session) and s._trajectory_end is not None: - s._trajectory_end += delta_sec - return - s = getattr(s, '_inner', None) - - def _cancel_trajectories(self) -> None: - """Drop any in-flight chunk from drivers and from the recording's tail. - - Emits ``[]`` on every command channel, so each driver's ``TrajectoryPlayer`` clears its buffer - (devices hold position) and ``TrajectoryOverrideSerializer`` drops its uncommitted tail. Must precede - ``STOP_EPISODE``, which flushes the recording's serializers and would otherwise commit canceled - waypoints. Also cancels the session's scheduling state so the next inference is not held back. + each observation instantly. Real: sleep to the next waypoint, capped at the poll period, so a + waypoint is emitted at its own time and a round rarely finds more than one due.""" + if self._embodiment.simulated: + return pimm.Yield() + due = min((sched[0][0] for sched in self._schedules.values() if sched), default=None) + if due is None: + return pimm.Sleep(POLL_PERIOD_SEC) + return pimm.Sleep(min(POLL_PERIOD_SEC, max(due - clock.now_ns(), 1) / 1e9)) + + def _cancel_session(self) -> None: + """Drop everything the episode has going: the schedule being played, and the call on the worker. + + The call is let go of rather than waited for, so a model that hangs cannot hold up the recording's + stop or the home. Devices hold their last commanded position; nothing downstream is buffered. """ - self._emit_commands([]) - if self._policy_session is not None: - self._policy_session.cancel() + for schedule in self._schedules.values(): + schedule.clear() + self._retire_worker() + + def _retire_worker(self) -> None: + """Let go of this episode's worker and the call it is running, keeping it for ``_reap_worker``: + ending an episode must not wait for a model that hangs.""" + if self._worker is not None: + self._worker.abandon() + self._retiring, self._worker = self._worker, None + + def _reap_worker(self) -> None: + """Join the retired worker and close the session its abandoned call was inside.""" + if self._retiring is not None: + self._retiring.join() + self._retiring = None def _finalize_recording( self, clock: pimm.Clock, payload: dict[str, Any] | None = None ) -> Generator[pimm.Command, None, None]: """Commit the live episode: cancel the in-flight chunk, stop the recorder — stamping the episode's full static meta (plus any terminal payload) — then close its span.""" - self._cancel_trajectories() - self.ds_command.emit(DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})})) + # Stamped before the worker is retired: the meta overlays what its session reports. + stop = DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})}) + self._cancel_session() + self.ds_command.emit(stop) virtual_now = clock.now() # before the round below, whose sim-clock advance belongs to no rollout # Give the recorder a round to commit the STOP before the next START (they share ``ds_command``, where # last-value-wins would drop one) and before the home command, so homing stays out of the recording. - yield self._pace() - # After that round, so the recorder's STOP-time record.io span is still in flight and parents to the - # episode. Accepted skew: a producer stepping in that shared round charges one span (≤ one control - # period) to the closing episode — the cooperative scheduler cannot give the recorder a turn alone. + yield self._pace(clock) + # After that round, so the recorder's STOP-time record.io span still parents to the episode. Skew: a + # producer stepping in that shared round charges ≤ one control period to the closing episode. self._telemetry.end(virtual_now) def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. - A resettable task's ``reset`` only arms the producer, which publishes the first observation on a later - round. The recorder drains its channels the turn it opens, so the pre-reset frame and the - inter-episode home command drop out and its first sample is the post-reset scene. The deadline is - armed here and moved to that first observation once it lands, so an episode that never gets one is - still bounded. + A resettable task's ``reset`` only arms the producer; the first observation lands a later round. The + recorder drains its channels the turn it opens, so the pre-reset frame and the inter-episode home + command drop out. The deadline is armed here and moved to that first observation once it lands. """ + # Before the span opens, so the wait for a call the last episode abandoned is inter-episode wall + # rather than overhead the timing reducer attributes to this one. + self._reap_worker() self.context = context - self._inference_latency = self.context.get('inference_latency', False) + latency = self.context.get(keys.INFERENCE_LATENCY, False) + if not isinstance(latency, bool): + raise ValueError(f'{keys.INFERENCE_LATENCY} is a flag, got {latency!r}') + # A real rig pays wall time whatever the trial asks: the knob is sim-only, and the eval CLI writes it + # into every trial. + charge_wall = latency or not self._embodiment.simulated self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -304,58 +398,62 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._policy_session = self.policy.new_session(self.context, clock.now) - self._running = True + self._worker = _InferenceWorker(self.policy, self.context, charge_wall, clock) self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) def _end_episode( self, clock: pimm.Clock, payload: dict[str, Any] | None = None, *, abort: bool = False ) -> Generator[pimm.Command, None, None]: - """Close the live episode: finalize (or abort) the recording, release the session, home devices. + """Close the live episode: finalize (or abort) the recording, retire the session, home devices. - Releasing the session here, not only at shutdown, closes a ``RemoteSession``'s websocket promptly, so - the offboard server's per-session cleanup (active-session decrement, idle watchdog) runs now. + The worker is retired rather than joined here, so a ``RemoteSession``'s websocket outlives the call + still using it. """ - if self._running: + if self._worker is not None: if abort: - self._cancel_trajectories() # abort has no finalize to do it — stop drivers before the home + self._cancel_session() # abort has no finalize to do it — stop the episode before the home self.ds_command.emit(DsWriterCommand.ABORT()) - yield self._pace() # the settling round a finalize also takes, before the home command + yield self._pace(clock) # the settling round a finalize also takes, before the home command self._telemetry.abort() else: yield from self._finalize_recording(clock, payload) - if self._policy_session: - self._policy_session.close() - self._policy_session = None - self._home(clock) - self._running = False + self._home() def _handle_directive(self, directive: Directive, clock: pimm.Clock) -> Generator[pimm.Command, None, None]: - """Dispatch a directive to the episode lifecycle; updates ``_running``.""" + """Dispatch a directive to the episode lifecycle.""" match directive.type: case DirectiveType.RUN: - if not self._running: # a RUN mid-trial is ignored — the operator finishes before starting anew + if self._worker is None: # a RUN mid-trial is ignored — the operator finishes before starting anew self._begin_episode(directive.payload or {}, clock) case DirectiveType.FINISH: - if self._running: # a FINISH while idle is ignored — nothing to finalize + if self._worker is not None: # a FINISH while idle is ignored — nothing to finalize yield from self._end_episode(clock, directive.payload) case DirectiveType.ABORT: yield from self._end_episode(clock, abort=True) case _: raise ValueError(f'Unknown directive type: {directive.type}') + @staticmethod + def _is_faulted(value: Any) -> bool: + """Whether a raw observation is an arm reporting a fault. Every other not-ready sample is simply absent.""" + return isinstance(value, roboarm.State) and value.status is roboarm.RobotStatus.ERROR + def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: """Read every observation channel and assemble the policy input dict. - Raises ``NoValueException`` if any channel has no value yet. Returns ``None`` while a serializer - reports a sample is not ready (``robot_state`` during a ``RESETTING`` arm) or a channel still holds a - pre-reset value — either way the harness skips inference rather than feed a partial or stale obs. + Raises ``NoValueException`` if any channel has no value yet. Returns ``None`` while a sample is not + ready (``robot_state`` during a ``RESETTING`` arm) or a channel still holds a pre-reset value, rather + than feed a partial or stale obs. A faulted arm is the exception: the plan being played was made for + an arm that is now somewhere else, so its observation reaches the stack carrying ``keys.ROBOT_FAULT`` + and without the arm's own entries. """ # Against the live model, not the one known at episode start: a remote env publishes its ``robot_meta`` # a turn after the reset that produced it, so at episode start there is no model to check. assert_default_frame(self._statics()) inputs: dict[str, Any] = {} + faulted = False + not_ready = False for name, obs in self._embodiment.observations.items(): message = self.observations[name].read() if message is None: @@ -365,115 +463,134 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: value = message.data if obs.serializer is not None: value = obs.serializer(value) - if value is None: - return None - for full_name, v in expand_suffixed(name, value): - if v is not None: - inputs[full_name] = v - if self._awaiting_obs: + if value is None: # no sample to give: a resetting or faulted arm alike + # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so + # the fault is recovered from the raw sample and stapled on as `keys.ROBOT_FAULT` — a name + # already claiming to be part of `robot_state`. Emit it from the serializer and this + # branch, the raw-type check and the flag all go, and the fault reaches the recording too. + faulted = faulted or self._is_faulted(message.data) + not_ready = True + continue + inputs.update({full: v for full, v in expand_suffixed(name, value) if v is not None}) + # Every channel is read before this decision, so a bimanual rig cannot hide one arm's fault behind + # another arm's not-ready sample: whichever channel comes first, the fault still reaches the stack. + if (not_ready and not faulted) or self._awaiting_obs: return None + # The trial's context goes under what the harness read this round, never over it: a context carries + # whatever keys the RUN directive puts in it, and a ``robot_state.fault`` among them must not tell + # ``StopOnFault`` that a faulted arm is sound. + inputs = {**self.context, **inputs} + inputs[keys.ROBOT_FAULT] = faulted inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() - inputs.update(self.context) - inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it - return inputs - - def _emit_commands(self, actions: list[dict[str, Any]]) -> None: - """Republish-all demux: emit every command channel from this action chunk. - - Each channel gets the ``(ts_ns, value)`` waypoints the chunk carries for it; a channel the chunk - omits gets ``[]``, overwriting its last-value-wins signal so the driver holds. An empty chunk - therefore cancels every channel. - """ - for name, emitter in self.commands.items(): - # Wrappers do action-timing math in float seconds; every pimm channel client expects ns. This is - # the single explicit seconds->ns seam. - traj = [(int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in actions if name in a] - emitter.emit(traj) - - def _inference_delay(self, wall_start: float) -> float: - """The inference cost to simulate: measured wall time (``True``), a fixed float, or 0 (``False``).""" - if self._inference_latency is True: # bool is an int subclass — check identity first - return time.monotonic() - wall_start - return float(self._inference_latency) - - def _step(self, clock: pimm.Clock) -> Generator[pimm.Sleep, None, None]: - """Build obs, call the session, demux its chunk into per-channel emissions. - - The chunk already carries absolute timestamps, stamped by the outermost scheduling wrapper. - """ - obs = self._build_obs(clock) - if obs is None: - return + inputs[keys.DESCRIPTOR] = self._embodiment.descriptor if not self._rollout_started: - # The rollout begins at its first observation, not when the reset returned: a reset only asks the - # producer for a scene, and the turns spent delivering it are neither the trial's budget nor its - # duration. + # The rollout begins at its first observation, not when the reset returned: the turns spent + # delivering the scene are neither the trial's budget nor its duration. self._rollout_started = True self._telemetry.start_rollout(clock.now()) if self._task is not None: self._deadline = clock.now() + self._task.timeout + return inputs - # Advance the sim clock by the inference cost so rollouts feel the model's latency, but only on - # cycles where inference ran — sleeping on blocked cycles would slow directive handling. The chunk - # was anchored pre-sleep, so it is post-shifted and the wrapper's gate bumped to match. - wall_start = time.monotonic() - actions = self._policy_session(frozen_view(obs)) - if actions is None: - return - delay = self._inference_delay(wall_start) - if delay > 0.0: - yield pimm.Sleep(delay) - actions = [{**a, keys.ACTION_TIMESTAMP: a[keys.ACTION_TIMESTAMP] + delay} for a in actions] - self._bump_schedule_end(delay) - - # The latency sleep (or a slow inference call on a real clock) may have crossed the deadline. Drop - # the chunk rather than emit past the advertised self-termination point; the run loop fires FINISH - # next cycle. + def _step(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: + """One round of inference: throttle the loop for the call in flight as the trial's mode requires and + reschedule on what it returns; with no call in flight, submit one on a fresh observation and give it + the rest of the round to return.""" + if not worker.idle: + self._throttle_and_reschedule(worker, clock) + if worker.idle: + obs = self._build_obs(clock) + if obs is not None: + worker.submit(obs) + self._throttle_and_reschedule(worker, clock) + + def _throttle_and_reschedule(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: + worker.throttle() + if worker.done and (trajectory := worker.result()) is not None: + self._reschedule(trajectory, clock) + + @staticmethod + def _assert_anchored(trajectory: list[dict[str, Any]], now: float) -> None: + """Reject a chunk whose timestamps are not times on the harness clock.""" + skew = max((abs(action[keys.ACTION_TIMESTAMP] - now) for action in trajectory), default=0.0) + if skew > MAX_ACTION_SKEW_SEC: + raise ValueError( + f'Action scheduled {skew:.0f}s from now, over the {MAX_ACTION_SKEW_SEC:.0f}s bound: the ' + f'rig-side stack is not anchoring chunks to the harness clock' + ) + + def _reschedule(self, trajectory: list[dict[str, Any]], clock: pimm.Clock) -> None: + """Replace the schedule being played with the session's trajectory. Every channel it names gets that + channel's waypoints; one it omits is cleared and holds. The timestamps are already absolute, stamped + by the scheduling wrapper for the instant the call's output takes effect. + """ if self._deadline is not None and clock.now() >= self._deadline: + # The world reached the deadline while the call was in flight, so its chunk is dropped rather than + # placed past the point the trial advertises it stops at; ``_run`` finishes the trial next round. return - + self._assert_anchored(trajectory, clock.now()) self._telemetry.step() - _assert_anchored(actions, clock.now()) - self._emit_commands(actions) + # The single explicit seconds->ns seam: wrappers time actions in float seconds, the schedules and + # every pimm channel are in ns. + for name, schedule in self._schedules.items(): + schedule.clear() + schedule.extend((int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in trajectory if name in a) + + def _play(self, clock: pimm.Clock) -> None: + """Emit each channel's command due this round, and nothing on a channel with none. + + A channel with several waypoints due emits the last: exact for an absolute setpoint, lossy for a + relative one. Pacing keeps one due per round wherever a round is shorter than the waypoint spacing. + """ + now_ns = clock.now_ns() + for name, schedule in self._schedules.items(): + value = None + while schedule and schedule[0][0] <= now_ns: + value = schedule.popleft()[1] + if value is not None: + self.commands[name].emit(value) def _trial_terminal(self, clock: pimm.Clock) -> dict[str, Any] | None: """The terminal static payload if a self-driven trial has ended this round, else ``None``. - The deadline is hard: a truthy ``done`` delivered within budget records ``eval.terminated`` True plus - its payload, the budget passing records False, and a terminal past the deadline is a timeout rather - than a late success. Only a freshly delivered ``done`` counts — the receiver latches its last value, - so a prior trial's terminal would otherwise re-fire; gating on delivery clears it without asking the - producer to republish. Reached only for a task with a deadline. + The deadline is hard: a truthy ``done`` within budget records ``eval.terminated`` True plus its + payload, the budget passing records False, and a terminal past the deadline is a timeout rather than + a late success. Only a freshly delivered ``done`` counts, or the receiver's latched value would + re-fire a prior trial's terminal. """ + deadline = self._deadline + if deadline is None: # a task-less session has no budget and ends on directives alone + return None done_msg = self.done.read() - if done_msg.updated and done_msg.data and done_msg.ts <= self._deadline * 1e9: + assert done_msg is not None # the receiver carries a default, so ``read`` always yields a message + if done_msg.updated and done_msg.data and done_msg.ts <= deadline * 1e9: return {**done_msg.data, keys.EVAL_TERMINATED: True} - if clock.now() >= self._deadline: + if clock.now() >= deadline: return {keys.EVAL_TERMINATED: False} return None def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # Home the embodiment before the first episode; each ``_end_episode`` re-homes for the next one, so # every episode begins from the home pose (a real arm gets the inter-episode gap to reach it). - self._home(clock) + self._home() try: yield from self._run(should_stop, clock) except BaseException: - # A failure mid-rollout unwinds past the normal span close. Seal the open span before the - # exception reaches ``bind``'s exit flush, or it never exports and its finished children orphan, - # losing their phases and charging the episode's wall to between_episodes. + # Seal the open span before the exception reaches ``bind``'s exit flush: an unended span never + # exports, orphaning its finished children and charging the episode's wall to between_episodes. self._telemetry.seal(clock.now()) raise + finally: + self._retire_worker() + self._reap_worker() - def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # noqa: C901 + def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: while not should_stop.value: # One action per round, mutually exclusive: handle a directive, start the next trial (or exit - # when the plan is done), finish a self-driven trial that is out of budget or done, or step the - # policy. Starting takes its own round so a begin never shares one with a step — inference waits - # for the producer's post-reset observation, which the recorder logs once its open-turn drain has - # cleared the channels. + # when the plan is done), finish one that is out of budget or done, or step the policy. Starting + # takes its own round, so inference waits for the producer's post-reset observation. directive_msg = self.directive.read() # Read every round so the flag clears mid-episode; a press during a trial is consumed, not replayed. manual_msg = self.manual_command.read() @@ -481,27 +598,24 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ assert directive_msg is not None and manual_msg is not None if directive_msg.updated: yield from self._handle_directive(directive_msg.data, clock) - elif not self._running: + elif self._worker is None: if manual_msg.updated and manual_msg.data is not None: - self._emit_now(manual_msg.data, clock) + self._emit(manual_msg.data) elif self._trials is not None: trial = next(self._trials, None) if trial is None: # plan exhausted — let the recorder commit the final episode, then exit yield pimm.Sleep(0.5) break self._begin_episode(trial, clock) - elif self._deadline is not None and (terminal := self._trial_terminal(clock)) is not None: + elif (terminal := self._trial_terminal(clock)) is not None: yield from self._end_episode(clock, terminal) else: try: - yield from self._step(clock) + self._step(self._worker, clock) except pimm.NoValueException: pass - yield self._pace() + self._play(clock) + yield self._pace(clock) - if self._running: + if self._worker is not None: yield from self._finalize_recording(clock) - if self._policy_session: - self._policy_session.close() - # The harness does not own the policy's lifetime: the caller may run several harnesses over - # one policy (a multi-eval sweep), so it closes the policy once, after the last run. diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index ba621f2e4..b62c7da0e 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -25,6 +25,8 @@ class ObservationCodec(Codec): task_field: output key carrying the language prompt at inference. """ + WIRE_NAME = 'observation_codec' + def __init__( self, state: dict[str, dict[str, int]], @@ -98,6 +100,6 @@ def to_spec(self): # Normalized to lists so the spec is identical before and after a wire round-trip. images = {name: [key, list(size)] for name, (key, size) in self._image_configs.items()} return { - 'name': 'observation_codec', + 'name': self.WIRE_NAME, 'args': {'state': self._state, 'images': images, 'task_field': self._task_field}, } diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index ed6755cef..99cfdc4ea 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -222,7 +222,7 @@ def _log_action_series(path: str, arr: np.ndarray, horizon: np.ndarray, base_ns: Each action is stamped at ``base_ns + horizon_i``, where ``base_ns`` is the inference-request time; successive chunks lay out along one clock so a ``TimeSeriesView`` has real axes. This precedes true execution by the inference latency: the harness's - ``ChunkedSchedule`` anchors commands at ``clock.now()`` *after* inference, which a recorder + ``ChunkedSchedule`` anchors commands at ``now()`` *after* inference, which a recorder tap sitting inside it cannot observe. """ arr = np.asarray(arr, dtype=np.float64) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 142539a0d..cb5174436 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -26,9 +26,9 @@ def __init__(self, ws_session: InferenceSession, compress_images: bool = False): self._session = ws_session self._compress_images = compress_images - def _prepare_obs(self, obs: dict[str, Any]) -> dict[str, Any]: + def _prepare_obs(self, obs: cabc.Mapping[str, Any]) -> dict[str, Any]: if not self._compress_images: - return obs + return dict(obs) return {key: self._prepare_value(key, value) for key, value in obs.items()} def _prepare_value(self, key: str, value: Any) -> Any: @@ -42,7 +42,7 @@ def _prepare_value(self, key: str, value: Any) -> Any: return type(value)(self._prepare_value(key, v) for v in value) return value - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: cabc.Mapping[str, Any]) -> list[dict[str, Any]] | None: """Forwards the observation to the remote server and returns the action trajectory. Single-action server responses are wrapped into a 1-element list to honor diff --git a/positronic/policy/spec.py b/positronic/policy/spec.py index f98b65199..f9447a388 100644 --- a/positronic/policy/spec.py +++ b/positronic/policy/spec.py @@ -46,7 +46,7 @@ RestrictImageSize, ) from positronic.policy.observation import ObservationCodec -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack class RemoteMarker(PolicyWrapper): @@ -142,27 +142,25 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) return self._policy -# TODO(hardcoded-keys): every name here is written twice — once as a literal in the wrapper's own ``to_spec`` -# and once as a key below — so a rename in one place silently desyncs the wire. Let each wrapper own its wire -# name as a class attribute and build this mapping from it. ``test_wire_names_match_table`` is what catches a -# desync until then. -# rules-allow: hardcoded-keys — fixing one of the thirteen leaves the rest on the old pattern; the TODO above -# names the whole-table fix, and the test above catches a desync meanwhile WIRE_WRAPPERS: dict[str, type[PolicyWrapper]] = { - 'chunked_schedule': ChunkedSchedule, - 'temporal_stack': TemporalStack, - 'action_timestamp': ActionTimestamp, - 'action_horizon': ActionHorizon, - 'binarize_grip_training': BinarizeGripTraining, - 'binarize_grip_inference': BinarizeGripInference, - 'flip_grip': FlipGrip, - 'restrict_image_size': RestrictImageSize, - 'change_ee_frame': ChangeEEFrame, - 'observation_codec': ObservationCodec, - 'absolute_position_action': AbsolutePositionAction, - 'absolute_joints_action': AbsoluteJointsAction, - 'relative_position_action': RelativePositionAction, - 'joint_delta_action': JointDeltaAction, + wrapper.WIRE_NAME: wrapper + for wrapper in ( + ChunkedSchedule, + StopOnFault, + TemporalStack, + ActionTimestamp, + ActionHorizon, + BinarizeGripTraining, + BinarizeGripInference, + FlipGrip, + RestrictImageSize, + ChangeEEFrame, + ObservationCodec, + AbsolutePositionAction, + AbsoluteJointsAction, + RelativePositionAction, + JointDeltaAction, + ) } diff --git a/positronic/policy/tests/golden_pipeline.json.gz b/positronic/policy/tests/golden_pipeline.json.gz index 43b3afddc..ac1cfbb81 100644 Binary files a/positronic/policy/tests/golden_pipeline.json.gz and b/positronic/policy/tests/golden_pipeline.json.gz differ diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index e7aceab6f..302d5e554 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -13,8 +13,8 @@ * horizon/gating regression -> the state trajectory diverges Everything runs on CPU in a virtual-time world only: no GL/GPU/MuJoCo, no wall -clock in the asserted path (a fixed-float ``inference_latency`` is used so -inference latency is deterministic). +clock in the asserted path (``_SimulatedLatency`` charges a fixed world-time +latency, so inference latency is deterministic). Regenerate the golden after an intentional behavior change: @@ -37,13 +37,14 @@ from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import CartesianPosition, Reset, TrajectoryPlayer +from positronic.drivers.roboarm.command import CartesianPosition, CommandType, Reset +from positronic.drivers.roboarm.tests.fakes import make_robot_state from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation from positronic.geom import Rotation, Transform3D -from positronic.policy.base import Policy, Session +from positronic.policy.base import DelegatingPolicy, DelegatingSession, Now, Policy, Session from positronic.policy.codec import ActionTiming from positronic.policy.harness import Directive, Harness -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.tests.testing_coutils import ManualDriver, drive_scheduler GOLDEN_FILE = Path(__file__).parent / 'golden_pipeline.json.gz' @@ -52,7 +53,7 @@ INITIAL_Q = np.array([0.10, -0.20, 0.30, -0.40, 0.50, -0.60, 0.70], dtype=np.float32) TARGET_POS = np.array([0.50, 0.00, 0.45], dtype=np.float32) -# Fixed deterministic inference latency. Spans >1 control tick (harness loop is +# Fixed deterministic inference latency in world time. Spans >1 control tick (harness loop is # 0.01 s) so the post-inference anchoring effect is observable in recorded ts. INFERENCE_LATENCY_S = 0.05 ACTION_FPS = 15.0 @@ -86,35 +87,55 @@ def new_session(self, context=None, now=None): return _ScriptedSession() -class _FakeRobotState: - """Lossless re-expression of the last applied command over sim-time.""" +class _SimulatedLatency(DelegatingPolicy): + """A fixed inference latency in world time: every chunk is stamped for, and reaches the harness at, + ``latency_sec`` after the call that produced it, whatever the machine took. Sits outermost, so nothing + below sees an observation while a chunk is held. A skip or a stop passes through at once.""" + + def __init__(self, inner: Policy, latency_sec: float): + super().__init__(inner) + self._latency_ns = round(latency_sec * 1e9) + + class _Session(DelegatingSession): + def __init__(self, inner: Session, now: Now, latency_ns: int): + super().__init__(inner) + self._now = now + self._latency_ns = latency_ns + self._held: list | None = None + self._release_at_ns = 0 + + def __call__(self, obs): + # Integer ns, the world's own timeline: a float compare can miss the release instant by one ULP. + now_ns = round(self._now() * 1e9) + if self._held is not None: + if now_ns < self._release_at_ns: + return None + held, self._held = self._held, None + return held + result = self._inner(obs) + if not result: + return result + self._held, self._release_at_ns = result, now_ns + self._latency_ns + return None + + def cancel(self): + self._held = None + super().cancel() - def __init__(self, pos: np.ndarray, q: np.ndarray, status: RobotStatus): - self._pos = pos - self._q = q - self.status = status - - @property - def q(self) -> np.ndarray: - return self._q.copy() - - @property - def dq(self) -> np.ndarray: - return np.zeros(7, dtype=np.float32) - - @property - def ee_pose(self) -> Transform3D: - return Transform3D(translation=self._pos.copy(), rotation=Rotation.identity) + def new_session(self, context=None, now=None): + assert now is not None, 'the harness supplies the clock' + latency_ns = self._latency_ns + return _SimulatedLatency._Session( + self._inner.new_session(context, lambda: now() + latency_ns / 1e9), now, latency_ns + ) class FakeRobot(pimm.ControlSystem): - """Deterministic closed-loop arm: applies the latest command immediately. + """Deterministic closed-loop arm: applies each command as it arrives. - Mirrors ``MujocoSim``'s arm loop (read latest command, apply, emit - state). ``ee_pose`` becomes the applied ``CartesianPosition`` target and the - first three joints track it, so recorded state is a lossless re-expression - of applied commands. Closed loop: the policy's next chunk evolves with this - feedback. + Mirrors ``MujocoSim``'s arm loop (execute on updated, emit state). ``ee_pose`` becomes the applied + ``CartesianPosition`` target and the first three joints track it, so recorded state is a lossless + re-expression of applied commands. Closed loop: the policy's next chunk evolves with this feedback. """ def __init__(self): @@ -122,7 +143,7 @@ def __init__(self): self._q = INITIAL_Q.copy() self._status = RobotStatus.AVAILABLE self._error_pending = False - self.commands = pimm.ControlSystemReceiver(self, default=[]) + self.commands = pimm.ControlSystemReceiver[CommandType](self) self.state = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) @@ -142,23 +163,16 @@ def _apply(self, cmd): def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): self.robot_meta.emit({}) - player = TrajectoryPlayer() while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) if self._status == RobotStatus.ERROR: self._status = RobotStatus.AVAILABLE - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale waypoint. - player.set([]) - else: - cmd = player.advance(clock.now_ns()) - if cmd is not None: - self._apply(cmd) + elif cmd_msg is not None and cmd_msg.updated: + self._apply(cmd_msg.data) if self._error_pending: self._status = RobotStatus.ERROR self._error_pending = False - self.state.emit(_FakeRobotState(self._pos, self._q, self._status)) + self.state.emit(make_robot_state(self._pos, self._q, self._status)) yield pimm.Sleep(CONTROL_PERIOD_S) @@ -167,18 +181,14 @@ class FakeGripper(pimm.ControlSystem): def __init__(self): self._grip = 0.0 - self.target_grip = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.grip = pimm.ControlSystemEmitter(self) def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): - player = TrajectoryPlayer() while not should_stop.value: msg = self.target_grip.read() - if msg.updated: - player.set(msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - self._grip = float(grip) + if msg is not None and msg.updated: + self._grip = float(msg.data) self.grip.emit(self._grip) yield pimm.Sleep(CONTROL_PERIOD_S) @@ -193,7 +203,7 @@ def _run_pipeline(tmp_path: Path) -> dict: embodiment = Embodiment( descriptor='', observations={ - 'robot_state': Observation(robot.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot.state, Serializers.robot_state), keys.GRIP: Observation(gripper.grip, None), }, commands={ @@ -202,8 +212,13 @@ def _run_pipeline(tmp_path: Path) -> dict: }, static_meta=dict(ROBOT_STATIC_META), meta_source=robot.robot_meta, + # The fake robot's control-period sleep is this world's sole time-master — the shape a sim eval + # runs in. + simulated=True, + ) + harness = Harness( + _SimulatedLatency((StopOnFault() | ChunkedSchedule()).wrap(policy), INFERENCE_LATENCY_S), embodiment ) - harness = Harness(ChunkedSchedule().wrap(policy), embodiment) ds_agent = wire.wire_embodiment(world, harness, embodiment, ds_writer, TimeMode.MESSAGE) world.connect(harness.ds_command, ds_agent.command) directive_em = world.pair(harness.directive) @@ -211,7 +226,7 @@ def _run_pipeline(tmp_path: Path) -> dict: # Robot/gripper emit state every tick, so the script only drives the # episode lifecycle and the one-shot error injection. script = [ - (partial(directive_em.emit, Directive.RUN(task='golden', inference_latency=INFERENCE_LATENCY_S)), 0.0), + (partial(directive_em.emit, Directive.RUN(task='golden')), 0.0), (None, 1.5), # several reactive inference + chunk/horizon cycles (robot.inject_error, 0.0), # one-shot error: that frame is dropped, then inference resumes (None, 0.5), diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 68ddd9008..df300c832 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,3 +1,5 @@ +import threading +import time from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -12,27 +14,17 @@ from positronic.dataset.serializers import Serializers from positronic.drivers import roboarm from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import ( - CartesianDelta, - CartesianPosition, - JointDelta, - JointPosition, - Reset, - TrajectoryPlayer, - _compose_delta, - from_wire, - reduce, - to_wire, -) +from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model +from positronic.drivers.roboarm.tests.fakes import make_robot_state from positronic.eval import Command, Embodiment, Observation, Task from positronic.geom import Rotation, Transform3D from positronic.offboard.client import InferenceSession -from positronic.policy.base import Policy, Session +from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import Directive, DirectiveType, Harness, _assert_anchored +from positronic.policy.harness import Directive, DirectiveType, Harness, _InferenceWorker from positronic.policy.remote import RemoteSession -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -52,7 +44,7 @@ def _eval_pass(run_id: str): CAM = 'image.cam' -def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None) -> Embodiment: +def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simulated=False) -> Embodiment: """Minimal Franka-shaped embodiment for harness unit tests. The sources/dests are no-ops: these tests pair the harness ports directly @@ -60,7 +52,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None) -> E home values, descriptor — is read by the Harness. """ observations = { - 'robot_state': Observation(pimm.NoOpEmitter(), Serializers.robot_state), + keys.ROBOT_STATE: Observation(pimm.NoOpEmitter(), Serializers.robot_state), keys.GRIP: Observation(pimm.NoOpEmitter(), None), } for cam in cameras: @@ -69,7 +61,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None) -> E keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command), 'target_grip': Command(pimm.NoOpReceiver(), 0.0, None), } - return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter()) + return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter(), simulated=simulated) class _SpySession(Session): @@ -128,16 +120,16 @@ def __init__( self.command = command self.target_grip = float(target_grip) self.last_obs: dict[str, Any] | None = None - self.observations: list[dict[str, object]] = [] + self.observations: list[dict[str, Any]] = [] self.reset_calls = 0 - self.last_reset_context = None + self.last_reset_context: dict[str, Any] | None = None self._meta: dict[str, object] = meta or {} @property def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None) -> Session: self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -207,41 +199,38 @@ def new_session(self, context=None, now=None) -> RemoteSession: return RemoteSession(_FakeInferenceSession(action)) -class FakeRobotState: - def __init__(self, translation: np.ndarray, joints: np.ndarray, status: RobotStatus) -> None: - self.ee_pose = Transform3D(translation=translation, rotation=Rotation.identity) - self.q = joints - self.dq = np.zeros_like(joints) - self.status = status - - @pytest.fixture def world(): with pimm.World(virtual_time=True) as w: yield w -def make_robot_state(translation, joints, status=RobotStatus.AVAILABLE) -> FakeRobotState: - translation = np.asarray(translation, dtype=np.float32) - joints = np.asarray(joints, dtype=np.float32) - return FakeRobotState(translation, joints, status) - - def emit_ready_payload(frame_emitter, robot_emitter, grip_emitter, robot_state): - frame_adapter = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.uint8) + frame_adapter = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) frame_adapter.array[:] = np.zeros((2, 2, 3), dtype=np.uint8) frame_emitter.emit(frame_adapter) robot_emitter.emit(robot_state) grip_emitter.emit(0.25) +class _Pacer(pimm.ControlSystem): + """Stands in for the simulator: the sole time-master, sleeping one control period every turn.""" + + def __init__(self, period: float = 0.005): + self._period = period + + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): + while not should_stop.value: + yield pimm.Sleep(self._period) + + def _pair_all(world, harness): """Pair all harness signals and return a dict of test handles.""" ds_recorder = RecordingEmitter() harness.ds_command._bind(ds_recorder) return { 'frame_em': world.pair(harness.observations[CAM]), - 'robot_em': world.pair(harness.observations['robot_state']), + 'robot_em': world.pair(harness.observations[keys.ROBOT_STATE]), 'grip_em': world.pair(harness.observations[keys.GRIP]), 'directive_em': world.pair(harness.directive), 'command_rx': world.pair(harness.commands[keys.ROBOT_COMMAND]), @@ -260,39 +249,27 @@ def _ds_types(p) -> list[DsWriterCommandType]: def _last_command(p): - """Extract the last robot command from the trajectory signal.""" + """The latest robot command the harness put on the channel.""" msg = p['command_rx'].read() - if msg is None or msg.data is None: - return None - traj = msg.data # list[tuple[float, CommandType]] - return traj[-1][1] if traj else None + assert msg is not None, 'no robot command was emitted' + return msg.data def _last_grip(p): - """Extract the last grip value from the grip trajectory signal.""" - msg = p['grip_rx'].read() - if msg is None or msg.data is None: - return None - traj = msg.data # list[tuple[float, float]] - return traj[-1][1] if traj else None - - -def _all_grips(p): - """Extract all grip values from the grip trajectory signal.""" + """The latest grip target the harness put on the channel.""" msg = p['grip_rx'].read() - if msg is None or msg.data is None: - return [] - return [g for _, g in msg.data] + assert msg is not None, 'no grip target was emitted' + return msg.data def _emitted_commands(recorder): - """All robot commands across a recorder's non-empty emitted trajectories.""" - return [cmd for _ts, traj in recorder.emitted if traj for _t, cmd in traj] + """Every robot command a recorder saw, in emission order.""" + return [cmd for _ts, cmd in recorder.emitted] def _emitted_grips(recorder): - """All grip targets across a recorder's non-empty emitted trajectories.""" - return [g for _ts, traj in recorder.emitted if traj for _t, g in traj] + """Every grip target a recorder saw, in emission order.""" + return [grip for _ts, grip in recorder.emitted] @pytest.mark.timeout(3.0) @@ -307,7 +284,7 @@ def test_harness_emits_cartesian_move(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -331,7 +308,7 @@ def test_harness_emits_cartesian_move(world): np.testing.assert_allclose(obs[keys.JOINT_VEL], np.zeros_like(robot_state.q)) assert obs[keys.GRIP] == pytest.approx(0.25) assert obs[keys.TASK] == 'stack-blocks' - assert obs['descriptor'] == '' # no descriptor passed -> empty string reaches the policy + assert obs[keys.DESCRIPTOR] == '' # no descriptor passed -> empty string reaches the policy # Recording == canonical policy I/O: the policy sees the same ``robot_state`` serializer # the dataset records. wall/obs timestamps carry volatile values, so lock the stable key set. assert set(obs) - {keys.WALL_TIME_NS, keys.OBS_TIME_NS} == { @@ -340,11 +317,11 @@ def test_harness_emits_cartesian_move(world): keys.JOINT_VEL, keys.EE_POSE, keys.GRIP, + keys.ROBOT_FAULT, keys.TASK, - 'descriptor', + keys.DESCRIPTOR, } - # Last non-empty command (a trailing ``[]`` cancel is emitted on shutdown). cmds = _emitted_commands(cmd_recorder) assert cmds, 'no robot command emitted' cmd = cmds[-1] @@ -366,7 +343,7 @@ def test_harness_passes_descriptor_to_policy(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -381,7 +358,7 @@ def test_harness_passes_descriptor_to_policy(world): drive_scheduler(scheduler, steps=20) assert policy.last_obs is not None - assert policy.last_obs['descriptor'] == 'mujoco.franka' + assert policy.last_obs[keys.DESCRIPTOR] == 'mujoco.franka' @pytest.mark.timeout(3.0) @@ -396,7 +373,7 @@ def test_robot_model_stays_out_of_the_observation(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations['grip']) directive_em = world.pair(harness.directive) @@ -467,7 +444,7 @@ def test_harness_waits_for_complete_inputs(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -856,18 +833,15 @@ def run(self, should_stop, clock): @pytest.mark.timeout(3.0) def test_done_after_deadline_is_a_timeout(world): - """The deadline is hard: a ``done`` delivered past it (here during the latency sleep) records as a - timeout — ``eval.terminated`` False, payload dropped — not a late stop-signal success.""" + """The deadline is hard: a ``done`` delivered past it records as a timeout — ``eval.terminated`` False, + payload dropped — not a late stop-signal success.""" policy = StubPolicy() - harness = Harness( - policy, make_embodiment(), task=Task(instruction='t', timeout=0.05), trials=[{'inference_latency': 0.2}] - ) + harness = Harness(policy, make_embodiment(), task=Task(instruction='t', timeout=0.05), trials=[{}]) p = _pair_all(world, harness) done_em = world.pair(harness.done) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) - # Obs starts inference + the 0.2s latency sleep; the 0.05s deadline lapses during it, and done is - # delivered at ~0.1s — past the deadline but before the harness next polls. The timeout must win. + # The 0.05s deadline lapses first; done lands at ~0.1s, after the trial has already timed out. driver = ManualDriver([ (partial(emit_ready_payload, p['frame_em'], p['robot_em'], p['grip_em'], robot_state), 0.1), (partial(done_em.emit, {keys.EVAL_SUCCESS: True}), 0.3), @@ -929,12 +903,14 @@ def test_trial_plan_self_drives(world): @pytest.mark.timeout(3.0) -def test_timeout_crossed_during_latency_sleep_drops_chunk(world): - """A chunk whose latency sleep crosses the deadline is dropped, never emitted.""" - policy = StubPolicy() - # The 0.2s latency sleep crosses the 0.05s deadline before the chunk is emitted. +def test_timeout_during_inference_drops_the_chunk(world): + """A trial whose deadline lapses while the model is still inside its call ends with the call in flight: + the trajectory it eventually returns is discarded, never emitted past the advertised termination point.""" harness = Harness( - policy, make_embodiment(), task=Task(instruction='test', timeout=0.05), trials=[{'inference_latency': 0.2}] + ChunkedSchedule().wrap(SlowPolicy(wall_sec=0.3)), # the call runs well past the deadline + make_embodiment(simulated=True), + task=Task(instruction='test', timeout=0.05), + trials=[{keys.INFERENCE_LATENCY: True}], ) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() @@ -944,21 +920,20 @@ def test_timeout_crossed_during_latency_sleep_drops_chunk(world): harness.ds_command._bind(ds_recorder) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([(partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), (None, 0.3)]) - scheduler = world.start([harness, driver]) - drive_scheduler(scheduler, steps=200) + scheduler = world.start([harness, driver, _Pacer()]) + drive_scheduler(scheduler, steps=2000) stops = [data for _, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] assert len(stops) == 1 assert stops[0].static_data[keys.EVAL_TERMINATED] is False - # The post-deadline chunk must not reach the drivers: the only non-empty emissions are the homing - # Reset / home grip from the startup home and the timeout FINISH. + # The only commands are the homing Reset / home grip from the startup home and the timeout FINISH. assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) assert _emitted_grips(grip_recorder) == [0.0, 0.0] @@ -1011,6 +986,105 @@ def test_run_while_running_is_ignored(world): assert policy.reset_calls == 1 +class _FrameWatchingPolicy(Policy): + """Reads its camera frame at both ends of a slow call, so a rewrite underneath it shows up as a difference.""" + + def __init__(self, wall_sec: float): + self._wall_sec = wall_sec + self.seen: list[tuple[np.ndarray, np.ndarray]] = [] + + def new_session(self, context=None, now=None): + return _FrameWatchingPolicy._Session(self) + + class _Session(Session): + def __init__(self, policy: '_FrameWatchingPolicy'): + self._policy = policy + + def __call__(self, obs): + entry = np.array(obs[CAM]) + time.sleep(self._policy._wall_sec) + self._policy.seen.append((entry, np.array(obs[CAM]))) + return None + + +@pytest.mark.timeout(10.0) +def test_a_producer_reusing_its_buffer_cannot_rewrite_a_pending_observation(world): + """A camera renders into the array behind the adapter it re-emits, and a wall-charged call runs while the + loop yields — so the observation handed to the worker has to be its own copy.""" + policy = _FrameWatchingPolicy(wall_sec=0.3) + harness = Harness(policy, make_embodiment()) + p = _pair_all(world, harness) + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + frame = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) + + def emit_frame(fill: int): + frame.array[:] = np.full((2, 2, 3), fill, dtype=np.uint8) + p['frame_em'].emit(frame) # the same adapter every time, as a camera does + p['robot_em'].emit(robot_state) + p['grip_em'].emit(0.25) + + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: True})), 0.0), + (partial(emit_frame, 1), 0.01), + (partial(emit_frame, 9), 0.05), # rewrites the buffer while the first call is still running + (None, 0.4), + ]) + + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=60) + + assert policy.seen, 'the policy was never called' + entry, exit_ = policy.seen[0] + np.testing.assert_array_equal(entry, exit_, 'the observation was rewritten while the call was in flight') + + +class _AbandonedCallPolicy(Policy): + """Records the order of session openings and call completions, with a call that outlives its episode.""" + + def __init__(self, wall_sec: float): + self._wall_sec = wall_sec + self.events: list[str] = [] + + def new_session(self, context=None, now=None): + self.events.append('open') + return _AbandonedCallPolicy._Session(self) + + class _Session(Session): + def __init__(self, policy: '_AbandonedCallPolicy'): + self._policy = policy + + def __call__(self, obs): + time.sleep(self._policy._wall_sec) + self._policy.events.append('answered') + return None + + +@pytest.mark.timeout(10.0) +def test_a_new_episode_waits_out_the_call_the_last_one_abandoned(world): + """An in-process policy is one model across episodes, so opening a session must not overtake a call + still inside the previous one — ``new_session`` resets the object that call is using.""" + policy = _AbandonedCallPolicy(wall_sec=0.4) + harness = Harness(policy, make_embodiment()) + p = _pair_all(world, harness) + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + emit_obs = partial(emit_ready_payload, p['frame_em'], p['robot_em'], p['grip_em'], robot_state) + + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='ep1')), 0.0), + (emit_obs, 0.01), # the observation that puts a call on the worker + (partial(p['directive_em'].emit, Directive.FINISH()), 0.02), # while that call is still running + (partial(p['directive_em'].emit, Directive.RUN(task='ep2')), 0.02), + (None, 0.02), + ]) + + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=40) + + assert policy.events[:3] == ['open', 'answered', 'open'], ( + f'the second session opened before the abandoned call answered: {policy.events}' + ) + + @pytest.mark.timeout(3.0) def test_run_calls_policy_reset_with_context(world): policy = StubPolicy() @@ -1043,37 +1117,35 @@ def reset(_context): scheduler = world.start([harness]) drive_scheduler(scheduler, steps=200) + assert policy.last_reset_context is not None assert policy.last_reset_context[keys.TASK] == 'resolved-on-reset' -@pytest.mark.timeout(3.0) -def test_finish_cancels_buffered_trajectory_before_stop_episode(world): - """FINISH must cancel the recording's trajectory tail *before* `STOP_EPISODE`. +class _LabeledRecorder(pimm.SignalEmitter): + """Records emissions from several channels into one shared list, so their order is comparable.""" - `STOP_EPISODE` calls `flush()` on `TrajectoryOverrideSerializer`, which - commits whatever is still buffered. The harness must emit `[]` on - `robot_command`/`target_grip` first, so the serializer drops its tail and - canceled waypoints are not recorded. - """ + def __init__(self, label, events): + self._label = label + self._events = events - class _LabeledRecorder(pimm.SignalEmitter): - def __init__(self, label, events): - self._label = label - self._events = events + def emit(self, data, ts: int = -1): + self._events.append((self._label, data)) - def emit(self, data, ts: int = -1): - self._events.append((self._label, data)) - events: list[tuple[str, object]] = [] +@pytest.mark.timeout(3.0) +def test_finish_stops_playing_the_live_chunk(world): + """FINISH drops the schedule the harness is playing: the chunk's remaining waypoints never reach the + devices, and the only command after the recorder's STOP is the home the close emits.""" policy = ChunkPolicy() wrapped = ActionTimestamp(fps=5.0).wrap(policy) # 1.8 s chunk — won't drain before FINISH harness = Harness(wrapped, make_embodiment()) + events: list[tuple[str, object]] = [] harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) harness.commands['target_grip']._bind(_LabeledRecorder('target_grip', events)) harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1083,33 +1155,21 @@ def emit(self, data, ts: int = -1): (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), (None, 0.1), (partial(directive_em.emit, Directive.FINISH()), 0.0), - (None, 0.1), + (None, 0.5), ] scheduler = world.start([harness, ManualDriver(script)]) - drive_scheduler(scheduler, steps=200) + drive_scheduler(scheduler, steps=400) - cancels = [i for i, (lbl, data) in enumerate(events) if lbl == keys.ROBOT_COMMAND and data == []] - stops = [ - i - for i, (lbl, data) in enumerate(events) - if lbl == 'ds_command' and getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE - ] - assert cancels, 'FINISH did not emit a cancel on robot_command' + stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'FINISH did not emit STOP_EPISODE' - assert cancels[0] < stops[0], ( - f'cancel ({cancels[0]}) must precede STOP_EPISODE ({stops[0]}); otherwise flush() commits canceled waypoints' - ) + grips_after = [data for lbl, data in events[stops[0] :] if lbl == keys.TARGET_GRIP] + assert grips_after == [0.0], f'the cancelled chunk kept playing past FINISH: {grips_after}' @pytest.mark.timeout(3.0) -def test_empty_chunk_cancels_both_robot_and_grip(world): - """A session returning ``[]`` must cancel *both* driver buffers. - - Empty action chunk is the session-level cancel signal (per the - ``Session.__call__`` contract). If only ``robot_command`` gets ``[]`` while - ``target_grip`` is skipped, the gripper ``TrajectoryPlayer`` keeps draining - stale waypoints — a partial cancel that's worse than no cancel. - """ +def test_empty_trajectory_leaves_every_channel_holding(world): + """A trajectory with no waypoints schedules nothing on any channel, so every device holds where the + startup home left it rather than one channel draining on while another stops.""" class _EmptyChunkSession(Session): def __call__(self, obs): @@ -1127,7 +1187,7 @@ def new_session(self, context=None, now=None): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1140,15 +1200,12 @@ def new_session(self, context=None, now=None): scheduler = world.start([harness, ManualDriver(script)]) drive_scheduler(scheduler, steps=200) - cmd_emits = [data for _ts, data in cmd_recorder.emitted] - grip_emits = [data for _ts, data in grip_recorder.emitted] - assert [] in cmd_emits, 'empty chunk did not cancel robot_command buffer' - assert [] in grip_emits, 'empty chunk did not cancel target_grip buffer' + assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) # only the startup home + assert _emitted_grips(grip_recorder) == [0.0] @pytest.mark.timeout(3.0) -def test_harness_clears_trajectory_on_home(world): - """Verify that HOME resets trajectory state so next RUN gets a fresh chunk.""" +def test_harness_clears_trajectory_on_abort(world): policy = ChunkPolicy() harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) @@ -1162,8 +1219,7 @@ def test_harness_clears_trajectory_on_home(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=5) - grips = _all_grips(p) - assert grips[0] >= 100.0, f'Expected chunk 1, got {grips}' + assert _last_grip(p) >= 100.0, 'Expected chunk 1' p['directive_em'].emit(Directive.ABORT()) drive_scheduler(scheduler, steps=2) @@ -1174,13 +1230,11 @@ def test_harness_clears_trajectory_on_home(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=4) - grips = _all_grips(p) - assert grips[0] >= 200.0, f'Expected chunk 2 (>= 200.0), got {grips}. Trajectory clearing failed!' + assert _last_grip(p) >= 200.0, 'Expected chunk 2; trajectory clearing failed' @pytest.mark.timeout(3.0) def test_harness_clears_trajectory_on_run(world): - """Verify that RUN resets trajectory state so a fresh chunk is emitted.""" policy = ChunkPolicy() harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) @@ -1194,8 +1248,7 @@ def test_harness_clears_trajectory_on_run(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=5) - grips = _all_grips(p) - assert grips[0] >= 100.0 + assert _last_grip(p) >= 100.0 p['directive_em'].emit(Directive.RUN(task='test-restart')) drive_scheduler(scheduler, steps=1) @@ -1203,8 +1256,7 @@ def test_harness_clears_trajectory_on_run(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=4) - grips = _all_grips(p) - assert grips[0] >= 200.0, f'Expected chunk 2 (>= 200.0), got {grips}. Trajectory clearing on RUN failed!' + assert _last_grip(p) >= 200.0, 'Expected chunk 2; trajectory clearing on RUN failed' @pytest.mark.timeout(3.0) @@ -1224,8 +1276,7 @@ def test_harness_skips_inference_on_error(world): drive_scheduler(scheduler, steps=1) emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], state_ok) drive_scheduler(scheduler, steps=3) - grips = _all_grips(p) - assert grips[0] >= 100.0 + assert _last_grip(p) >= 100.0 obs_before = len(policy.observations) p['robot_em'].emit(state_err) @@ -1233,9 +1284,8 @@ def test_harness_skips_inference_on_error(world): assert len(policy.observations) == obs_before # the errored state is never fed to the policy emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], state_ok) - drive_scheduler(scheduler, steps=3) - grips = _all_grips(p) - assert grips[0] >= 200.0 + drive_scheduler(scheduler, steps=20) # long enough for the first chunk to play out and the next to land + assert _last_grip(p) >= 200.0 def test_directive_preserves_payload(): @@ -1284,64 +1334,6 @@ def test_cartesian_delta_applies_in_world_frame(): assert not np.allclose(target.translation, (current * delta).translation) # guards against body-frame compose -def test_reduce_accumulates_due_cartesian_deltas(): - # Rotations about different axes so the world-frame compose is non-commutative -- this pins the fold order - # (apply d0 then d1), not just that a fold happened. - d0 = Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.array([0.3, 0.0, 0.0]))) - d1 = Transform3D(np.array([0.02, 0.01, 0.0]), Rotation.from_rotvec(np.array([0.0, 0.0, 0.2]))) - out = reduce([(10, CartesianDelta(d0)), (20, CartesianDelta(d1))]) - assert isinstance(out, CartesianDelta) - expected = _compose_delta(d0, d1) # two due deltas catch up as their world-frame compose, not last-wins - np.testing.assert_allclose(out.delta.translation, expected.translation) - np.testing.assert_allclose(out.delta.rotation.as_quat, expected.rotation.as_quat, atol=1e-12) - assert not np.allclose(out.delta.rotation.as_quat, _compose_delta(d1, d0).rotation.as_quat) - - -def test_reduce_sums_due_joint_deltas(): - out = reduce([(10, JointDelta(np.array([0.1, -0.2, 0.3]))), (20, JointDelta(np.array([0.0, 0.2, -0.1])))]) - assert isinstance(out, JointDelta) - np.testing.assert_allclose(out.velocities, [0.1, 0.0, 0.2]) - - -def test_reduce_absolute_run_keeps_last(): - p0 = CartesianPosition(Transform3D(np.array([0.1, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - p1 = JointPosition(np.array([0.2, 0.0, 0.0])) - assert reduce([(10, p0), (20, p1)]) is p1 - - -def test_reduce_raises_on_absolute_delta_mix(): - cart_pos = CartesianPosition(Transform3D(np.zeros(3), Rotation.from_rotvec(np.zeros(3)))) - cart_delta = CartesianDelta(Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - joint_pos = JointPosition(np.zeros(3)) - joint_delta = JointDelta(np.array([0.1, 0.0, 0.0])) - with pytest.raises(ValueError): - reduce([(10, cart_pos), (20, cart_delta)]) - with pytest.raises(ValueError): - reduce([(10, cart_delta), (20, cart_pos)]) - with pytest.raises(ValueError): # JointPosition then JointDelta: the delta has no faithful anchor to fold onto - reduce([(10, joint_pos), (20, joint_delta)]) - - -def test_reduce_raises_on_mixed_delta_spaces(): - cart_delta = CartesianDelta(Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - joint_delta = JointDelta(np.array([0.1, 0.0, 0.0])) - with pytest.raises(ValueError): - reduce([(10, cart_delta), (20, joint_delta)]) - with pytest.raises(ValueError): - reduce([(10, joint_delta), (20, cart_delta)]) - - -def test_trajectory_player_accumulates_missed_deltas(): - d0 = Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3))) - d1 = Transform3D(np.array([0.02, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3))) - player = TrajectoryPlayer(reduce=reduce) - player.set([(10, CartesianDelta(d0)), (20, CartesianDelta(d1))]) - out = player.advance(20) # both waypoints due in one tick -> summed, not dropped to the last - assert isinstance(out, CartesianDelta) - np.testing.assert_allclose(out.delta.translation, [0.03, 0.0, 0.0]) - assert player.advance(30) is None - - @pytest.mark.parametrize('status', [RobotStatus.RESETTING, RobotStatus.ERROR]) def test_robot_state_serializer_drops_not_ready(status): state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=status) @@ -1350,39 +1342,29 @@ def test_robot_state_serializer_drops_not_ready(status): def test_robot_state_serializer_available_has_no_error_key(): state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.AVAILABLE) - assert set(Serializers.robot_state(state)) == {'.q', '.dq', '.ee_pose'} + serialized = Serializers.robot_state(state) + assert serialized is not None + assert set(serialized) == {'.q', '.dq', '.ee_pose'} @pytest.mark.timeout(3.0) -def test_shutdown_cancels_trajectory_before_stop(world): - """Shutdown while recording must cancel buffered trajectories before STOP_EPISODE. - - ``STOP_EPISODE`` flushes ``TrajectoryOverrideSerializer``; without a prior - cancel it would commit the unexecuted tail of an in-flight chunk (the - FINISH/RUN paths already cancel first). - """ +def test_shutdown_stops_playing_the_live_chunk(world): + """Shutdown while recording drops the schedule too: the unplayed tail of the live chunk never reaches + the devices after the recorder's STOP.""" events: list[tuple[str, object]] = [] - - class _LabeledRecorder(pimm.SignalEmitter): - def __init__(self, label): - self._label = label - - def emit(self, data, ts: int = -1): - events.append((self._label, data)) - wrapped = ActionTimestamp(fps=5.0).wrap(ChunkPolicy()) # 1.8 s chunk — won't drain before shutdown harness = Harness(wrapped, make_embodiment()) - harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND)) - harness.commands['target_grip']._bind(_LabeledRecorder('target_grip')) - harness.ds_command._bind(_LabeledRecorder('ds_command')) + harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) + harness.commands[keys.TARGET_GRIP]._bind(_LabeledRecorder(keys.TARGET_GRIP, events)) + harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) - # RUN + a complete obs buffers a chunk; the driver then ends, which makes the + # RUN + a complete obs schedules a chunk; the driver then ends, which makes the # world signal shutdown while still recording — exercising the run() finalizer. driver = ManualDriver([ (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), @@ -1392,15 +1374,9 @@ def emit(self, data, ts: int = -1): scheduler = world.start([harness, driver]) drive_scheduler(scheduler, steps=200) - cancels = [i for i, (lbl, data) in enumerate(events) if lbl == keys.ROBOT_COMMAND and data == []] - stops = [ - i - for i, (lbl, data) in enumerate(events) - if lbl == 'ds_command' and getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE - ] - assert cancels, 'shutdown did not cancel robot_command' + stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'shutdown did not emit STOP_EPISODE' - assert cancels[0] < stops[0], 'cancel must precede STOP_EPISODE on shutdown' + assert not [lbl for lbl, _ in events[stops[0] :] if lbl == keys.TARGET_GRIP] @pytest.mark.timeout(5.0) @@ -1638,15 +1614,569 @@ def episode(gap_steps: int) -> float: def test_unanchored_chunk_is_refused(): """A stack that never anchored leaves chunk-relative stamps, which read as decades before now.""" with pytest.raises(ValueError, match='not anchoring'): - _assert_anchored([{'timestamp': 0.0}], now=1.7e9) + Harness._assert_anchored([{'timestamp': 0.0}], now=1.7e9) def test_doubly_anchored_chunk_is_refused(): """Two schedulers each add the clock, putting the chunk a lifetime ahead.""" with pytest.raises(ValueError, match='not anchoring'): - _assert_anchored([{'timestamp': 3.4e9}], now=1.7e9) + Harness._assert_anchored([{'timestamp': 3.4e9}], now=1.7e9) def test_anchored_chunk_passes(): """A real chunk spans seconds around now, and a late action sits just behind it.""" - _assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) + Harness._assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) + + +@pytest.mark.parametrize(('expired', 'scheduled'), [(True, False), (False, True)]) +def test_a_reply_is_scheduled_only_while_the_trial_still_has_budget(world, expired, scheduled): + """A trial advertises the instant it stops at. A call whose rounds in flight carried the world past that + instant has its chunk dropped instead of placed, and ``_run`` finishes the trial on the next round.""" + policy = StubPolicy() + harness = Harness(policy, make_embodiment()) + now = world.clock.now() + harness._deadline = now - 1.0 if expired else now + 1.0 + harness._worker = _InferenceWorker(policy, {}, charge_wall=True, clock=world.clock) + harness._worker.submit({}) + + harness._throttle_and_reschedule(harness._worker, world.clock) + + assert bool(harness._schedules[keys.ROBOT_COMMAND]) is scheduled + + +class _SlowSession(Session): + """A session whose inference costs ``wall_sec`` of real time and returns a fixed-length chunk.""" + + def __init__(self, wall_sec: float, span_sec: float, steps: int): + self._wall_sec = wall_sec + self._span_sec = span_sec + self._steps = steps + + def __call__(self, obs): + time.sleep(self._wall_sec) + dt = self._span_sec / self._steps + pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) + return [ + { + keys.ROBOT_COMMAND: CartesianPosition(pose=pose), + keys.TARGET_GRIP: float(i), + keys.ACTION_TIMESTAMP: i * dt, + } + for i in range(self._steps) + ] + + +class SlowPolicy(Policy): + def __init__(self, wall_sec: float = 0.0, span_sec: float = 0.2, steps: int = 10): + self._wall_sec = wall_sec + self._span_sec = span_sec + self._steps = steps + + def new_session(self, context=None, now=None): + return _SlowSession(self._wall_sec, self._span_sec, self._steps) + + +class _ReplanEarly(PolicyWrapper): + """Infers on the first observation and again halfway through the chunk it returned. + + The re-query-before-exhaustion shape (RTC, temporal ensembling) that the substrate exists for: unlike + ``ChunkedSchedule`` it leaves waypoints to play while a call is in flight. + """ + + class _Session(DelegatingSession): + def __init__(self, inner: Session, now): + super().__init__(inner) + self._now = now + self._replan_at: float | None = None + + def __call__(self, obs): + t0 = obs[keys.OBS_TIME_NS] / 1e9 + if self._replan_at is not None and t0 < self._replan_at: + return None + result = self._inner(obs) + assert result is not None, 'the inner policy of this test wrapper always returns a chunk' + anchor = self._now() + result = [{**action, keys.ACTION_TIMESTAMP: anchor + action[keys.ACTION_TIMESTAMP]} for action in result] + self._replan_at = t0 + (result[-1][keys.ACTION_TIMESTAMP] - t0) / 2 + return result + + def wrap_session(self, inner: Session, context, now): + assert now is not None # the harness always passes its clock + return _ReplanEarly._Session(inner, now) + + +class _TimedRecorder(pimm.SignalEmitter): + """Records each emission against the world clock, so a test can read when a command went out.""" + + def __init__(self, clock: pimm.Clock): + self._clock = clock + self.emitted: list[tuple[float, Any]] = [] + + def emit(self, data, ts: int = -1): + self.emitted.append((self._clock.now(), data)) + + +def _run_episode( + world, policy, wrapper, *, latency, simulated=True, steps=4000, run_sec=1.5 +) -> list[tuple[float, Any]]: + """One trial whose context asks for ``latency``; returns the grip commands with the world time each went + out at. A sim trial runs against a pacer, the sole time-master a real rig doesn't need.""" + harness = Harness(wrapper.wrap(policy), make_embodiment(simulated=simulated)) + grip_recorder = _TimedRecorder(world.clock) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: latency})), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, run_sec), + ]) + systems = [harness, driver, _Pacer()] if simulated else [harness, driver] + drive_scheduler(world.start(systems), steps=steps) + return grip_recorder.emitted[1:] # drop the startup home + + +@pytest.mark.timeout(20.0) +def test_default_latency_pauses_the_world_for_the_call(world): + """Sim's default charges nothing: the world does not advance while the model runs, so the chunk is + anchored at the observation's own instant however long the call really took.""" + played = _run_episode(world, SlowPolicy(wall_sec=0.05), ChunkedSchedule(), latency=False) + + assert played, 'no command was played' + assert played[0][0] < 0.01, f'the world advanced during the call: first command at {played[0][0]}s' + + +@pytest.mark.timeout(20.0) +def test_measured_latency_charges_the_calls_own_wall_duration(world): + """``inference_latency=True`` charges the world what the model really took, so a slow server is scored + as slow — at the cost of a trace that inherits the machine's noise.""" + played = _run_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=True) + + assert played, 'no command was played' + assert played[0][0] >= 0.2, f'first command at {played[0][0]}s, under the 0.2s the call took' + + +@pytest.mark.timeout(20.0) +def test_a_real_rig_pays_wall_time_whatever_the_trial_asks_for(world): + """The knob simulates a trial; a real rig pays what its calls take, so a context asking for the world to + be held (the eval CLI writes the flag into every trial) does not hold it.""" + played = _run_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=False, simulated=False) + + assert played, 'no command was played' + assert played[0][0] >= 0.2, f'first command at {played[0][0]}s, under the 0.2s the call took' + + +class _ObservedTicks(PolicyWrapper): + """Records the observation instant of every call that reaches it.""" + + def __init__(self): + self.seen: list[float] = [] + + class _Session(DelegatingSession): + def __init__(self, inner: Session, seen: list[float]): + super().__init__(inner) + self._seen = seen + + def __call__(self, obs): + self._seen.append(obs[keys.OBS_TIME_NS] / 1e9) + return self._inner(obs) + + def wrap_session(self, inner: Session, context, now): + return _ObservedTicks._Session(inner, self.seen) + + +@pytest.mark.timeout(30.0) +def test_the_wrappers_see_every_tick(world): + """What a temporal stack records: nothing keeps an observation from the wrappers above the scheduler — + not the machine inside the model call, and not the rounds in between.""" + ticks = _ObservedTicks() + _run_episode( + world, SlowPolicy(wall_sec=0.01, span_sec=0.3, steps=15), ticks | ChunkedSchedule(), latency=False, run_sec=1.0 + ) + + period = 0.005 # the pacer's control period + gaps = [round(b - a, 4) for a, b in zip(ticks.seen, ticks.seen[1:], strict=False)] + assert gaps and all(gap <= period for gap in gaps), gaps + + +@pytest.mark.timeout(20.0) +def test_harness_keeps_playing_while_a_call_is_in_flight(world): + """A wrapper that replans before its chunk is exhausted leaves waypoints due during inference, and the + harness emits them on time instead of standing still until the model answers.""" + played = _run_episode(world, SlowPolicy(wall_sec=0.15, span_sec=0.4, steps=20), _ReplanEarly(), latency=True) + + # The first chunk lands at ~0.15s and spans 0.4s; the second call starts halfway through it (~0.28s) and + # runs for 0.15s the world runs through; the waypoints due in that window have to keep going out. + during = [t for t, _ in played if 0.29 <= t < 0.4] + assert len(during) >= 3, f'the harness stopped playing during inference: {[t for t, _ in played]}' + + +@pytest.mark.timeout(3.0) +def test_a_faulted_arm_reaches_the_policy_without_its_state(world): + """A fault is not swallowed as "no observation": it goes up to the policy stack, which owns what to do + about it. The arm's own entries are absent — a faulted arm has no sample to give.""" + policy = SpyPolicy() + harness = Harness(policy, make_embodiment()) + p = _pair_all(world, harness) + + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='test')), 0.0), + ( + partial( + emit_ready_payload, + p['frame_em'], + p['robot_em'], + p['grip_em'], + make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR), + ), + 0.01, + ), + (None, 0.02), + ]) + + drive_scheduler(world.start([harness, driver]), steps=40) + + assert policy.last_obs is not None, 'the fault never reached the policy' + assert policy.last_obs[keys.ROBOT_FAULT] is True + assert keys.JOINTS not in policy.last_obs + assert keys.EE_POSE not in policy.last_obs + + +@pytest.mark.timeout(3.0) +def test_one_arm_resetting_does_not_hide_another_arms_fault(world): + """A resetting arm and a faulted one both serialize to nothing. The fault wins whichever channel the + embodiment happens to list first, so a bimanual rig cannot carry on driving a faulted arm.""" + left, right = f'{keys.ROBOT_STATE}.left', f'{keys.ROBOT_STATE}.right' + embodiment = Embodiment( + '', + { + left: Observation(pimm.NoOpEmitter(), Serializers.robot_state), + right: Observation(pimm.NoOpEmitter(), Serializers.robot_state), + keys.GRIP: Observation(pimm.NoOpEmitter(), None), + }, + {keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command)}, + {}, + pimm.NoOpEmitter(), + ) + policy = SpyPolicy() + harness = Harness(policy, embodiment) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + left_em = world.pair(harness.observations[left]) + right_em = world.pair(harness.observations[right]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + def emit_states(): + left_em.emit(make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.RESETTING)) + right_em.emit(make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR)) + grip_em.emit(0.0) + + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='test')), 0.0), + (emit_states, 0.01), + (None, 0.02), + ]) + + drive_scheduler(world.start([harness, driver]), steps=40) + + assert policy.last_obs is not None, 'the resetting arm swallowed the other arm’s fault' + assert policy.last_obs[keys.ROBOT_FAULT] is True + + +def test_the_trial_context_cannot_stand_in_for_what_the_harness_read(world): + """A context carries whatever the trial plan puts in it. What the harness read and stamped this round + outranks it, so a config key can neither report the arm sound nor supply its state.""" + policy = SpyPolicy() + harness = Harness(policy, make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + faulted = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR) + context = {'task': 'test', keys.ROBOT_FAULT: False, keys.GRIP: 'from the config'} + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(**context)), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, faulted), 0.01), + (None, 0.02), + ]) + + drive_scheduler(world.start([harness, driver]), steps=40) + + assert policy.last_obs is not None + assert policy.last_obs[keys.ROBOT_FAULT] is True, 'a context key reported the faulted arm sound' + assert policy.last_obs[keys.GRIP] != 'from the config', 'a context key stood in for a channel' + + +@pytest.mark.timeout(20.0) +def test_a_stop_clears_the_chunk_in_the_round_the_fault_is_seen(world): + """A stop has no waypoints to place: an arm that faults mid-chunk stops in the round its fault is seen.""" + fault_at, period = 0.5, 0.005 + stack = StopOnFault() | ChunkedSchedule() + harness = Harness(stack.wrap(SlowPolicy(span_sec=1.0, steps=50)), make_embodiment(simulated=True)) + grip_recorder = _TimedRecorder(world.clock) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + pose, joints = [0.1, 0.2, 0.3], [0.4, 0.5, 0.6] + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, make_robot_state(pose, joints)), fault_at), + (partial(robot_em.emit, make_robot_state(pose, joints, status=RobotStatus.ERROR)), 0.5), + ]) + drive_scheduler(world.start([harness, driver, _Pacer(period)]), steps=4000) + + played = [t for t, _ in grip_recorder.emitted[1:]] # drop the startup home + assert [t for t in played if t < fault_at], 'the chunk was not playing when the arm faulted' + late = [t for t in played if t > fault_at + 3 * period] + assert not late, f'the faulted arm was still being driven at {late}' + + +@pytest.mark.timeout(20.0) +def test_finish_does_not_wait_for_the_call_in_flight(): + """FINISH ends the episode where it lands: the recording stops and the arm homes while the model is + still inside its call, and the failure that call ends in reaches nobody. + + Real time, real rig: a wall-charged call is the only one the harness leaves in flight across rounds. + """ + hang_sec = 1.0 + + class _HangingSession(Session): + def __call__(self, obs): + time.sleep(hang_sec) + raise RuntimeError('inference boom') + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + cmd_recorder = RecordingEmitter() + ds_recorder = _TimedRecorder(world.clock) + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(ds_recorder) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + started = world.clock.now() + drive_scheduler(world.start([harness, driver]), steps=40) + + stops = [(t, data) for t, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] + assert len(stops) == 1 + assert stops[0][0] - started < hang_sec, 'the stop waited for the call to answer' + assert isinstance(cmd_recorder.emitted[-1][1], Reset) + + +@pytest.mark.timeout(20.0) +def test_the_run_ends_only_once_the_call_it_abandoned_is_out_of_the_policy(): + """A sweep runs a harness per eval over one shared policy, so a call still inside the model when a run + ends would meet the next run's ``new_session`` — or ``policy.close()``. The run outlives its own call.""" + hang_sec = 1.0 + left_the_model = threading.Event() + + class _HangingSession(Session): + def __call__(self, obs): + time.sleep(hang_sec) + left_the_model.set() + return None + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver]), steps=40) + + assert left_the_model.is_set(), 'the run returned with a call still inside the shared policy' + + +@pytest.mark.timeout(20.0) +def test_the_session_is_closed_only_once_its_call_has_left_it(): + """``RemoteSession.close`` shuts the websocket the call in flight is talking over, and ``Session`` asks + for no thread safety, so the session is retired with its worker and closed once that worker is joined.""" + inside_at_close = [] + + class _HangingSession(Session): + def __init__(self): + self.inside = False + + def __call__(self, obs): + self.inside = True + try: + time.sleep(1.0) + return None + finally: + self.inside = False + + def close(self): + inside_at_close.append(self.inside) + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver]), steps=40) + + assert inside_at_close == [False], 'the session was closed while its own call was still inside it' + + +@pytest.mark.timeout(3.0) +def test_a_rescheduled_trajectory_clears_the_channels_it_omits(world): + """A trajectory naming only one channel replaces the whole schedule: the omitted channel stops being + played rather than draining the previous trajectory's tail.""" + + class _GripThenArm(Session): + """First a two-channel chunk, then an arm-only one that must silence the gripper.""" + + def __init__(self): + self._calls = 0 + + def __call__(self, obs): + self._calls += 1 + pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) + command = CartesianPosition(pose=pose) + if self._calls == 1: + return [ + {keys.ROBOT_COMMAND: command, keys.TARGET_GRIP: 0.5, keys.ACTION_TIMESTAMP: i * 0.01} + for i in range(10) + ] + return [{keys.ROBOT_COMMAND: command, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10)] + + class _GripThenArmPolicy(Policy): + def new_session(self, context=None, now=None): + return _GripThenArm() + + harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) + grip_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, 0.5), + ]) + drive_scheduler(world.start([harness, driver]), steps=1000) + + grips = _emitted_grips(grip_recorder) + assert grips[0] == 0.0 # the startup home + assert set(grips[1:]) == {0.5}, f'the second chunk kept the gripper playing: {grips}' + + +@pytest.mark.timeout(3.0) +def test_home_and_manual_commands_are_emitted_as_plain_values(world): + """Homing and operator commands bypass the schedule: they are the command, not a plan to play.""" + harness = Harness(StubPolicy(), make_embodiment()) + cmd_recorder = RecordingEmitter() + grip_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + manual_em = world.pair(harness.manual_command) + + pose = Transform3D(translation=np.array([0.1, 0.1, 0.1], dtype=np.float32), rotation=Rotation.identity) + manual = CartesianPosition(pose=pose) + driver = ManualDriver([(partial(manual_em.emit, {keys.ROBOT_COMMAND: manual}), 0.01), (None, 0.02)]) + drive_scheduler(world.start([harness, driver]), steps=50) + + assert _emitted_commands(cmd_recorder) == [Reset(), manual] + assert _emitted_grips(grip_recorder) == [0.0] + + +@pytest.mark.timeout(20.0) +def test_abort_discards_a_call_that_is_still_in_flight(world): + """An ABORT while the model is still inside its call throws that answer away: the trajectory it carries + never reaches the devices.""" + harness = Harness(ChunkedSchedule().wrap(SlowPolicy(wall_sec=1.0)), make_embodiment(simulated=True)) + cmd_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: True})), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, 0.05), # well inside the 1.0s the call takes + (partial(directive_em.emit, Directive.ABORT()), 0.0), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver, _Pacer()]), steps=2000) + + assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 6a19ab5fd..6b30b4637 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -27,7 +27,7 @@ RestrictImageSize, ) from positronic.policy.observation import ObservationCodec -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack class _FakeClock: @@ -63,28 +63,62 @@ def new_session(self, context=None, now=None): return self._session -def _obs(now_sec=0.0): - return {keys.OBS_TIME_NS: int(now_sec * 1e9)} +def _obs(now_sec=0.0, fault=False): + return {keys.OBS_TIME_NS: int(now_sec * 1e9), keys.ROBOT_FAULT: fault} + + +class TestStopOnFault: + def test_a_faulted_arm_stops_what_is_executing(self): + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session(_obs(0.0, fault=True)) == [] + assert inner.call_count == 0, 'the model was asked about an arm that is not tracking it' + + def test_a_sound_arm_reaches_the_model(self): + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session(_obs(0.0)) is not None + assert inner.call_count == 1 + + def test_an_observation_with_no_flag_is_not_a_fault(self): + """Only the harness stamps the flag; a probe replaying a recording has no arm to fault.""" + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session({keys.OBS_TIME_NS: 0}) is not None + assert inner.call_count == 1 + + def test_recovery_plans_afresh_instead_of_resuming(self): + """The fault resets the scheduler below it, so the first sound observation infers again rather than + waiting out the chunk stamped before the fault.""" + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 1.0}]) + session = (StopOnFault() | ChunkedSchedule()).wrap(inner).new_session(now=_FakeClock(t=0.0).now) + + assert session(_obs(0.0)) is not None # a chunk that runs until 1.0 + assert session(_obs(0.2, fault=True)) == [] + assert session(_obs(0.3)) is not None class TestChunkedSchedule: def test_first_call_runs_inference(self): # Relative timestamps: trajectory of duration 0.5s clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) policy = ChunkedSchedule().wrap(inner) session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None assert len(result) == 2 # Timestamps stamped to absolute by ChunkedSchedule. - assert result[0]['timestamp'] == 1.0 - assert result[1]['timestamp'] == 1.5 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 + assert result[1][keys.ACTION_TIMESTAMP] == 1.5 def test_returns_none_while_trajectory_active(self): # Trajectory starts at clock=1.0, ends at 1.0+0.5=1.5. clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) policy = ChunkedSchedule().wrap(inner) session = policy.new_session(now=clock.now) session(_obs()) @@ -95,26 +129,33 @@ def test_returns_none_while_trajectory_active(self): def test_re_infers_after_trajectory_consumed(self): clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) session = ChunkedSchedule().wrap(inner).new_session(now=clock.now) - session(_obs()) # trajectory ends at clock=1.5 - clock.t = 1.3 - assert session(_obs()) is None + session(_obs(1.0)) # trajectory ends at clock=1.5 + assert session(_obs(1.3)) is None clock.t = 1.6 - result = session(_obs()) + result = session(_obs(1.6)) assert result is not None assert inner._session.call_count == 2 def test_single_action_refires_immediately_after(self): """Single action at ts=0 → trajectory_end = now → next tick re-infers.""" clock = _FakeClock(t=1.0) - policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) - session(_obs()) + session(_obs(1.0)) clock.t = 1.01 - result = session(_obs()) + result = session(_obs(1.01)) assert result is not None + def test_expiry_is_judged_at_the_observation_instant(self): + """Whether the trajectory has run out is a question about the observation, not about ``now``.""" + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) + session = ChunkedSchedule().wrap(inner).new_session(now=_FakeClock(t=2.0).now) + session(_obs(1.0)) # anchored at now()=2.0, so the trajectory ends at 2.5 + assert session(_obs(2.4)) is None + assert session(_obs(2.6)) is not None + class TestPipelineComposition: """Test | operator across PolicyWrapper and Codec types.""" @@ -123,7 +164,7 @@ def test_wrapper_pipe_wrapper(self): clock = _FakeClock(t=1.0) pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = pipeline.wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None @@ -134,7 +175,7 @@ def test_codec_pipe_wrapper(self): codec = ActionTimestamp(fps=10.0) pipeline = codec | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy([{'action': 'test', 'timestamp': 0.0}])) + policy = pipeline.wrap(_ConstPolicy([{'action': 'test', keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None @@ -150,7 +191,7 @@ def test_full_pipeline(self): session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None - assert result[0]['timestamp'] == 1.0 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 # Second call within trajectory window returns None (ChunkedSchedule). clock.t = 1.2 assert session(_obs()) is None @@ -368,9 +409,12 @@ def test_non_deliverable_wrapper_fails_loudly(self): with pytest.raises(NotImplementedError, match='not deliverable'): IKJointsAction(solver_cls=None).to_spec() - def test_wire_names_match_table(self): + def test_the_table_publishes_these_exact_wire_names(self): + """The strings a deployed server already declares its local stack with. Spelled out here rather than + read off ``WIRE_NAME``, so renaming an attribute cannot quietly rename the wire.""" instances = { 'chunked_schedule': ChunkedSchedule(), + 'stop_on_fault': StopOnFault(), 'temporal_stack': TemporalStack(('v',), (0.0,)), 'action_timestamp': ActionTimestamp(fps=10.0), 'action_horizon': ActionHorizon(1.0), @@ -462,17 +506,17 @@ def test_inline_full_pipe(self): session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None - assert result[0]['timestamp'] == 1.0 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 clock.t = 1.2 assert session(_obs()) is None def test_inline_tolerates_marker_less_pipe(self): clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) session = policy.new_session(now=clock.now) result = session(_obs()) - assert result is not None and result[0]['timestamp'] == 1.0 + assert result is not None and result[0][keys.ACTION_TIMESTAMP] == 1.0 def test_inline_bare_source_pipe_is_the_loaded_policy(self): inner = _ConstPolicy([]) diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index 7b0e1d7d2..abbbabf07 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -1,4 +1,4 @@ -"""Composable policy wrappers — scheduling and temporal frame stacking. +"""Composable policy wrappers — scheduling, fault handling and temporal frame stacking. Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only @@ -20,6 +20,34 @@ def _obs_time(obs) -> float: return obs[keys.OBS_TIME_NS] / 1e9 +class StopOnFault(PolicyWrapper): + """Stop the arm while it is faulted, and plan afresh once it recovers. + + A faulted arm is not tracking the plan it was given, so the plan is worthless: this answers the empty + trajectory — stop what is executing — and resets the sessions below, so the first sound observation after + recovery reaches the model instead of resuming a chunk stamped before the fault. It belongs outside the + scheduling wrapper, which would otherwise answer "keep playing" without ever seeing the fault. + + The fault is ``keys.ROBOT_FAULT``, which the harness stamps on every observation it builds. An + observation from anywhere else — a probe replaying a recording — carries no arm to fault. + """ + + WIRE_NAME = 'stop_on_fault' + + class _Session(DelegatingSession): + def __call__(self, obs): + if not obs.get(keys.ROBOT_FAULT, False): + return self._inner(obs) + self.cancel() + return [] + + def wrap_session(self, inner: Session, context, now: Now | None): + return StopOnFault._Session(inner) + + def to_spec(self): + return {'name': self.WIRE_NAME} + + class ChunkedSchedule(PolicyWrapper): """Wait for the current trajectory to finish before calling the inner policy again. @@ -29,6 +57,8 @@ class ChunkedSchedule(PolicyWrapper): until the last action's timestamp is reached, then calls the inner policy. """ + WIRE_NAME = 'chunked_schedule' + class _Session(DelegatingSession): """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" @@ -44,7 +74,7 @@ def __call__(self, obs): 'new_session. The harness supplies it; a direct RemotePolicy.new_session() outside the harness ' 'must too.' ) - if self._trajectory_end is not None and self._now() < self._trajectory_end: + if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: return None result = self._inner(obs) if result is not None: @@ -68,7 +98,7 @@ def wrap_session(self, inner: Session, context, now: Now | None): return ChunkedSchedule._Session(inner, now) def to_spec(self): - return {'name': 'chunked_schedule'} + return {'name': self.WIRE_NAME} class _StackBuffer: @@ -134,6 +164,8 @@ class TemporalStack(PolicyWrapper): chunk-0 empty prefix) never engage it on a padded full-length stack. """ + WIRE_NAME = 'temporal_stack' + class _Session(DelegatingSession): def __init__(self, inner: Session, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_start: bool): super().__init__(inner) @@ -158,11 +190,11 @@ def __init__(self, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_st 'in-range targets and the stack would be empty' ) - def wrap_session(self, inner: Session, context, now: Now): + def wrap_session(self, inner: Session, context, now: Now | None): return TemporalStack._Session(inner, self._keys, self._offsets_sec, self._pad_start) def to_spec(self): return { - 'name': 'temporal_stack', + 'name': self.WIRE_NAME, 'args': {'keys': list(self._keys), 'offsets_sec': list(self._offsets_sec), 'pad_start': self._pad_start}, } diff --git a/positronic/replay_record.py b/positronic/replay_record.py index 1525e561b..04014a3c0 100644 --- a/positronic/replay_record.py +++ b/positronic/replay_record.py @@ -25,13 +25,6 @@ from positronic.utils.logging import init_logging -class _TrajectoryEmitter(pimm.ControlSystemEmitter): - """Sends each replayed value as the one-waypoint trajectory the command channels carry, at its playback time.""" - - def emit(self, data: Any, ts: int = -1) -> None: - super().emit([(ts, data)], ts) - - class Replay(DsPlayerAgent): """Adapts `DsPlayerAgent` to be used as a policy control system.""" @@ -42,12 +35,12 @@ def __init__(self, poll_hz: float = 100.0): self.gripper_state = pimm.FakeReceiver(self) self.robot_meta_in = pimm.FakeReceiver(self) self.frames = pimm.ReceiverDict(self, fake=True) - self.outputs['robot_commands'] = _TrajectoryEmitter(self) - self.outputs[keys.TARGET_GRIP] = _TrajectoryEmitter(self) + self.outputs[keys.ROBOT_COMMAND] = pimm.ControlSystemEmitter(self) + self.outputs[keys.TARGET_GRIP] = pimm.ControlSystemEmitter(self) @property def robot_commands(self) -> pimm.ControlSystemEmitter: - return self.outputs['robot_commands'] + return self.outputs[keys.ROBOT_COMMAND] @property def target_grip(self) -> pimm.ControlSystemEmitter: @@ -55,8 +48,13 @@ def target_grip(self) -> pimm.ControlSystemEmitter: class RestoreCommand(Derive): + """Derives the column ``Replay`` plays the arm's commands from: a player matches a column to a port by + name, so this must carry the name that port is registered under.""" + def __init__(self): - super().__init__(robot_commands=self._commands_from_episode) + # ``Derive`` takes its columns as ``**kwargs`` beside a keyword-only ``meta``, so a checker reading an + # unpacked dict cannot rule out ``meta`` and matches the transform against its type. + super().__init__(**{keys.ROBOT_COMMAND: self._commands_from_episode}) # pyright: ignore[reportArgumentType] @staticmethod def _commands_from_episode(episode: Episode) -> Any: diff --git a/positronic/robot_controller.py b/positronic/robot_controller.py index 07fbc67f2..4df95f6f8 100644 --- a/positronic/robot_controller.py +++ b/positronic/robot_controller.py @@ -1,4 +1,5 @@ import configuronic as cfn +import numpy as np import pimm import positronic.cfg.hardware.roboarm @@ -18,15 +19,15 @@ def main(robot): match text_command.split(' '): case ['reset']: - command_emmiter.emit([(world.clock.now_ns(), command.Reset())]) + command_emmiter.emit(command.Reset()) case ['move', x, y, z, qw, qx, qy, qz]: pos = [float(x) for x in [x, y, z]] quat = geom.Rotation.from_quat([float(qw), float(qx), float(qy), float(qz)]) move = command.CartesianPosition(geom.Transform3D(translation=pos, rotation=quat)) - command_emmiter.emit([(world.clock.now_ns(), move)]) + command_emmiter.emit(move) case ['joint_move', *args]: args = [float(x) for x in args] - command_emmiter.emit([(world.clock.now_ns(), command.JointPosition(positions=args))]) + command_emmiter.emit(command.JointPosition(positions=np.asarray(args))) case ['info']: print('Q', state_receiver.value.q) print('DQ', state_receiver.value.dq) diff --git a/positronic/simulator/env_server/adapter.py b/positronic/simulator/env_server/adapter.py index eed79fee5..14d096e73 100644 --- a/positronic/simulator/env_server/adapter.py +++ b/positronic/simulator/env_server/adapter.py @@ -1,14 +1,13 @@ """The ``EnvAdapter`` interface: the per-benchmark canonical<->raw mappings, on the client side. ``RemoteEnvControlSystem`` is a dumb translator — it moves data between pimm signals and this adapter. -The adapter is the smart half: it turns the Harness's command trajectories into the env's raw action -(owning trajectory playing and how to hold between waypoints), maps raw observations back to canonical -signals — policy-facing and privileged ground-truth kept separate — and reads the terminal. Each -benchmark ships one adapter (``vendors/``-style); the native ``MujocoSim`` fixture is the reference. +The adapter is the smart half: it turns the Harness's commands into the env's raw action (owning what to +hold between them), maps raw observations back to canonical signals — policy-facing and privileged +ground-truth kept separate — and reads the terminal. Each benchmark ships one adapter (``vendors/``-style); +the native ``MujocoSim`` fixture is the reference. """ from abc import ABC, abstractmethod -from collections import defaultdict from typing import Any, final import numpy as np @@ -18,14 +17,6 @@ from positronic.drivers.roboarm import command as roboarm_command -def fresh_command_players() -> defaultdict[str, roboarm_command.TrajectoryPlayer]: - """A trajectory player per command channel: ``robot_command`` accumulates the deltas due in one tick (a - missed tick catches up instead of dropping motion), every other channel keeps the last value due.""" - players = defaultdict(roboarm_command.TrajectoryPlayer) - players[keys.ROBOT_COMMAND] = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - return players - - class EnvAdapter(ABC): """The mappings between the canonical embodiment contract and an env's raw wire payloads.""" @@ -38,11 +29,11 @@ def reset_token(self, context: dict[str, Any]) -> Any: """ @abstractmethod - def action(self, commands: dict[str, pimm.Message], now_ns: int) -> dict[str, Any]: - """The latest per-channel command messages + the clock -> the raw action the env steps. + def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: + """The latest per-channel command messages -> the raw action the env steps. - The adapter owns trajectory playing (sampling each channel's waypoints down to ``now_ns``) and - what to do between waypoints — e.g. hold the last commanded value, the absolute-mode invariant. + A channel delivers a command only when one comes due, so the adapter owns what happens in between — + e.g. hold the last commanded value, the absolute-mode invariant. """ @abstractmethod @@ -102,10 +93,9 @@ def _wire_command(cmd: Any) -> dict[str, Any]: class WireCommandAdapter(EnvAdapter): """An adapter whose action is the shared wire payload ``{'command': <tagged dict>, 'grip': float}``. - The command side of every remote benchmark adapter: it plays each command channel's trajectory down to - the clock — holding an absolute setpoint between waypoints, firing a relative delta once — and flattens - the held arm command (a pose as ``[t(3), R(9)]``, joint positions, or per-step joint deltas) plus the - gripper closure into one payload. + The command side of every remote benchmark adapter: it holds an absolute setpoint until the next command + arrives and fires a relative delta once, and flattens the held arm command (a pose as ``[t(3), R(9)]``, + joint positions, or per-step joint deltas) plus the gripper closure into one payload. All action *encoding* — how the tagged command becomes the env's native action — stays server-side with the env's own model. Subclasses implement ``_reset_token`` (the base clears the per-trial command state around it) and keep the observation and terminal mappings to themselves. @@ -118,11 +108,7 @@ def __init__(self, env_control_frame: geom.Transform3D | None = None): self._reset_command_state() def _reset_command_state(self) -> None: - self._players = fresh_command_players() - self._held: dict[str, Any] = {} # last sampled waypoint per channel — re-sent until it changes - # Last commanded gripper closure, held across a cancelled grip trajectory: grip is an absolute [0, 1] - # value with no 'hold' command to fall back on (unlike the arm), so cancelling must freeze it, not reopen. - self._grip = 0.0 + self._held: dict[str, Any] = {} # last command per channel — re-sent until the next one arrives @final def reset_token(self, context: dict[str, Any]) -> Any: @@ -133,21 +119,14 @@ def reset_token(self, context: dict[str, Any]) -> Any: def _reset_token(self, context: dict[str, Any]) -> Any: """The per-trial RUN context -> the env's opaque reset token; the command state is already cleared.""" - def action(self, commands: dict[str, pimm.Message], now_ns: int) -> dict[str, Any]: + def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: for name, msg in commands.items(): - player = self._players[name] if msg.updated: - player.set(msg.data) - if not msg.data: # an empty trajectory cancels: stop replaying the held waypoint - self._held.pop(name, None) - value = player.advance(now_ns) - if value is not None: - self._held[name] = value + self._held[name] = msg.data # The server maps the held command into its controller's action. Reset has no env-side action, so it # forwards as a hold; a delta — Cartesian or joint — is a one-shot relative motion, forwarded once then # dropped. Re-sending a stale delta would re-compose it against the moving arm every tick (the eef - # drifts, or the joints walk toward their limits), so once a delta's trajectory is exhausted the arm - # holds its measured pose. + # drifts, or the joints walk toward their limits), so after one step the arm holds its measured pose. cmd = self._held.get(keys.ROBOT_COMMAND) match cmd: case roboarm_command.Reset(): @@ -155,6 +134,5 @@ def action(self, commands: dict[str, pimm.Message], now_ns: int) -> dict[str, An cmd = None case roboarm_command.CartesianDelta() | roboarm_command.JointDelta(): self._held.pop(keys.ROBOT_COMMAND) - if keys.TARGET_GRIP in self._held: - self._grip = float(self._held[keys.TARGET_GRIP]) - return {'command': _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), 'grip': self._grip} + grip = float(self._held.get(keys.TARGET_GRIP, 0.0)) + return {'command': _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), 'grip': grip} diff --git a/positronic/simulator/env_server/client.py b/positronic/simulator/env_server/client.py index 4f6ca6f96..ee8dd2438 100644 --- a/positronic/simulator/env_server/client.py +++ b/positronic/simulator/env_server/client.py @@ -4,6 +4,7 @@ control system; tests use it directly to compare a socket rollout against an in-process one. """ +import logging import time from typing import Any @@ -12,6 +13,13 @@ from .protocol import decode, encode +logger = logging.getLogger(__name__) + +# How long ``close`` waits to be acknowledged. Teardown often runs while the peer is on its way out, and a +# simulator wedged in its own destructor holds the socket open without ever answering — an unbounded wait there +# hangs the run in place of ending it. +_CLOSE_ACK_TIMEOUT = 5.0 + class EnvConnection: """One websocket to an ``EnvServer``, opened with retry. ``reset``/``step`` block on the round-trip. @@ -52,12 +60,14 @@ def _request(self, msg: dict[str, Any]) -> dict[str, Any]: return result def close(self) -> None: - # Best-effort: ask the server to release, but a peer that is already gone (a crashed or killed server) is - # success too — the socket is closed regardless. try: self._ws.send(encode({'cmd': 'close'})) - self._ws.recv() + self._ws.recv(timeout=_CLOSE_ACK_TIMEOUT) except ConnectionClosed: - pass + pass # a peer already gone has released whatever the acknowledgement would have reported + except TimeoutError: + # Abandoning it is still better than hanging the run here, but a server that took the request and + # never answered is wedged rather than finished, and its resources are nobody's to reclaim now. + logger.error('Env server did not acknowledge close within %.1fs; abandoning it', _CLOSE_ACK_TIMEOUT) finally: self._ws.close() diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index e8a39aae7..4896cf99c 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -5,7 +5,7 @@ to, but no command logic. Each control period it hands the latest command messages to the ``EnvAdapter``, round-trips the raw action it returns over the wire, and re-emits the canonical signals the adapter maps back — so only raw arrays cross the boundary and the World's virtual clock advances by the env's -``control_dt`` per step. The adapter owns trajectory playing, holding, and the canonical<->raw mappings; +``control_dt`` per step. The adapter owns holding and the canonical<->raw mappings; ``control_dt`` is whatever the latest observation reports (``reset`` and every ``step``). """ @@ -36,7 +36,7 @@ def __init__(self, adapter: EnvAdapter, serve: AbstractContextManager[tuple[str, self._cleanup = ExitStack() self._conn: EnvConnection | None = None - self.commands: pimm.ReceiverDict = pimm.ReceiverDict(self, default=[]) + self.commands: pimm.ReceiverDict = pimm.ReceiverDict(self) self.observations: pimm.EmitterDict = pimm.EmitterDict(self) self.privileged: pimm.EmitterDict = pimm.EmitterDict(self) self.robot_meta: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) @@ -119,7 +119,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p # observation assembly (shared-memory image allocation + camera copies) inside it, so the # reduce can split materialisation out of the wire cost. with telemetry.span(telemetry_keys.SPAN_ENV_STEP): - self._frame = self._step_env(clock) + self._frame = self._step_env() with telemetry.span(telemetry_keys.SPAN_MATERIALIZE): self._emit_payload(self._frame['obs']) finally: @@ -127,9 +127,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p # ever connected. self._cleanup.close() - def _step_env(self, clock: pimm.Clock) -> dict[str, Any]: - commands = {name: receiver.read() for name, receiver in self.commands.items()} - result = self._conn.step(self._adapter.action(commands, clock.now_ns())) + def _step_env(self) -> dict[str, Any]: + reads = ((name, receiver.read()) for name, receiver in self.commands.items()) + commands = {name: msg for name, msg in reads if msg is not None} + result = self._conn.step(self._adapter.action(commands)) payload = self._adapter.terminal(result) if payload: # truthy-valued done: a non-empty payload ends the trial, an empty/``None`` one continues self.done.emit(payload) @@ -152,7 +153,7 @@ def remote_franka_embodiment( server cannot import positronic to emit it via ``robot_meta``). """ observations = { - 'robot_state': Observation(proxy.observations['robot_state'], Serializers.robot_state), + keys.ROBOT_STATE: Observation(proxy.observations[keys.ROBOT_STATE], Serializers.robot_state), keys.GRIP: Observation(proxy.observations[keys.GRIP], None), **{logical: Observation(proxy.observations[logical], Serializers.camera_images) for logical in camera_dict}, } diff --git a/positronic/simulator/env_server/tests/mujoco_env.py b/positronic/simulator/env_server/tests/mujoco_env.py index fdff9ec66..1a6e29e03 100644 --- a/positronic/simulator/env_server/tests/mujoco_env.py +++ b/positronic/simulator/env_server/tests/mujoco_env.py @@ -133,10 +133,9 @@ def step(self, action: dict[str, Any]) -> dict[str, Any]: cmd = roboarm_command.CartesianDelta(geom.Transform3D.from_vector(command['delta'], _ROTMAT)) case other: raise ValueError(f'MujocoEnv got unsupported command type {other!r}') - now_ns = self._clock.now_ns() if cmd is not None: - self._cmd_emit.emit([(now_ns, cmd)]) - self._grip_emit.emit([(now_ns, float(action['grip']))]) + self._cmd_emit.emit(cmd) + self._grip_emit.emit(float(action['grip'])) self._advance(self._timestep) return {'obs': self._read_obs(), 'done': False, 'control_dt': self._timestep} @@ -172,7 +171,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = geom.Transform3D(raw_obs['ee_pos'], geom.Rotation.from_quat(raw_obs['ee_quat'])) state.encode(raw_obs['q'], raw_obs['dq'], ee_pose) state.array[14 + 7] = float(raw_obs['status']) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} for logical, model_name in self._camera_dict.items(): frame = raw_obs['cameras'][model_name] adapter = pimm.shared_memory.NumpySMAdapter(shape=frame.shape, dtype=frame.dtype) diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 76bc0818f..cb37caff2 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -1,9 +1,12 @@ -from contextlib import nullcontext +import threading +import time +from contextlib import contextmanager, nullcontext from dataclasses import replace import numpy as np import pos3 import pytest +from websockets.sync.server import serve as websocket_serve import pimm from positronic import geom, keys @@ -20,7 +23,8 @@ from positronic.policy.tests.test_harness import StubPolicy from positronic.policy.wrappers import ChunkedSchedule from positronic.simulator.env_server.adapter import EnvAdapter, _in_env_control_frame, _wire_command -from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.env_server.client import _CLOSE_ACK_TIMEOUT, EnvConnection +from positronic.simulator.env_server.launcher import free_port from positronic.simulator.env_server.proxy import RemoteEnvControlSystem from positronic.simulator.env_server.server import EnvProtocol from positronic.simulator.env_server.tests.conftest import serve_env @@ -95,6 +99,41 @@ def test_transport_is_transparent(env_server): assert direct_step['control_dt'] == socket_step['control_dt'] +@contextmanager +def _mute_server(): + """A peer that accepts the connection and then answers nothing, holding the socket open. + + What a simulator looks like once it is wedged in its own teardown: the process is past serving but the + socket outlives it, so nothing ever closes the connection from that end. + """ + host, port = 'localhost', free_port() + release = threading.Event() + + def handler(connection): + for _ in connection: + release.wait(timeout=60.0) + + server = websocket_serve(handler, host, port) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield host, port + finally: + release.set() + server.shutdown() + thread.join(timeout=5.0) + + +@pytest.mark.timeout(60.0) +def test_close_gives_up_on_a_peer_that_never_answers(): + """Teardown ends the run: an unanswered goodbye is as good as a closed socket, and is not waited out.""" + with _mute_server() as (host, port): + conn = EnvConnection(host, port) + started = time.monotonic() + conn.close() + assert time.monotonic() - started < _CLOSE_ACK_TIMEOUT + 10.0 + + _HOLD = {'command': {'type': 'hold'}, 'grip': 0.0} @@ -144,7 +183,7 @@ def test_robolab_reports_and_drives_the_same_frame(self): 'joint_vel': np.zeros(7), 'grip': 0.0, } - reported = adapter.observations(raw)['robot_state'].ee_pose + reported = adapter.observations(raw)[keys.ROBOT_STATE].ee_pose commanded = _in_env_control_frame(roboarm_command.CartesianPosition(reported), adapter.env_control_frame).pose np.testing.assert_allclose(commanded.as_vector(self.rotmat), eef.as_vector(self.rotmat), atol=1e-6) @@ -211,7 +250,7 @@ class _CountdownAdapter(EnvAdapter): def reset_token(self, context): return context.get('eval.seed') - def action(self, commands, now_ns): + def action(self, commands): return {} def observations(self, raw_obs): diff --git a/positronic/simulator/libero/adapter.py b/positronic/simulator/libero/adapter.py index acb88ceac..02d2681bb 100644 --- a/positronic/simulator/libero/adapter.py +++ b/positronic/simulator/libero/adapter.py @@ -42,7 +42,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = geom.Transform3D(raw_obs['eef_pos'], geom.Rotation.from_quat_xyzw(raw_obs['eef_quat'])) state = MujocoFrankaState() state.encode(raw_obs['joint_pos'], raw_obs['joint_vel'], ee_pose) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} for logical, env_key in self._camera_dict.items(): # robosuite renders bottom-up; flip to standard top-down orientation (LIBERO's own video path # flips the same way). diff --git a/positronic/simulator/mujoco/sim.py b/positronic/simulator/mujoco/sim.py index 8ab1852a6..6d41d3077 100644 --- a/positronic/simulator/mujoco/sim.py +++ b/positronic/simulator/mujoco/sim.py @@ -99,9 +99,9 @@ class MujocoSim(pimm.ControlSystem): ``reset`` rebuilds the scene and flags frame-0 publication; the run loop publishes that post-reset scene on its next turn — in sequence, before any step — so the first inference reads it and the - recorder logs it. Every other turn applies the due command waypoints, steps once, and emits the due - streams (post-step, Gym-style). The sim sleeps one control period each turn, so it is the eval's sole - time-master. Each stream has an independent rate (``*_fps``, ``None`` = every physics tick). + recorder logs it. Every other turn applies whatever command has just arrived, steps once, and emits + the due streams (post-step, Gym-style). The sim sleeps one control period each turn, so it is the + eval's sole time-master. Each stream has an independent rate (``*_fps``, ``None`` = every physics tick). """ def __init__( @@ -140,20 +140,14 @@ def __init__( self._home() self._error = False self._adapters: dict[str, pimm.shared_memory.NumpySMAdapter] | None = None - self._arm_player = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - self._grip_player = roboarm_command.TrajectoryPlayer() self._last_grip = 0.0 # Set by ``reset``; the run loop publishes frame-0 (instead of stepping) on its next turn and clears it. self._reset_pending = False - self.commands: pimm.SignalReceiver[roboarm_command.Trajectory[roboarm_command.CommandType]] = ( - pimm.ControlSystemReceiver(self, default=[]) - ) + self.commands = pimm.ControlSystemReceiver[roboarm_command.CommandType](self) self.state: pimm.SignalEmitter[MujocoFrankaState] = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.SignalReceiver[roboarm_command.Trajectory[float]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.grip: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) self.cameras: pimm.EmitterDict = pimm.EmitterDict(self) # Privileged ground truth: the full ``save_state`` dict, spec keys prefixed with '.' so the @@ -162,12 +156,14 @@ def __init__( # replays these states through it (``mj_setState`` + ``mj_forward``). self.sim_state: pimm.SignalEmitter[dict[str, np.ndarray]] = pimm.ControlSystemEmitter(self) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: self._emit_robot_meta() - state_due = _Cadence(self._state_fps) - grip_due = _Cadence(self._grip_fps) - sim_state_due = _Cadence(self._sim_state_fps) - cameras_due = _Cadence(self._camera_fps) + streams = [ + (_Cadence(self._state_fps), self._emit_state), + (_Cadence(self._grip_fps), self._emit_grip), + (_Cadence(self._sim_state_fps), self._emit_sim_state), + (_Cadence(self._camera_fps), self._emit_cameras), + ] while not should_stop.value: yield pimm.Sleep(self.model.opt.timestep) @@ -183,23 +179,13 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p continue now = clock.now() cmd_msg = self.commands.read() - if cmd_msg.updated: - self._arm_player.set(cmd_msg.data) if self._error: self._error = False - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale - # waypoint once the error clears. - self._arm_player.set([]) - else: - cmd = self._arm_player.advance(clock.now_ns()) - if cmd is not None: - self._apply_command(cmd) + elif cmd_msg is not None and cmd_msg.updated: + self._apply_command(cmd_msg.data) grip_msg = self.target_grip.read() - if grip_msg.updated: - self._grip_player.set(grip_msg.data) - grip = self._grip_player.advance(clock.now_ns()) - if grip is not None: - self._last_grip = grip + if grip_msg is not None and grip_msg.updated: + self._last_grip = grip_msg.data self._apply_grip(self._last_grip) # An env step is the sim advance plus the observations it produces, rendering included @@ -208,14 +194,9 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p with telemetry.span(telemetry_keys.SPAN_ENV_STEP): self.step() self.fps_counter.tick() - if state_due(now): - self._emit_state() - if grip_due(now): - self._emit_grip() - if sim_state_due(now): - self._emit_sim_state() - if cameras_due(now): - self._emit_cameras() + for due, emit in streams: + if due(now): + emit() if self._renderer is not None: self._renderer.close() @@ -228,17 +209,14 @@ def reset(self, seed: int | None = None): colors, cameras) re-randomize too; the renderer and IK physics rebind lazily. The run loop publishes the prepared scene as frame-0 on its next turn — in sequence, before any step — so the first inference reads the reset state and the recorder logs it. Stale commands queued while idle - (e.g. the inter-episode home) are dropped and the run-loop's trajectory players and held grip are - cleared, so the first step neither applies a queued command nor replays the previous episode's - trajectory on the freshly reset scene. + (e.g. the inter-episode home) are dropped and the held grip is cleared, so the first step does not + apply a queued command on the freshly reset scene. """ self._load_scene(seed) self._home() self._error = False self.commands.read() self.target_grip.read() - self._arm_player.set([]) - self._grip_player.set([]) self._last_grip = 0.0 self._reset_pending = True diff --git a/positronic/simulator/robolab/adapter.py b/positronic/simulator/robolab/adapter.py index 34a4413ab..e4afbda3e 100644 --- a/positronic/simulator/robolab/adapter.py +++ b/positronic/simulator/robolab/adapter.py @@ -34,7 +34,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = eef_pose * self.env_control_frame.inv state = MujocoFrankaState() state.encode(raw_obs['joint_pos'], raw_obs['joint_vel'], ee_pose) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} # TODO: honour a camera_dict naming any other RoboLab camera. env.py renders only the WRIST_LEFT # preset (over_shoulder_left + wrist) and hard-codes emitting those two, so a request for e.g. # over_shoulder_right_camera raises below. The full fix threads the requested set end-to-end: carry diff --git a/positronic/tests/test_data_collection.py b/positronic/tests/test_data_collection.py index f114fad74..b84babab2 100644 --- a/positronic/tests/test_data_collection.py +++ b/positronic/tests/test_data_collection.py @@ -7,12 +7,7 @@ import pimm from positronic import keys, wire from positronic.data_collection import DataCollectionController, OperatorPosition, controller_positions_serializer -from positronic.dataset.ds_writer_agent import ( - DsWriterAgent, - DsWriterCommand, - DsWriterCommandType, - TrajectoryOverrideSerializer, -) +from positronic.dataset.ds_writer_agent import DsWriterAgent, DsWriterCommand, DsWriterCommandType from positronic.dataset.episode import Episode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers @@ -162,18 +157,18 @@ def test_data_collection_with_mujoco_robot_gripper(tmp_path): writer_cm = LocalDatasetWriter(tmp_path) agent = DsWriterAgent(writer_cm.__enter__()) - agent.add_signal('target_grip', TrajectoryOverrideSerializer(None)) - agent.add_signal(keys.ROBOT_COMMAND, TrajectoryOverrideSerializer(Serializers.robot_command)) + agent.add_signal(keys.TARGET_GRIP) + agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) agent.add_signal('controller_positions', controller_positions_serializer) - agent.add_signal('robot_state', Serializers.robot_state) + agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) agent.add_signal(keys.GRIP) world.connect(sim.state, dc.robot_state) - world.connect(sim.state, agent.inputs['robot_state']) + world.connect(sim.state, agent.inputs[keys.ROBOT_STATE]) world.connect(dc.robot_commands, sim.commands) world.connect(dc.robot_commands, agent.inputs[keys.ROBOT_COMMAND]) world.connect(dc.target_grip, sim.target_grip) - world.connect(dc.target_grip, agent.inputs['target_grip']) + world.connect(dc.target_grip, agent.inputs[keys.TARGET_GRIP]) world.connect(sim.grip, agent.inputs[keys.GRIP]) world.connect(dc.ds_agent_commands, agent.command) @@ -213,7 +208,7 @@ def stop_episode(): ep = ds[0] assert isinstance(ep, Episode) - expected = {'target_grip', 'controller_positions.right', keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP} + expected = {keys.TARGET_GRIP, 'controller_positions.right', keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP} assert expected.issubset(set(ep.keys())) # Robot/gripper signals should have at least one sample @@ -258,9 +253,9 @@ def finger_qpos() -> float: grip = world.pair(sim.grip) driver = ManualDriver([ - (lambda: target_grip.emit([(world.clock.now_ns(), 1.0)]), 0.5), + (lambda: target_grip.emit(1.0), 0.5), (lambda: snapshots.update(closed=finger_qpos(), closed_grip=grip.read().data), 0.0), - (lambda: target_grip.emit([(world.clock.now_ns(), 0.0)]), 0.5), + (lambda: target_grip.emit(0.0), 0.5), (lambda: snapshots.update(opened=finger_qpos(), opened_grip=grip.read().data), 0.0), ]) diff --git a/positronic/tests/test_inference_integration.py b/positronic/tests/test_inference_integration.py index 8b3d1826b..c6a08a874 100644 --- a/positronic/tests/test_inference_integration.py +++ b/positronic/tests/test_inference_integration.py @@ -144,7 +144,7 @@ def close(self): assert keys.EE_POSE in last_obs # The task's instruction is injected by the harness. assert last_obs[keys.TASK] == 'integration-test' - assert last_obs['descriptor'] == 'mujoco.franka' + assert last_obs[keys.DESCRIPTOR] == 'mujoco.franka' class _CountdownProducer(pimm.ControlSystem): diff --git a/positronic/utils/registration.py b/positronic/utils/registration.py index 5cfcb30ab..cc2c8ad25 100644 --- a/positronic/utils/registration.py +++ b/positronic/utils/registration.py @@ -186,7 +186,7 @@ def perform_registration(webxr, robot_arm): while not w.should_stop: if move_throttler.wait_time() <= 0: waypoint = roboarm_command.CartesianPosition(WAYPOINTS[current_point]) - commands.emit([(w.clock.now_ns(), waypoint)]) + commands.emit(waypoint) current_point += 1 if current_point >= len(WAYPOINTS): break diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 6071fb46e..04b2805de 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -19,7 +19,7 @@ from positronic.policy import Policy, Session from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints from positronic.utils.logging import init_logging from positronic.vendors import gr00t @@ -354,7 +354,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=codecs.ee_quat, source=gr00t_source) def pipeline(codec, source): - return ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source # Each entry pairs the codec with the matching GR00T modality config; they must agree with training. diff --git a/positronic/vendors/lerobot/server.py b/positronic/vendors/lerobot/server.py index 09f4784ec..443534182 100644 --- a/positronic/vendors/lerobot/server.py +++ b/positronic/vendors/lerobot/server.py @@ -12,7 +12,7 @@ from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, Pipeline, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.utils.logging import init_logging from positronic.vendors.lerobot import codecs as lerobot_codecs @@ -61,7 +61,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=lerobot_codecs.ee, source=lerobot_source) def pipeline(codec: Codec, source: ModelSource) -> Pipeline: - return ChunkedSchedule() | RestrictImageSize(512, 512) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(512, 512) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 5e608c1b1..3128394fb 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -15,7 +15,7 @@ from positronic.policy import Codec, Policy, Session from positronic.policy.observation import TASK_FIELD from positronic.policy.spec import PolicySource, inline -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import resolve_checkpoint from positronic.vendors.lerobot_0_3_3.backbone import register_all @@ -127,4 +127,4 @@ def act(checkpoints_dir: str, checkpoint: str | None, n_action_steps: int | None ) def act_absolute(base: Policy, codec: Codec): """ACT with the absolute-position codec, composed in-process.""" - return inline(ChunkedSchedule() | codec | PolicySource(base)) + return inline(StopOnFault() | ChunkedSchedule() | codec | PolicySource(base)) diff --git a/positronic/vendors/lerobot_0_3_3/server.py b/positronic/vendors/lerobot_0_3_3/server.py index 1996a9bb4..a6f740442 100644 --- a/positronic/vendors/lerobot_0_3_3/server.py +++ b/positronic/vendors/lerobot_0_3_3/server.py @@ -15,7 +15,7 @@ from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.utils.logging import init_logging from positronic.vendors.lerobot_0_3_3 import codecs as lerobot_codecs @@ -84,7 +84,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=lerobot_codecs.ee, source=lerobot_source) def pipeline(codec: Codec, source: ModelSource): - return ChunkedSchedule() | RestrictImageSize(224, 224) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/molmoact2/server.py b/positronic/vendors/molmoact2/server.py index d43a3f4c9..96cfb487d 100644 --- a/positronic/vendors/molmoact2/server.py +++ b/positronic/vendors/molmoact2/server.py @@ -9,7 +9,7 @@ from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.logging import init_logging from positronic.vendors.molmoact2 import codecs as molmoact2_codecs from positronic.vendors.molmoact2.policy import MolmoAct2Policy, warm_observation @@ -60,7 +60,7 @@ def meta(self, model_id: str) -> dict[str, Any]: @cfn.config(codec=molmoact2_codecs.droid, source=molmoact2_source) def pipeline(codec: Codec, source: ModelSource): - return ChunkedSchedule() | RestrictImageSize() | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize() | remote | codec | source droid = pipeline diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 519b0ad37..bdb39b8b1 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -16,7 +16,7 @@ from positronic.policy import Codec, Policy, Session from positronic.policy.codec import ChangeEEFrame, RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import get_latest_checkpoint, list_checkpoints from positronic.utils.logging import init_logging from positronic.vendors import openpi @@ -258,7 +258,7 @@ def pipeline(codec: Codec, source: ModelSource, ee_frame: geom.Transform3D | Non ``ee_frame`` places the end-effector frame this checkpoint's poses live in relative to ``DEFAULT_FRAME`` (``models.DROID_EE_FRAME``); ``None`` for a checkpoint trained in ``default``, or one speaking joints. """ - local = ChunkedSchedule() | RestrictImageSize(224, 224) + local = StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) if ee_frame is not None: # Outermost, so everything downstream — the wire, the server's codec — sees poses already in ``ee_frame``. local = ChangeEEFrame(ee_frame) | local diff --git a/positronic/wire.py b/positronic/wire.py index cdc15af7d..0be4dbe79 100644 --- a/positronic/wire.py +++ b/positronic/wire.py @@ -3,7 +3,7 @@ import pimm from positronic import keys, telemetry, telemetry_keys from positronic.dataset import DatasetWriter -from positronic.dataset.ds_writer_agent import DsWriterAgent, TimeMode, TrajectoryOverrideSerializer +from positronic.dataset.ds_writer_agent import DsWriterAgent, TimeMode from positronic.dataset.serializers import Serializers, StatefulSerializer from positronic.eval import ROBOT_STATIC_META, Embodiment, Observation @@ -44,19 +44,17 @@ def wire( # noqa: C901 for signal_name in cameras.keys(): ds_agent.add_signal(signal_name, Serializers.camera_images) if robot_arm is not None: - # Command channels carry whole trajectories; flatten with last-writer-wins so the - # recording is a dense per-command stream. See TrajectoryOverrideSerializer. - ds_agent.add_signal(keys.ROBOT_COMMAND, TrajectoryOverrideSerializer(Serializers.robot_command)) - ds_agent.add_signal('robot_state', Serializers.robot_state) + ds_agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) + ds_agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) if gripper is not None: - ds_agent.add_signal(keys.TARGET_GRIP, TrajectoryOverrideSerializer(None)) + ds_agent.add_signal(keys.TARGET_GRIP) ds_agent.add_signal(keys.GRIP) for signal_name, emitter in cameras.items(): world.connect(emitter, ds_agent.inputs[signal_name]) if robot_arm is not None: world.connect(harness.robot_commands, ds_agent.inputs[keys.ROBOT_COMMAND]) - world.connect(robot_arm.state, ds_agent.inputs['robot_state']) + world.connect(robot_arm.state, ds_agent.inputs[keys.ROBOT_STATE]) if gripper is not None: world.connect(harness.target_grip, ds_agent.inputs[keys.TARGET_GRIP]) world.connect(gripper.grip, ds_agent.inputs[keys.GRIP]) @@ -109,9 +107,7 @@ def wire_embodiment( ds_agent.add_signal(name, obs.serializer) world.connect(obs.source, ds_agent.inputs[name]) for name, cmd in embodiment.commands.items(): - # Command channels carry whole trajectories; flatten with last-writer-wins so the - # recording is a dense per-command stream. See TrajectoryOverrideSerializer. - ds_agent.add_signal(name, TrajectoryOverrideSerializer(cmd.serializer)) + ds_agent.add_signal(name, cmd.serializer) world.connect(harness.commands[name], ds_agent.inputs[name]) for name, priv in privileged.items(): ds_agent.add_signal(name, priv.serializer) diff --git a/utilities/fake_dataset_generator.py b/utilities/fake_dataset_generator.py index 3523927c8..aa0aba0ab 100644 --- a/utilities/fake_dataset_generator.py +++ b/utilities/fake_dataset_generator.py @@ -138,7 +138,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): 'eval.tote_placement': random.choice(['left', 'right']), 'eval.external_camera': random.choices(['left', 'right', 'NA'], [5, 5, 1])[0], 'inference.policy_fps': self.fps, - 'inference_latency': False, + keys.INFERENCE_LATENCY: False, **self.policy_meta, }