Distribute Arena Experiment Runs on OSMO

[cvolkcvolk]

Summary

Run Arena Experiment Runs in parallel as independent OSMO jobs and combine their outputs.

Detailed description

• Create one independently scheduled OSMO group per Run. A Pi0 Run gets its own policy-server task; a local-policy Run executes by itself.
• Keep Arena’s normal timestamped layout inside every runner task: {{output}}/<timestamp>/<run-name>/....
• After all runners finish, use a CPU task to find each exact Run directory, copy it into the final Experiment output, and generate the combined index.html report.
• Let OSMO fail the workflow when a Run or final output task fails; retries and partial-result handling remain out of scope.
• Build directly on Run OpenPi ArenaExperiments on OSMO sequentially #897. Each OSMO task handles one typed Run and does not use legacy --chunk_size execution.

[arena-review-bot]

🔵 lead_list / task_cls_list no longer drive rendering

After the fan-out, _get_group_dicts builds every task directly with hard-coded lead= flags, so lead_list and task_cls_list no longer shape what this workflow renders — they only satisfy the base __init__ asserts (_assert_single_lead_task, len(task_cls_list) > 0), which still describe a single 2-task group rather than the N groups actually emitted. Are they still meaningful here, or is there a cleaner way to satisfy the base contract without leaving these two attrs describing a shape the workflow no longer produces?

[arena-review-bot]

🤖 Isaac Lab-Arena Review Bot

Summary

This PR fans a single-group Experiment workflow out into one OSMO group per Run (each Pi0 Run gets a dedicated server), then adds an aggregate-results CPU task that copies every runner output into one directory and rebuilds a combined report. The seam rename (build_and_run → execute_run), the task_name/resource/output_url plumbing on the tasks, and the new aggregation path are all clean, backward-compatible, and well-covered by the updated tests. Overall this reads as a solid, well-scoped change.

I verified the likely risk areas and they hold up: build_report only pulls stdlib + a pure filename parser, so running it on the gpu: 0 aggregation resource is fine; only the aggregate task publishes to DATASET_SWIFT_URL (runners pass output_url=None), so there is no output collision; and failed Runs fail the workflow through OSMO input dependencies. Copyright years are correct.

Findings

🔵 Improvement: osmo/workflows/arena_experiment_workflow.py:31 — After the refactor, _get_group_dicts builds tasks directly with hard-coded lead= flags, so lead_list/task_cls_list no longer describe the rendered workflow and only feed the base __init__ asserts. Worth deciding whether they still earn their place here.

Test Coverage

Good. aggregate_experiment_outputs gets a dedicated test for the happy path plus the missing/ambiguous-output asserts; the workflow test now checks the per-Run group/task/host-token wiring, the aggregate inputs/outputs, and the CPU-only aggregation resource; the sim-touching execute_run test keeps the inner/outer pattern. No gaps stood out.

Verdict

Ship it

[greptile-apps]

Greptile Summary

This PR distributes each Arena Experiment Run into its own independent OSMO group, replacing the previous single-group design where all Runs shared one Experiment Runner and one Pi0 policy server. Each Run now gets a singleton ExperimentRunnerTask (with output_url=None) and, when using Pi0RemotePolicyCfg, a dedicated Pi0ServerTask; a final ExperimentResultsTask group collects all runner outputs via OSMO {{input:…}} tokens and publishes a combined report to Swift.

• Workflow fan-out: Pi0ArenaExperimentWorkflow._get_group_dicts() replaces _get_tasks(), emitting one arena-run-{i} group per Run followed by an aggregate-results group whose task inputs list every runner, letting OSMO enforce ordering and propagate failures.
• Aggregation script: aggregate_experiment_outputs.py locates each run's output directory via an episode_results_rebuild0.jsonl sentinel, copies it under the combined output tree, and calls build_report().
• Rename: build_and_run → execute_run clarifies that the caller owns the SimulationApp lifecycle.

Confidence Score: 4/5

Safe to merge; the distributed fan-out design is well-tested and the changes are structurally sound.

The group_name constructor parameter on Pi0ArenaExperimentWorkflow is accepted and stored but never read by the overriding _get_group_dicts(), making it a silent no-op that could mislead future callers. Additionally, aggregate_experiment_outputs.py contains a redundant mkdir call on a directory that was already created moments earlier. Neither issue affects correctness in the current codebase, but the misleading API surface on the workflow class is worth addressing before this pattern gets copied.

osmo/workflows/arena_experiment_workflow.py (group_name parameter) and isaaclab_arena/evaluation/aggregate_experiment_outputs.py (redundant mkdir)

Important Files Changed

Filename	Overview
osmo/workflows/arena_experiment_workflow.py	Core refactor: replaces single-group workflow with one group per Run + aggregation group; group_name constructor parameter is stored but silently ignored by the overridden _get_group_dicts().
osmo/tasks/experiment_results_task.py	New CPU-only aggregation task; correctly declares inter-group task inputs, embeds OSMO token-resolved run-input JSON, and generates entry script with set -euo pipefail.
isaaclab_arena/evaluation/aggregate_experiment_outputs.py	New aggregation script; correctly uses episode_results_rebuild0.jsonl as a sentinel to locate and copy each Run output dir, but contains a redundant mkdir call on a directory already created earlier in the same function.
osmo/tasks/base_task.py	Adds optional task_name and resource constructor parameters; host_token() now accepts an explicit task name override; backward-compatible change.
osmo/workflows/workflow.py	Extracts _get_group_dicts() and _create_resources_dict() from create_workflow_dict() so subclasses can override them independently; no logic changes to base behavior.
osmo/tasks/experiment_runner_task.py	Adds task_name and output_url constructor parameters; output_url=None suppresses the Swift URL output so runners upload nothing (aggregator handles the final upload).
isaaclab_arena/evaluation/run_execution.py	Renames build_and_run to execute_run and improves docstring to clarify SimulationApp ownership; pure rename with no logic changes.
isaaclab_arena/tests/test_osmo_experiment_workflow.py	Tests updated to cover per-Run groups, per-Run Pi0 servers, aggregation task structure, and resource overrides; comprehensive coverage of the new distributed design.
isaaclab_arena/tests/test_aggregate_experiment_outputs.py	New tests cover the happy path (two runs, one report) and error cases (missing / duplicate run output sentinel); good coverage.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller as Caller / CLI
    participant W as Pi0ArenaExperimentWorkflow
    participant G0 as arena-run-0 Group
    participant Gn as arena-run-N Group
    participant Agg as aggregate-results Group

    Caller->>W: Pi0ArenaExperimentWorkflow(experiment_cfg)
    W->>W: deepcopy(experiment_cfg)
    W->>W: _assert_pi0_server_compatible()
    Caller->>W: generate_workflow()
    W->>W: _get_group_dicts()

    loop For each Run
        W->>W: deepcopy(run_cfg) to singleton_experiment
        alt Pi0RemotePolicy
            W->>G0: ExperimentRunnerTask + Pi0ServerTask
        else Local policy
            W->>Gn: ExperimentRunnerTask only
        end
    end

    W->>Agg: ExperimentResultsTask(run_task_names)
    W-->>Caller: workflow dict

    Note over G0,Gn: OSMO runs groups in parallel
    G0->>Agg: task output dir
    Gn->>Agg: task output dir
    Agg->>Agg: aggregate_experiment_outputs.py
    Agg-->>Caller: combined report at DATASET_SWIFT_URL

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 Caller as Caller / CLI
    participant W as Pi0ArenaExperimentWorkflow
    participant G0 as arena-run-0 Group
    participant Gn as arena-run-N Group
    participant Agg as aggregate-results Group

    Caller->>W: Pi0ArenaExperimentWorkflow(experiment_cfg)
    W->>W: deepcopy(experiment_cfg)
    W->>W: _assert_pi0_server_compatible()
    Caller->>W: generate_workflow()
    W->>W: _get_group_dicts()

    loop For each Run
        W->>W: deepcopy(run_cfg) to singleton_experiment
        alt Pi0RemotePolicy
            W->>G0: ExperimentRunnerTask + Pi0ServerTask
        else Local policy
            W->>Gn: ExperimentRunnerTask only
        end
    end

    W->>Agg: ExperimentResultsTask(run_task_names)
    W-->>Caller: workflow dict

    Note over G0,Gn: OSMO runs groups in parallel
    G0->>Agg: task output dir
    Gn->>Agg: task output dir
    Agg->>Agg: aggregate_experiment_outputs.py
    Agg-->>Caller: combined report at DATASET_SWIFT_URL

Loading

_{Reviews (1): Last reviewed commit: "Distribute Arena Runs across OSMO groups" | Re-trigger Greptile}

[greptile-apps]

The run_output_dir.parent is always output_dir, which was already created unconditionally with mkdir(parents=True, exist_ok=True) three lines earlier. This call is a no-op in every execution path.

Suggested change

	run_output_dir = output_dir / run_name
	assert not run_output_dir.exists(), f"Run output destination already exists: '{run_output_dir}'"
	run_output_dir.parent.mkdir(parents=True, exist_ok=True)
	shutil.copytree(result_paths[0].parent, run_output_dir)
	run_output_dir = output_dir / run_name
	assert not run_output_dir.exists(), f"Run output destination already exists: '{run_output_dir}'"
	shutil.copytree(result_paths[0].parent, run_output_dir)

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!

[greptile-apps]

group_name parameter is silently ignored

Pi0ArenaExperimentWorkflow.__init__ accepts and forwards group_name to Workflow.__init__, where it is stored as self.group_name. However, the overridden _get_group_dicts() never reads self.group_name — it unconditionally emits "arena-run-{index}" and "aggregate-results". A caller passing group_name="my-exp" gets groups named arena-run-0, arena-run-1, etc. with no indication that the parameter was ignored. Consider either removing the parameter from this subclass's signature or using self.group_name as a prefix in the generated names.