feat(agents): add Jupyter-backed code policy module - #3259
Conversation
Codecov Report❌ Patch coverage is
@@ Coverage Diff @@
## main #3259 +/- ##
==========================================
+ Coverage 75.20% 75.26% +0.06%
==========================================
Files 1128 1131 +3
Lines 107954 108307 +353
Branches 9751 9771 +20
==========================================
+ Hits 81188 81521 +333
- Misses 23955 23965 +10
- Partials 2811 2821 +10
Flags with carried forward coverage won't be shown. Click here to find out more.
... and 3 files with indirect coverage changes 🚀 New features to boost your workflow:
|
TomCC7
changed the title
|
cleaned up, can check again @leshy |
e25829b to
06a77a2
Compare
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
github-actions
Bot
added
the
ready-to-merge
github-actions
Bot
removed
the
ready-to-merge
| return f"In [{execution_count}] {state} in {duration_s:.2f}s\n\n{body}" | ||
|
|
||
|
|
||
| code_policy_module = CodePolicyModule.blueprint |
| super().start() | ||
|
|
||
| @skill | ||
| def python_exec(self, code: str, timeout_s: float = MAX_EXECUTION_TIMEOUT_S) -> str: |
There was a problem hiding this comment.
Note that this cannot exceed the RPC timeout, which I think is 30 seconds. You might want to extend that for this call. I think Jeff added a way to extend RPC timeouts for particular methods.
| # Runtime objects remain None until this instance has been deployed into | ||
| # its worker, so Module serialization needs no custom hooks. |
There was a problem hiding this comment.
I'm not sure what this means... Why not just do:
self._execution_lock: threading.Lock = threading.Lock()
here?
| lock = self._execution_lock | ||
| if lock is None: | ||
| lock = self._execution_lock = threading.Lock() |
There was a problem hiding this comment.
This is quite odd.
If you're mutating self._execution_lock from different threads, you need a lock for it... but it's odd to need a lock for a lock.
Of course self._execution_lock won't be none in practice because self.start() initializes it. But you can avoid all of this by just initializing it in __init__.
| """ | ||
|
|
||
|
|
||
| class CodePolicyModule(Module): |
There was a problem hiding this comment.
Hmm, the architecture is a bit odd.
You have to run dimos run ..., which spawns several processes, one of which runs CodePolicyModule. CodePolicyModule spawns a Jupyter python process which runs Dimos.connect() which connects to the dimos run... process. A bit circular.
There was a problem hiding this comment.
Agree. You actually raised a very good opportunity to do separation here. We can actually just write a script to expose a single execute_python mcp skill to agent so agent won't get confused by all other mcp skills. I'll check if this is better
|
|
||
|
|
||
| class CodePolicyModule(Module): | ||
| """Execute trusted agent-authored Python against a running DimOS system.""" |
There was a problem hiding this comment.
I haven't read the full paper, but I think a key part is being able to generate code when an undefined function is called.
Central to our approach is prompting hierarchical code-gen (recursively defining undefined functions),
There was a problem hiding this comment.
The original code as policy paper is in 2022, which is even pre-chatgpt era (1000 years ago in AI lol). So you might want to check this instead: https://research.nvidia.com/labs/gear/aspire/
And there is indeed a naming issue of this module, it's only a small part of the code-as-policy flow.
|
pretty cool feature |
|
Really nice direction — a persistent, agent-authored Python action space is a big unlock beyond one-shot skill calls. I ran the branch on a sim xArm setup (minimal blueprint + an OpenAI-compatible LLM backend) and the core loop works well: one On the naming / scope discussion (paul-nechifor's point about hierarchical code-gen): agreed this module is not full CaP — but I'd argue hierarchical code-gen is not the part of the 2022 paper worth chasing. Two more recent data points reframe where the leverage is:
What the loop's other half would look like (v2 candidates, non-blocking):
Related work worth a look if you haven't: CodeAct (Wang et al., ICML 2024, arXiv:2402.01030) — evidence that code as the action format outperforms JSON/text tool calls for multi-step control flow in agents, which is exactly the bet this PR makes; and Graph-as-Policy (2026) for a contrasting persistent-execution representation. Two operational notes from my run: (a) while testing I hit an unrelated daemon-mode Zenoh bug that makes the kernel bootstrap fail with "No running DimOS coordinator found" under Happy to prototype the structured-trace piece against this branch if there's interest. |
|
closed in lieu of #3378 |
5 participants
Contribution path
Problem
DimOS agents can call predefined skills, but they cannot submit a Python program that processes native observations, branches or retries, and composes deployed RPCs as one synchronous policy rollout.
Solution
CodePolicyModuleexposing one synchronouspython_exec(code, timeout_s)MCP skill.jupyter_client+ipykernelstack instead of a custom worker protocol.Dimosapp handle and a Memory2SqliteStoreattached to the explicitly configured recorder database.Recorderbeside the xArm simulation-agent blueprint. It records the enabled joint-state and color-image streams to a stable database path without changing Memory2 or perception.ipykernelandjupyter-clientare explicit dependencies of theagentsextra and are imported lazily. The v1 interface is foreground-only, text-only, single-call-at-a-time, and intentionally not a security sandbox. Real-hardware blueprints are unchanged.How to Test
Install and run the xArm simulation:
From another terminal:
uv run dimos mcp call python_exec --json-args \ '{"code":"print(memory.list_streams()); print(app.skills.get_robot_state())"}'Focused automated validation:
The focused suite passes 19 tests. Ruff, strict mypy over the changed source files, blueprint registry generation, staged pre-commit hooks, and
git diff --checkalso pass. A live MCP smoke test verified thatpython_execcan read the latest joint state and color image from the active Recorder database and query bounded observation history. Earlier validation covered DimOS RPC calls, namespace persistence, and interrupt-first timeout recovery.AI assistance
OpenAI Codex with GPT-5 was substantially involved in design discussion, implementation, tests, documentation, and validation. The author reviewed and approved the direction interactively.
Checklist