Distribute Arena Experiment Runs with exact outputs

[cvolkcvolk]

Summary

Run every Arena Run independently on OSMO and publish one Experiment report.

Detailed description

• Create one independently scheduled OSMO group per Run. A Pi0 Run gets its own policy-server task in the same group.
• Keep the timestamped --output_base_dir behavior used by existing local commands. OSMO instead passes --experiment_output_directory so each Experiment Runner writes to the exact directory owned by its task.
• After every Run succeeds, the CPU-only build-experiment-output task collects the Run directories, generates index.html, and publishes one complete Experiment output. If any Run fails, the workflow fails and no combined report is published; retries and partial reports remain out of scope.

Run it on OSMO

This submits the nine Runs in the OpenPI Experiment. The Experiment Runner uses the default nvcr.io/nvstaging/isaac-amr/isaaclab_arena:latest image.

/isaac-sim/python.sh -m osmo.submit_arena_experiment \
    --experiment_cfg isaaclab_arena_environments/experiment_configs/droid_pnp_srl_openpi_experiment.yaml \
    --policy_server pi0 \
    osmo.workflow_name=arena-parallel-test \
    osmo.pool=isaac-dev-l40-03 \
    osmo.platform=ovx-l40 \
    osmo.memory=120Gi \
    experiment_runner.image=nvcr.io/nvstaging/isaac-amr/isaaclab_arena:osmo_parallel

[arena-review-bot]

🤖 Isaac Lab-Arena Review Bot

Summary

This PR fans a Pi0 Arena Experiment out into one independently scheduled OSMO group per Run (each Pi0 Run getting its own dedicated policy-server task), plus a CPU-only aggregation group that stages every runner’s workflow-local output and builds one combined index.html. On the Experiment Runner side it adds a mutually-exclusive --experiment_output_directory for an exact, empty destination while preserving the timestamped --output_base_dir behavior. The change is cleanly factored, the run/aggregation boundary is well separated (OSMO staging stays in osmo/, core build_report is reused), and test coverage is genuinely thorough.

Design, Boundaries & Scope

• Base-class scaffolding is now vestigial (inline). Pi0ArenaExperimentWorkflow overrides _get_group_dicts, so task_cls_list / lead_list / the inherited _get_tasks no longer describe the real task set and survive only to pass the base constructor’s _assert_single_lead_task. Worth dropping or annotating so they aren’t read as the workflow’s actual tasks.
• Resource-model change for existing users (confirm intended). The old workflow ran one shared Pi0 server for all matching Runs; this now allocates a dedicated server GPU per Pi0 Run and runs the runner groups in parallel. That is the explicit point of the PR (independent scheduling requires per-group servers), but it multiplies GPU footprint for anyone already using this workflow — flagging it as a deliberate blast-radius change rather than a bug.

Findings

🔵 Improvement: osmo/workflows/arena_experiment_workflow.py:24 — task_cls_list / lead_list no longer match the fanned-out task set; drop or annotate them.

Test Coverage

Strong. New test_osmo_aggregate_experiment_results.py covers the loader, the missing-Run assertion, and end-to-end combined-report assembly (asserting the copied per-Run files and 2 job(s) in the report); it is pure-Python (Phase 1) and correctly needs no simulation app since build_report pulls in no sim-dependent modules. test_osmo_experiment_workflow.py is substantially expanded to assert the per-Run group/task naming, uniqueness, dedicated-server wiring, aggregation inputs/outputs, and the exact runner command. The runner-side flag is covered in test_experiment_runner.py (mutual exclusivity, exact-vs-timestamped parsing, produced files) and the reject paths (multi-chunk + non-empty exact dir) fire before sim start. No gaps spotted.

Verdict

Ship it — the one design note (vestigial task_cls_list / lead_list) is an optional cleanup, and the resource-model change is worth an explicit confirm but appears intended.

[greptile-apps]

Greptile Summary

This PR distributes Arena Experiment Runs across independently scheduled OSMO groups: each Run gets its own OSMO group (plus a dedicated Pi0 policy-server task for Pi0 Runs), and a final CPU-only aggregation group copies per-run outputs and generates one combined index.html. A new --experiment_output_directory CLI option lets managed tooling supply an exact, empty output path instead of the default timestamped subdirectory.

• Pi0ArenaExperimentWorkflow replaces the single-group workflow with N run groups + 1 aggregation group, embedding a serialized single-Run ArenaExperimentCfg per runner task and connecting each Pi0 runner to its own dedicated policy-server task.
• ExperimentResultsTask writes an OSMO {{input:task-name}}→run-name mapping into a JSON file that aggregate_experiment_results.py uses to copy per-run directories into the combined output and regenerate index.html.
• experiment_runner.py gains --experiment_output_directory as a mutually exclusive alternative to --output_base_dir, validated before SimulationApp starts to fail fast on non-empty directories.

Confidence Score: 3/5

The run-distribution logic is well-tested and the CLI change is backward-compatible, but the aggregation task embeds OSMO template tokens inside a JSON file that the Python script reads as paths — if OSMO does not resolve tokens in arbitrary file contents, the aggregation step will always fail at runtime.

The per-run group fanout, Pi0 server pairing, and --experiment_output_directory plumbing are all correct and thoroughly unit-tested. The aggregation design in ExperimentResultsTask introduces an unverified assumption: {{input:...}} OSMO tokens are written into a JSON file via files[].contents, and aggregate_experiment_results.py reads those values as filesystem paths. Whether OSMO resolves template tokens in all files[].contents entries — not only the recognized entry script — is critical to the design but is not documented or tested end-to-end. If the assumption does not hold, every distributed experiment run will fail at the aggregation step.

osmo/tasks/experiment_results_task.py — specifically the _get_files_to_create method that embeds {{input:...}} tokens into the JSON mapping file.

Important Files Changed

Filename	Overview
osmo/tasks/experiment_results_task.py	New task that embeds {{input:task-name}} OSMO tokens inside a JSON file content block; relies on OSMO resolving template tokens in arbitrary files[*].contents, not just the entry script.
osmo/workflows/arena_experiment_workflow.py	New Pi0ArenaExperimentWorkflow that distributes each Run into its own OSMO group with a dedicated policy server, plus an aggregation group. Logic is well-structured; the OSMO token design in ExperimentResultsTask is the main open question.
osmo/tasks/experiment_runner_task.py	Switched to --experiment_output_directory {{output}} and published_output_url=None for workflow-local staging; unconditionally passes --enable_cameras regardless of experiment camera requirements.
osmo/scripts/aggregate_experiment_results.py	New script to copy per-run staged directories and build a combined report; logic is correct, assertions guard against missing run directories.
isaaclab_arena/evaluation/experiment_runner.py	Added --experiment_output_directory support with pre-flight empty-dir check before SimulationApp starts; mutual exclusion with --output_base_dir correctly enforced via argparse group.
isaaclab_arena/evaluation/experiment_runner_cli.py	Added --experiment_output_directory as a mutually exclusive option; argparse group correctly prevents simultaneous use with --output_base_dir.
isaaclab_arena/tests/test_osmo_experiment_workflow.py	Comprehensive new tests covering group structure, task naming, embedded experiment YAML, policy-server configuration, and resource assignment. All assertions verified against the experiment YAML.
isaaclab_arena/tests/test_osmo_aggregate_experiment_results.py	New tests covering directory loading, missing-run detection, and end-to-end multi-run output assembly including report generation.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant S as Submit Script
    participant W as Pi0ArenaExperimentWorkflow
    participant RG0 as OSMO Group: arena-run-0
    participant RG1 as OSMO Group: arena-run-1
    participant RGN as OSMO Group: arena-run-N (local policy)
    participant AG as OSMO Group: arena-aggregation

    S->>W: submit_workflow()
    W->>W: _get_group_dicts() — one group per Run
    W->>RG0: "group {experiment-runner-0 (lead), policy-server-0}"
    W->>RG1: "group {experiment-runner-1 (lead), policy-server-1}"
    W->>RGN: "group {experiment-runner-N (lead, no server)}"
    W->>AG: "group {aggregate-results (lead, resource=aggregation, gpu=0)}"

    Note over RG0: experiment-runner-0 writes
    Note over RG1: experiment-runner-1 writes
    Note over RGN: experiment-runner-N writes

    RG0-->>AG: OSMO stages input:experiment-runner-0
    RG1-->>AG: OSMO stages input:experiment-runner-1
    RGN-->>AG: OSMO stages input:experiment-runner-N

    AG->>AG: aggregate_experiment_results.py copies run dirs
    AG->>AG: build_report → index.html
    AG->>AG: publish output to Swift

Loading

%%{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 Script
    participant W as Pi0ArenaExperimentWorkflow
    participant RG0 as OSMO Group: arena-run-0
    participant RG1 as OSMO Group: arena-run-1
    participant RGN as OSMO Group: arena-run-N (local policy)
    participant AG as OSMO Group: arena-aggregation

    S->>W: submit_workflow()
    W->>W: _get_group_dicts() — one group per Run
    W->>RG0: "group {experiment-runner-0 (lead), policy-server-0}"
    W->>RG1: "group {experiment-runner-1 (lead), policy-server-1}"
    W->>RGN: "group {experiment-runner-N (lead, no server)}"
    W->>AG: "group {aggregate-results (lead, resource=aggregation, gpu=0)}"

    Note over RG0: experiment-runner-0 writes
    Note over RG1: experiment-runner-1 writes
    Note over RGN: experiment-runner-N writes

    RG0-->>AG: OSMO stages input:experiment-runner-0
    RG1-->>AG: OSMO stages input:experiment-runner-1
    RGN-->>AG: OSMO stages input:experiment-runner-N

    AG->>AG: aggregate_experiment_results.py copies run dirs
    AG->>AG: build_report → index.html
    AG->>AG: publish output to Swift

Loading

Comments Outside Diff (1)

1. osmo/tasks/experiment_runner_task.py, line 75-87 (link)

   --enable_cameras unconditionally passed

   The run script always appends --enable_cameras, even for single-Run experiments whose embedded ArenaExperimentCfg has no camera requirements. This causes AppLauncher to initialise the full camera subsystem on every OSMO runner task, adding startup overhead for non-camera Runs. The ExperimentRunnerTask already has access to self.experiment_cfg and could conditionally include the flag only when any Run's environment has enable_cameras=True (mirroring the existing _assert_camera_support_enabled check in experiment_runner.py).

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "Clarify exact output validation" | Re-trigger Greptile}

[alexmillane]

LGTM