Add YAML-env and variation support to PolicyRunnerTask

[alexmillane]

Summary

Add graph-spec YAML environment and variation support OSMO workflows.

Detailed description

• Split the arena_env (name) from the other args
• Add the appropriate flag if it is a yaml description of the env.
• Adds --env_variations for appending variations.

Not done

• Not too much thought given to the flags as we plan to completely get rid of the flag based system in a week or so.

[greptile-apps]

Greptile Summary

This PR splits the old combined --arena_env_args argument (which previously accepted "env_name [env_args...]" as a single string) into three separate arguments — --arena_env (required), --arena_env_args (optional trailing flags), and --variations (Hydra-style overrides) — and adds get_arena_env_token() to transparently map YAML paths to --env_graph_spec_yaml flags vs. passing registered env names verbatim.

• policy_runner_task.py: New --arena_env is required=True; get_arena_env_token() detects .yaml/.yml suffixes and injects the appropriate flag; _get_policy_runner_command() replaces the old inline f-string with a filter-and-join parts list.
• submit_gr00t_policy_runner_workflow.py: Usage examples correctly updated to the new interface.
• submit_zero_action_policy_runner_workflow.py / submit_pi0_plus_policy_runner_workflow.py: Not updated — their default examples omit the now-required --arena_env, and their custom examples still pass the env name bundled inside --arena_env_args, causing argparse failures at runtime.

Confidence Score: 4/5

The core logic in policy_runner_task.py is correct, but the zero-action and pi0 submit scripts were not updated and their invocations will fail at argparse.

The refactoring of PolicyRunnerTask is clean and the GR00T submit script is properly updated. However, submit_zero_action_policy_runner_workflow.py and submit_pi0_plus_policy_runner_workflow.py still carry the old --arena_env_args-only pattern: the default examples omit the now-required --arena_env entirely, and the custom examples bundle the env name inside --arena_env_args. Anyone following those examples will get an immediate argparse failure.

osmo/submit_zero_action_policy_runner_workflow.py and osmo/submit_pi0_plus_policy_runner_workflow.py need their usage examples updated to match the new --arena_env / --arena_env_args interface.

Important Files Changed

Filename	Overview
osmo/tasks/policy_runner_task.py	Splits --arena_env_args into --arena_env (required) + --arena_env_args + --variations; adds get_arena_env_token() for YAML/name detection. Minor docstring typo (missing opening backtick).
osmo/submit_gr00t_policy_runner_workflow.py	Usage examples correctly updated to use the new --arena_env / --arena_env_args split interface.
osmo/submit_zero_action_policy_runner_workflow.py	Not updated: default example omits the now-required --arena_env, and the custom example still uses the old bundled --arena_env_args format — both invocations fail at runtime.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["CLI: --arena_env <value>"] --> B{"value.endswith('.yaml' | '.yml')?"}
    B -- Yes --> C["get_arena_env_token()\n→ '--env_graph_spec_yaml <value>'"]
    B -- No --> D["get_arena_env_token()\n→ '<value>' (registered env name)"]
    C --> E["_get_policy_runner_command()"]
    D --> E
    F["--arena_env_args"] --> E
    G["--variations"] --> E
    H["--policy_runner_args"] --> E
    I["_get_policy_args() (abstract)"] --> E
    E --> J["policy_runner.py <args>\n  --output_base_dir ...\n  [policy args]\n  [env token]\n  [env args]\n  [variations]"]

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"}}}%%
flowchart TD
    A["CLI: --arena_env <value>"] --> B{"value.endswith('.yaml' | '.yml')?"}
    B -- Yes --> C["get_arena_env_token()\n→ '--env_graph_spec_yaml <value>'"]
    B -- No --> D["get_arena_env_token()\n→ '<value>' (registered env name)"]
    C --> E["_get_policy_runner_command()"]
    D --> E
    F["--arena_env_args"] --> E
    G["--variations"] --> E
    H["--policy_runner_args"] --> E
    I["_get_policy_args() (abstract)"] --> E
    E --> J["policy_runner.py <args>\n  --output_base_dir ...\n  [policy args]\n  [env token]\n  [env args]\n  [variations]"]

Loading

_{Reviews (4): Last reviewed commit: "Self review." | Re-trigger Greptile}

[arena-review-bot]

🤖 Isaac Lab-Arena Review Bot

Summary

Adds --env_variations (Hydra-style overrides) and YAML-graph-spec resolution to the base PolicyRunnerTask, so all policy-runner OSMO tasks inherit both. The mechanics line up with the core: --env_graph_spec_yaml is a real flag in isaaclab_arena/cli, and the appended key=value overrides are collected as hydra_overrides via parse_known_args. The change is opt-in (--env_variations defaults to ""), so the existing default path is unchanged. Clean and correct — just one small simplification and a note on test coverage.

Findings

🔵 Improvement: osmo/tasks/policy_runner_task.py:44 — the if env_variations else "" guard is redundant and inconsistent with the two lines above it (which call _normalize_args directly). --env_variations defaults to "", so the value is always a str, and _normalize_args("") already returns "". See inline suggestion.

Test Coverage

_resolve_env_source and _get_env_spec_args are pure string logic (no sim, no live env) with real branching — YAML vs. registered-name, trailing args present/absent, variations present/absent. They're cheap to unit-test, but there's currently no test harness under osmo/. Worth a small pytest covering the three branches (YAML path with/without trailing args, plain env name, and each with/without --env_variations) so a future refactor of the .yaml/.yml heuristic or the arg ordering can't silently break the generated command.

Verdict

Ship it — non-blocking suggestions only.

[alexmillane]

Self review.

[arena-review-bot]

🤖 Isaac Lab-Arena Review Bot

Summary

This cleanly splits the env identifier into a new required --arena_env flag and adds --env_variations for Hydra-style overrides, both on the base PolicyRunnerTask so every policy-runner subclass inherits them. The rendering logic is correct — the empty-segment join avoids stray spaces, and the appended variations flow through policy_runner.py's parse_known_args into hydra_overrides as intended. One follow-up before merge: the sibling submit scripts weren't updated for the new flag.

Findings

🟡 Warning: osmo/submit_pi0_plus_policy_runner_workflow.py & osmo/submit_zero_action_policy_runner_workflow.py (usage docstrings) — Both tasks inherit PolicyRunnerTask.add_task_arguments unchanged, so --arena_env is now required=True for them too, and --arena_env_args no longer carries the env name. Their usage examples still pass the name inside --arena_env_args (e.g. --arena_env_args 'kitchen_pick_and_place --object cracker_box --embodiment franka_ik') and omit --arena_env entirely — so those commands now fail argparse (--arena_env is required), and even with it added the name would be mis-routed into the env args. The gr00t submit script was updated to the new form; could the pi0 and zero_action docstrings get the same --arena_env <name> split so the examples still run?

🔵 Improvement: CLI contract change — Making --arena_env required (and shrinking --arena_env_args to env args only) is a good, more-explicit split, but it breaks any existing OSMO invocation using the old --arena_env_args 'name ...' form. Worth a heads-up to downstream users of these submit scripts; the split itself reads well.

Test Coverage

osmo/ has no test harness today, so the rendered run-script isn't unit-tested — that's the pre-existing state, not introduced here. If a lightweight check ever lands for these tasks, _get_run_script/_get_arena_env (YAML-path vs registered-name, empty-segment skipping) would be the natural things to assert on.

Verdict

Minor fixes needed

[alexmillane]

Self-review

[qianl-nv]

LGTM, thanks!