Pi0.5 on OSMO - #823
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>
alexmillane
left a comment
alexmillane
left a comment
There was a problem hiding this comment.
More self-review.
![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 pi0 policy-server OSMO workflow (server + policy-runner client co-scheduled over {{host:policy_server}}) and refactors the workflow/task layer so tasks declare their own CLI args via add_task_arguments and workflows compose tasks via task_cls_list/lead_list. The refactor is clean, the single-lead invariant is asserted, and the vendor-specific websocket override is correctly placed in the isaaclab_arena_openpi extension package rather than core. Since this is a DRAFT, the main things to resolve before un-drafting are a few temporary/debug leftovers in policy_runner_task.py.
Findings
See inline comments.
Test Coverage
The websocket-client override is exercised by the existing test_pi0_remote_policy.py (the monkeypatch target was correctly updated to the new WebsocketClientPolicy). OSMO submit-script wiring stays untested, consistent with the rest of osmo/, and no new sim-dependent code paths were added that would need the inner/outer pattern.
Verdict
Minor fixes needed
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
This draft adds a pi0 policy-server workflow (server co-scheduled with the policy-runner client over OSMO's {{host:policy_server}} substitution) and cleans up the OSMO workflow structure: workflows declare tasks via task_cls_list/lead_list, tasks own their CLI args via add_task_arguments, and per-policy PolicyRunnerTask subclasses replace the old workflow-type→args map. The refactor is coherent and the new _assert_single_lead_task invariant reads well. The main thing holding it back from merge is leftover debug/temporary scaffolding in policy_runner_task.py — flagged below (all confined to osmo/ deployment scripts, so no core-boundary concerns).
Findings
🟡 Warning: osmo/tasks/policy_runner_task.py:18 — DEFAULT_IMAGE points at a personal :alex_testing tag (was :latest). This looks like a temporary testing value that would ship to everyone submitting a policy-runner workflow — worth reverting to :latest (or a real pinned tag) before this leaves draft.
🟡 Warning: osmo/tasks/policy_runner_task.py:67 — _get_inputs carries a # REMOVE! marker and the arena-lfs-data input commented out, now returning []. Is dropping the LFS test-data input intended, or WIP scaffolding? Either way the commented line + # REMOVE! should go before merge.
🔵 Improvement: osmo/tasks/policy_runner_task.py:84 — the old _get_run_script body is left commented out and the live path is now just set -euxo pipefail\n{command}. Beyond deleting the dead block: are the dropped setup steps (ldconfig, the _isaac_sim symlink, the LFS cp) now handled by the image/entrypoint? If the runner relies on any of them this would fail at runtime — worth a confirming note.
Test Coverage
test_pi0_remote_policy.py is updated to patch the new WebsocketClientPolicy override, so the existing suite stays green. The new keepalive params on the override aren't directly exercised, but they're thin passthroughs to websockets.sync.client.connect — not worth a dedicated test. No new test needed for the OSMO submit scripts (deployment glue).
Verdict
Minor fixes needed — clear out the # REMOVE! / commented scaffolding and the alex_testing image tag before taking this out of draft.
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
Draft PR adds a pi0 policy-server OSMO workflow (server + policy-runner co-scheduled) and refactors the OSMO workflow/task structure so tasks own their CLI args, workflows declare a task_cls_list with per-task lead flags, and build logic is shared. The refactor is a clean simplification overall — dropping the WorkflowType/PolicyType enums and the single-lead assertion read well, and extracting build_server_image.sh de-duplicates the docker build. A few things below are worth confirming before this leaves draft, mostly around behavior that silently changed in the policy-runner task and an edge in the new arg-dedup.
Findings
🟡 Warning: osmo/workflows/workflow.py:77 — _deduplicate_args dedups by class identity, but add_task_arguments is inherited; two sibling subclasses sharing it would each add the same args and argparse would raise. Dedup on the resolved method instead.
🟡 Warning: osmo/tasks/policy_runner_task.py:83 — the env-setup boilerplate and the arena-lfs-data input were dropped from the run script; confirm the entrypoint now covers this.
🔵 Improvement: osmo/workflows/workflow_constants.py:16 — three URL constants are unused.
🔵 Improvement: osmo/tasks/pi0_server_task.py:14 — policy config/dir are hardcoded; consider exposing them.
Test Coverage
osmo/ has no pytest harness, so the new pure-logic bits (_assert_single_lead_task, _deduplicate_args, build_parser) are untested — consistent with the existing lack of coverage there, but they're now the kind of logic a small unit test would protect. On the policy side, test_pi0_remote_policy.py was correctly updated to patch the new WebsocketClientPolicy, but nothing asserts that ping_interval/ping_timeout actually reach the client — a light addition given they're the point of the override.
Verdict
Minor fixes needed (draft).
Greptile SummaryThis PR adds a pi0.5 inference-server workflow to OSMO, co-scheduling a
Confidence Score: 4/5Safe to merge with one fix: the documented minimal invocations in both submit-script docstrings will fail at argument parsing because --arena_env_args is now required with no default. The refactor is structurally clean and the new pi0 workflow logic is sound. The one concrete breakage is that both submit-script docstrings advertise a minimal invocation that will immediately error out because --arena_env_args has no default. osmo/tasks/policy_runner_task.py and the two submit scripts whose docstrings show the now-broken minimal invocation. Important Files Changed
Reviews (3): Last reviewed commit: "Bring memory limit back down to make it ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
Adds a co-scheduled pi0 policy-server + policy-runner OSMO workflow and refactors the OSMO generator into per-workflow classes with per-task CLI args and a per-task lead flag. The refactor reads cleanly and the layering is sound (the vendor-specific websocket override lives in isaaclab_arena_openpi, the generic mechanism stays in osmo/workflows). My main concern is that the refactor quietly changed the documented default policy-runner path so that the arg-less "default" invocations the submit scripts advertise no longer run.
Design, Boundaries & Scope
The move from a single enum-switched submit_evaluation_workflow.py to explicit per-policy task/workflow classes is justified — different policies genuinely need different args and run scripts (pi0 needs a co-scheduled server + connection flags), which the old WorkflowType/PolicyType enums couldn't model. Dropping the ldconfig / mkdir / _isaac_sim symlink boilerplate from the run script is a good cleanup — docker/setup/entrypoint.sh already does all of it.
The scope concern is default behavior: --policy_runner_args now defaults to "" (was --num_steps 100 --headless) and --arena_env_args is now required=True with no default. Together these mean the "default" commands the submit-script docstrings advertise now fail (details inline). Worth confirming this is intended and updating the docs.
Findings
See inline comments.
Test Coverage
test_pi0_remote_policy.py is correctly updated to patch the new WebsocketClientPolicy subclass, and the keepalive plumbing is exercised indirectly. The OSMO refactor itself (per-task lead flags, _assert_single_lead_task, build_parser/_deduplicate_args) has no test — a small --dry-run/render test asserting the generated dict (exactly one lead, expected task names, a non-failing default parse) would have caught the broken-default and required-arg issues above cheaply. Not a blocker for infra glue, but worth considering.
Verdict
Minor fixes needed
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
Adds a pi0 policy-server + policy-runner OSMO workflow and refactors the OSMO generator into per-workflow classes with declarative task_cls_list / lead_list, task-owned CLI args, and a shared build_parser. The refactor is clean and the class-attribute-driven workflow model reads well. A couple of points below are worth a look before merge — chiefly whether dropping the run-script's ldconfig step is safe on OSMO, plus some duplication and a fake default.
Findings
🟡 Warning: osmo/tasks/policy_runner_task.py:83 — the run script no longer runs ldconfig, which previously mirrored entrypoint.sh. OSMO overrides the image ENTRYPOINT, so /entrypoint.sh (and its ldconfig) never runs — see inline.
🔵 Improvement: osmo/submit_zero_action_policy_runner_workflow.py:39 — near-identical main() duplicated across the two submit scripts; could share one helper.
🔵 Improvement: osmo/tasks/base_task.py:26 — lead is effectively required but defaults to None, which would render lead: null.
Test Coverage
test_pi0_remote_policy.py is correctly re-pointed at the new WebsocketClientPolicy subclass and the existing reconnect/adapter tests still cover the client path. The new keepalive plumbing (ping_interval/ping_timeout flowing into the connection) and the overridden _wait_for_server aren't directly exercised, but that's hard to unit-test without a live server and is low-risk config plumbing — no blocker.
Verdict
Minor fixes needed
There was a problem hiding this comment.
🤖 Isaac Lab-Arena Review Bot
Summary
This PR adds a pi0 policy-runner + inference-server OSMO workflow and refactors the OSMO generator scripts: tasks now own their CLI args via add_task_arguments, workflows declare their task list with per-task lead flags, and common args are centralized. The refactor is a clear improvement (the WorkflowType/PolicyType enums and the repeated arg blocks are gone). A couple of default-behavior changes and a bit of residual duplication are worth a look before merge.
Design, Boundaries & Scope
submit_zero_action_policy_runner_workflow.py and submit_pi0_plus_policy_runner_workflow.py have main() bodies that are identical except for the workflow class and the description string. Could we collapse them into one shared entry point — e.g. a submit_workflow_cli(workflow_cls, description, epilog) helper in workflows/ (or a small Workflow.run_cli classmethod) — so a new workflow needs only its Workflow subclass and not another copy of this boilerplate?
Findings
See inline comments.
Test Coverage
test_pi0_remote_policy.py is correctly updated to patch the new WebsocketClientPolicy subclass, so the existing coverage still exercises the reconnect path. The new keepalive-ping override and the refactored OSMO task/workflow assembly (e.g. Workflow._assert_single_lead_task, create_task_dict output) are untested; a small unit test rendering a multi-task workflow dict and asserting the single-lead invariant would guard the new group logic. The OSMO scripts are host-side tooling (not part of the sim test phases), so no inner/outer sim-test pattern is needed here.
Verdict
Minor fixes needed
## Summary Add pi0 policy-server workflow to OSMO. Refactor OSMO generator scripts in the process. ## Detailed description - **pi0 workflow** - Schedules an openpi policy server alongside the policy-runner client - _Note that we had to override the one function from the official pi client to stop dropouts_ - **Refactor OSMO workflow structure** - Common args no longer repeated. - Workflows declare multiple tasks. - Adds multi-task groups with a per-task `lead` flag (exactly one lead per group). - Tasks own their CLI args via `add_task_arguments` --------- Signed-off-by: alex <amillane@nvidia.com> Co-authored-by: aiguldzh-nvidia <adzhumamurat@nvidia.com>
3 participants
Summary
Add pi0 policy-server workflow to OSMO. Refactor OSMO generator scripts in the process.
Detailed description
leadflag (exactly one lead per group).add_task_arguments