feat: add agent-triggered Go2 episode recording#3117
Conversation
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…sodeMonitor Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ord blueprint Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…guidance Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Greptile SummaryThis PR adds agent-triggered episode recording for the Go2 workflow. The main changes are:
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (4): Last reviewed commit: "fix(collection): preserve transition pub..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if event == "start" and task_label is not None: | ||
| with self._lock: | ||
| self._task_label = task_label | ||
| self._transition(cast("EpisodeCommand", event), time.time()) |
There was a problem hiding this comment.
RPC calls run concurrently, but this path releases _lock after assigning _task_label and reacquires it in _transition. Two simultaneous starts can publish one caller's episode with the other caller's label, so the returned message names one task while DataPrep persists another.
There was a problem hiding this comment.
Fixed. set_episode now passes the requested label into _transition, and the label assignment plus state mutation happen under the same lock. A concurrent-start regression test verifies that alpha and beta each retain their own emitted label.
| if event == "save": | ||
| return f"Episode saved. Total saved this session: {saved}." | ||
| return f"Episode discarded. Total discarded this session: {discarded}." |
There was a problem hiding this comment.
When save or discard arrives while idle, _transition changes no counter and DataPrep has no pending episode to commit. These unconditional responses still tell the agent that an episode was saved or discarded, creating a phantom successful operation for duplicate or out-of-order calls.
There was a problem hiding this comment.
Fixed. Save and discard while idle now return a clear no-op response and emit no EpisodeStatus event. Both paths have regression coverage.
| - `where_am_i` gives your current street/area and nearby landmarks | ||
| - `map_query` finds places on the OSM map by description and returns coordinates | ||
| ## Episode Recording |
There was a problem hiding this comment.
we inject skill prompt into agent context automatically, so this might not be necessary here.
There was a problem hiding this comment.
Removed the system-prompt addition. The skill docstrings provide the injected tool guidance.
TomCC7
left a comment
TomCC7
left a comment
There was a problem hiding this comment.
Hi, would you fill in the PR description so we have a better understanding of what solution you are providing here? Also it would be good if you can cleanup the modified files a bit (superpower, logs, etc)
| Returns: | ||
| A human-readable status string (also safe to surface to the agent). | ||
| """ | ||
| if event not in ("start", "save", "discard"): |
There was a problem hiding this comment.
Updated the public event parameter to EpisodeAction, a Literal restricted to start, save, and discard. The internal EpisodeCommand still includes toggle for button input.
There was a problem hiding this comment.
I don't think you are using it here. You defined the enum but this if statement still use the bare tuple
| if event not in ("start", "save", "discard"): | |
| if event not in EpisodeAction: |
There was a problem hiding this comment.
The previous EpisodeAction was a typing.Literal alias, so membership against it would not work at runtime. I replaced both aliases with a real EpisodeCommand StrEnum and normalize/validate the incoming value once with EpisodeCommand(event).
| from dimos.learning.collection.episode_monitor_spec import EpisodeControlSpec | ||
|
|
||
|
|
||
| class EpisodeRecordingSkillContainer(Module): |
There was a problem hiding this comment.
seems fine to just add these functions to episode monitor directly
There was a problem hiding this comment.
Done. The three recording skills now live directly on EpisodeMonitorModule; the separate skill container and RPC spec were removed.
| @@ -0,0 +1,18 @@ | |||
| Wed Jul 15 12:43:01 2026 | |||
There was a problem hiding this comment.
This should not be committed.
There was a problem hiding this comment.
Removed MUJOCO_LOG.TXT and the other unrelated local artifacts from the branch.
|
Please fill in the PR description. |
…recording # Conflicts: # dimos/robot/all_blueprints.py
yusufjarada
changed the title
|
Addressed the requested changes and pushed the revision:
Local verification: 47 collection/dataprep tests passed, generated blueprint check passed, Ruff passed, and the simulated MCP start → move → save flow produced all expected DB streams without runtime errors. The remaining Backport check requires a maintainer to add |
Codecov Report❌ Patch coverage is
@@ Coverage Diff @@
## main #3117 +/- ##
==========================================
- Coverage 73.23% 72.98% -0.25%
==========================================
Files 1037 1040 +3
Lines 93895 94031 +136
Branches 8528 8534 +6
==========================================
- Hits 68761 68629 -132
- Misses 22792 23085 +293
+ Partials 2342 2317 -25
Flags with carried forward coverage won't be shown. Click here to find out more.
... and 25 files with indirect coverage changes 🚀 New features to boost your workflow:
|
| EpisodeAction: TypeAlias = Literal["start", "save", "discard"] | ||
| # A button/keyboard press can also request `toggle`, which resolves to | ||
| # `start`/`save` based on the current state and never reaches the output. | ||
| EpisodeCommand: TypeAlias = Literal["start", "save", "discard", "toggle"] |
There was a problem hiding this comment.
These two definitions are nearly the same. Can you elaborate what's the difference between an action and a command? And is it possible to unify it?
There was a problem hiding this comment.
Unified them into one runtime EpisodeCommand StrEnum with start, save, discard, and toggle values. The separate EpisodeAction alias is removed.
| Returns: | ||
| A human-readable status string (also safe to surface to the agent). | ||
| """ | ||
| if event not in ("start", "save", "discard"): |
There was a problem hiding this comment.
I don't think you are using it here. You defined the enum but this if statement still use the bare tuple
| if event not in ("start", "save", "discard"): | |
| if event not in EpisodeAction: |
| if event == "start": | ||
| return ( | ||
| f"Recording started (task: {task_label})." if task_label else "Recording started." | ||
| ) | ||
| if event == "save": | ||
| return f"Episode saved. Total saved this session: {status.episodes_saved}." | ||
| return f"Episode discarded. Total discarded this session: {status.episodes_discarded}." |
There was a problem hiding this comment.
better using a match statement or just a simple dict here as these three cases are same level, and using any as fallback looks weird.
There was a problem hiding this comment.
Changed the response dispatch to an explicit match over the resolved status event, with separate start, save, and discard cases.
| # (inside _transition) so it happens regardless of entry point — button, | ||
| # keyboard, or the set_episode RPC. | ||
| if event in ("save", "discard"): | ||
| self._task_label = self.config.default_task_label |
There was a problem hiding this comment.
I'm not sure what's the settings here. Generally during data collection we do the same task multiple times instead of doing different tasks right? Let me know what you think.
There was a problem hiding this comment.
Updated this to session-level task semantics. Save and discard now preserve the active task label, and start_recording accepts an optional label so repeated takes reuse it. Supplying a new label changes tasks; reset_counters restores the configured default. The MuJoCo smoke test persisted two takes with the same label.
|
Pushed the follow-up review fixes in
Local verification: 51 tests passed, Ruff passed, blueprint generation passed, and a two-take MuJoCo MCP run persisted the same label across both episodes with all expected streams. Ready for re-review. |
| status = self._snapshot(event, ts) | ||
| self._emit(status) | ||
| status = self._snapshot(event.value, ts) | ||
| return self._emit(status) |
There was a problem hiding this comment.
The lock serializes label assignment and snapshot creation, but _emit runs after the lock is released. If concurrent starts use alpha and beta, the alpha transition can snapshot first while beta publishes first. DataPrep then processes the starts in the wrong order and can assign the intervening episode to the wrong label. Serialize publication with the state transition or send snapshots through an ordered publisher.
There was a problem hiding this comment.
Fixed in 550af88. Episode state mutation/snapshotting still uses the state lock, while a separate transition lock now serializes the complete snapshot-and-publish sequence. This preserves publication order without holding the state lock during subscriber I/O. Added a deterministic concurrent-transition regression test that reproduced the alpha/beta inversion before the fix. Focused suite: 52 passed; episode monitor suite: 19 passed; Ruff and diff checks pass.
4 participants
Contribution path
Problem
DimOS can already persist robot streams and export episodes to LeRobot/HDF5, but an agent cannot mark labeled training episodes at runtime. Existing collection workflows depend on teleop buttons.
Solution
start_recording,stop_recording, anddiscard_recordingdirectly fromEpisodeMonitorModule.EpisodeCommandenum for start/save/discard/toggle commands and normalize RPC input once.unitree-go2-agentic-recordblueprint.unitree-go2-agenticblueprint unchanged.cmd_velis produced by the existingMovementManagerin the transitive Go2 agentic blueprint. A blueprint invariant test verifies that producer and recorder input are both present.How to Test
In another terminal:
dimos mcp call start_recording --arg task_label=walk-test dimos mcp call relative_move --arg forward=0.2 dimos mcp call stop_recording # Omit task_label to record another take of walk-test dimos mcp call start_recordingVerification completed locally:
pytest dimos/learning/collection dimos/learning/dataprep/test_core.py dimos/robot/test_all_blueprints_generation.py -q— 51 passedgo2_recorder.pyfocused coverage — 100%AI assistance
Claude Code with Opus 4.8 produced the initial design and implementation. OpenAI Codex with GPT-5 reviewed maintainer feedback, revised the implementation, added regression tests, resolved upstream merges, and ran verification. I reviewed and understand the resulting changes.
Checklist