Add GR00T policy-runner OSMO workflow - #848
Conversation
Record every value drawn by an enabled variation's sampler so downstream sensitivity-analysis tooling has the input factors that produced each episode. This adds a sample-observer layer the recorder builds on: SamplerBase gains add_listener/remove_listener and a sample() template method that notifies listeners around the concrete _sample(); VariationBase gains add_sample_listener/remove_sample_listener, re-binding subscriptions onto the sampler rebuilt by apply_cfg so they survive cfg swaps. ArenaEnvBuilder constructs a VariationRecorder, attaches it after Hydra overrides but before any sampling, and exposes it on env.unwrapped. The recorder is stashed on the builder rather than the env cfg because the configclass __post_init__ deep-copies its attributes, which would orphan the listener closures. Signed-off-by: alex <amillane@nvidia.com>
…riations_recording
…riations_recording
…riations_recording
…ecording_per_episode
Move the inline RecordVideo / CameraObsVideoRecorder wrapping out of policy_runner and eval_runner into a single video_recording module (VideoRecordingCfg + wrap_env_for_video), so the gym-wrapper plumbing lives in one place instead of being duplicated across the two runners. Signed-off-by: alex <amillane@nvidia.com>
- Move video_recording.py to isaaclab_arena/utils/ and drop its module docstring. - Rename VideoRecordingCfg fields and the corresponding CLI flags: --video -> --record_viewport_video, --camera_video -> --record_camera_video, --video_dir -> --video_base_dir. - Write videos into a reverse-dated run subdirectory (timestamped_run_dir), shared across all jobs in an eval run, mirroring Isaac Lab's log layout. - Guard _resolve_video_length when both num_steps and num_episodes are None. - Tidy wrap_env_for_video docstring args. Signed-off-by: alex <amillane@nvidia.com>
Use underscore-only flag names (--record_viewport_video, --record_camera_video, --video_base_dir) to match the convention used by every other CLI argument. Signed-off-by: alex <amillane@nvidia.com>
Greptile SummaryIntroduces the GR00T inference-server task (
Confidence Score: 5/5Safe to merge; the new GR00T workflow is additive, the pi0 refactor keeps existing behaviour intact, and all current call sites pass All changed code paths are either new additions or straightforward refactors of the pi0 workflow to use shared constants. No existing workflows break: No files require special attention for this merge; the Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant S as submit_gr00t_policy_runner_workflow.py
participant W as Gr00tPolicyRunnerWorkflow
participant RT as Gr00tPolicyRunnerTask (lead)
participant ST as Gr00tServerTask
S->>W: build_parser() → parse args
S->>W: __init__(workflow_args, task_args)
W->>W: _get_tasks()
W->>ST: "host_token() → {{host:gr00t_server}}"
W->>RT: "__init__(..., remote_host={{host:gr00t_server}}, lead=True)"
W->>ST: "__init__(..., lead=False)"
S->>W: submit_workflow(dry_run, pool, priority)
W->>W: render_yaml()
Note over RT: Waits for GR00T server ready
RT-->>ST: gRPC connect :8000
ST-->>RT: policy inference responses
RT->>RT: Runs evaluation episodes → uploads to DATASET_SWIFT_URL
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant S as submit_gr00t_policy_runner_workflow.py
participant W as Gr00tPolicyRunnerWorkflow
participant RT as Gr00tPolicyRunnerTask (lead)
participant ST as Gr00tServerTask
S->>W: build_parser() → parse args
S->>W: __init__(workflow_args, task_args)
W->>W: _get_tasks()
W->>ST: "host_token() → {{host:gr00t_server}}"
W->>RT: "__init__(..., remote_host={{host:gr00t_server}}, lead=True)"
W->>ST: "__init__(..., lead=False)"
S->>W: submit_workflow(dry_run, pool, priority)
W->>W: render_yaml()
Note over RT: Waits for GR00T server ready
RT-->>ST: gRPC connect :8000
ST-->>RT: policy inference responses
RT->>RT: Runs evaluation episodes → uploads to DATASET_SWIFT_URL
Reviews (3): Last reviewed commit: "Self review and cleanup." | Re-trigger Greptile |
![arena-review-bot[bot]](https://avatars.githubusercontent.com/u/195794629?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
Adds a two-task OSMO workflow (GR00T inference server + Arena policy runner, lead) that gates the eval on the server coming up and uploads per-run eval outputs. The structure closely follows the established pi0 sibling files, the wait-for-server flags line up exactly with isaaclab_arena_gr00t.utils.wait_for_gr00t_server, and the runner reuses server_port for both the server bind and the client connect so the two can't drift. Only minor nits below.
Findings
🔵 Improvement: osmo/tasks/gr00t_policy_runner_task.py:41 — the getattr(task_args, "arena_image", image) fallback never fires (see inline).
🔵 Improvement: osmo/tasks/gr00t_server_task.py:54 — same dead getattr fallback (see inline).
🔵 Nit: osmo/tasks/gr00t_policy_runner_task.py:67 — stray paren in help text (see inline).
Test Coverage
These are OSMO submission-tooling files (rendered YAML + bash entry scripts) that mirror the test_gr00t_closedloop_e2e CI job the docstring references; like the pi0/zero-action siblings they aren't unit-tested here, which is consistent with the existing pattern. No new test needed.
Verdict
Ship it
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
alexmillane
left a comment
alexmillane
left a comment
There was a problem hiding this comment.
Partial self-review.
- Slim Gr00tPolicyRunnerTask to just its policy args + gr00t-specific flags; rely on the base for the run script, outputs, task name, and image (drops the out-of-scope image override and the wait-for-server preamble for now). - Align Gr00tServerTask with pi0_server_task: inline task name, drop the wait-for-server helper, and fold model/embodiment/port into constants. - Centralize the policy-server port in workflow_constants (POLICY_SERVER_PORT = 8000) and use it for both the GR00T and pi0 remote tasks. - Reverse the workflow task order to match the pi0 workflow. Signed-off-by: alex <amillane@nvidia.com>
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
Adds a GR00T policy-runner OSMO workflow (a GR00T inference server + a lead policy runner) that closely mirrors the existing pi0 two-task workflow, and centralizes the shared server port into POLICY_SERVER_PORT. The change is small, cleanly reuses the PolicyRunnerTask/Workflow/BaseTask bases, and the get_task_name() / {{host:gr00t_server}} tokens line up correctly. One minor duplication observation inline, plus a note on the PR description.
Findings
See the inline comment on osmo/tasks/gr00t_server_task.py.
Test Coverage
No tests are added. This is consistent with the sibling pi0 and zero-action workflows, which have none either — these are YAML-rendering submit scripts, and a --dry-run render smoke test would be the natural coverage. That's a pre-existing gap rather than something this PR must close, so not blocking.
Note on the PR description
The detailed description says the PR "Adds EVAL_OUTPUT_* constants for per-run evaluation output uploads," but the diff only adds POLICY_SERVER_PORT — no EVAL_OUTPUT_* constants exist in workflow_constants.py at HEAD. This looks stale after the "Address review" commit slimmed things down; worth updating the description so reviewers aren't looking for code that isn't there.
Verdict
Ship it
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
This PR adds a two-task OSMO workflow (GR00T inference server + Arena policy-runner) plus a submit script, and refactors the pi0 remote runner to share the same server-host wiring via a new BaseTask.host_token() helper and a centralized POLICY_SERVER_PORT. The composition is clean — server specifics live in extension-package config paths, the runner reaches the server through an OSMO host token, and the port is a single source of truth. My main comment is that the server↔runner wiring is now duplicated across the gr00t and pi0 workflows.
Design, Boundaries & Scope
The new _get_tasks override in gr00t_policy_runner_workflow.py is near-identical to the one this PR also adds to pi0_plus_policy_runner_workflow.py — both build [runner(remote_host=Server.host_token()), server]. Since task_cls_list already names the runner/server pair, this wiring could live once in a shared base (RemotePolicyRunnerWorkflow(Workflow)), with the subclasses only declaring the class pair. See the inline comment for a concrete shape. Not a blocker, but the two bodies will drift.
Findings
🟡 Warning: osmo/workflows/gr00t_policy_runner_workflow.py:20 — duplicated server-host wiring; lift into a shared base (inline).
🔵 Improvement: osmo/submit_gr00t_policy_runner_workflow.py:13 — "Default" usage example is missing the required --arena_env_args (inline).
🔵 Improvement: osmo/workflows/workflow_constants.py — POLICY_SERVER_PORT is now the single source for the client port and the GR00T server's --port, but Pi0ServerTask still relies on serve_policy.py's implicit default rather than the constant, so the pi0 server bind port and the client's POLICY_SERVER_PORT can silently diverge. Consider threading the constant into the pi0 server too (its file isn't touched here, so strictly out of this PR's scope).
Test Coverage
These are pure-Python workflow-config generators (they render YAML dicts, no sim), so the sim inner/outer pattern doesn't apply. There don't appear to be unit tests exercising the new wiring; a small test asserting that _get_tasks() gives the runner remote_host == "{{host:gr00t_server}}" would cheaply guard the runner↔server task-name sync (which host_token() couples). Optional given alpha, but it's the one thing most likely to break silently if a task name changes.
Verdict
Minor fixes needed
Note: the PR description mentions adding EVAL_OUTPUT_* constants, but the final diff only adds POLICY_SERVER_PORT — description looks stale.
- Quote --model_path/--embodiment_tag in the GR00T server run script. - Give the GR00T server image an explicit :latest tag. - Add nvidia-smi to the pi0 server run script for symmetry with GR00T. Signed-off-by: alex <amillane@nvidia.com>
|
/review |
🤖 Isaac Lab-Arena Review BotSummaryThis PR adds a GR00T server + policy-runner OSMO workflow ( Findings🟡 Warning — 🔵 Improvement — Test CoverageThis is OSMO submit/deploy tooling (workflow-dict assembly, no sim), so the inner/outer simulation-test pattern does not apply. No unit tests accompany the new task/workflow classes; given the deterministic dict-rendering surface, a light test that the gr00t workflow renders two tasks with the runner as lead and VerdictShip it — minor, non-blocking suggestions only. |
## Summary Add a GR00T policy-runner evaluation OSMO workflow. ## Detailed description - Adds `Gr00tServerTask` and `Gr00tPolicyRunnerTask`, their two-task workflow, and a submit script. - The server serves the GR00T droid checkpoint; the lead policy-runner task waits for it, then evaluates the GR00T remote closed-loop policy. - Adds `EVAL_OUTPUT_*` constants for per-run evaluation output uploads. --------- Signed-off-by: alex <amillane@nvidia.com> Co-authored-by: aiguldzh-nvidia <adzhumamurat@nvidia.com>
3 participants
Summary
Add a GR00T policy-runner evaluation OSMO workflow.
Detailed description
Gr00tServerTaskandGr00tPolicyRunnerTask, their two-task workflow, and a submit script.EVAL_OUTPUT_*constants for per-run evaluation output uploads.