diff --git a/isaaclab_arena/tests/test_osmo_experiment_workflow.py b/isaaclab_arena/tests/test_osmo_experiment_workflow.py index d4ca9db43..f5c106fb1 100644 --- a/isaaclab_arena/tests/test_osmo_experiment_workflow.py +++ b/isaaclab_arena/tests/test_osmo_experiment_workflow.py @@ -19,25 +19,25 @@ from isaaclab_arena.evaluation.legacy_graph_environment_cli import LegacyGraphEnvironmentCfg from isaaclab_arena.policy.zero_action_policy import ZeroActionPolicyCfg from isaaclab_arena_environments.pick_and_place_maple_table_environment import PickAndPlaceMapleTableEnvironmentCfg +from isaaclab_arena_gr00t.policy.gr00t_remote_closedloop_policy import Gr00tRemoteClosedloopPolicyCfg from isaaclab_arena_openpi.policy import pi0_remote_policy # noqa: F401 from isaaclab_arena_openpi.policy.pi0_remote_config import Pi0RemotePolicyCfg from osmo.submit_arena_experiment import ( - POLICY_SERVER_TASK_CFG_BY_NAME, - POLICY_SERVER_WORKFLOW_BY_CONFIG_TYPE, ArenaExperimentSubmissionCfg, build_arena_experiment_submission_cfg, main, + submission_cfg_to_str, submit_arena_experiment, ) -from osmo.tasks.base_task import TaskCfg from osmo.tasks.collect_experiment_outputs_task import ( _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH, _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH, experiment_runner_output_directory_input_token, ) from osmo.tasks.experiment_runner_task import REMOTE_EXPERIMENT_PATH, ExperimentRunnerTask, ExperimentRunnerTaskCfg +from osmo.tasks.gr00t_server_task import Gr00tServerTask, Gr00tServerTaskCfg from osmo.tasks.pi0_server_task import Pi0ServerTask, Pi0ServerTaskCfg -from osmo.workflows.arena_experiment_workflow import Pi0ArenaExperimentWorkflow +from osmo.workflows.arena_experiment_workflow import ArenaExperimentWorkflow from osmo.workflows.workflow import WorkflowCfg from osmo.workflows.workflow_constants import DATASET_SWIFT_URL, OSMO_TASK_OUTPUT_DIR, POLICY_SERVER_PORT @@ -50,6 +50,7 @@ REPOSITORY_ROOT / "isaaclab_arena_environments/experiment_configs/droid_pnp_srl_openpi_experiment.yaml" ) OPENPI_RUN_NAME = "droid_pnp_srl_openpi_billiard_hall" +GR00T_CONFIG_YAML_PATH = "isaaclab_arena_gr00t/policy/config/droid_manip_gr00t_closedloop_config.yaml" def _pi0_experiment_cfg(first_variant: str = "pi05") -> ArenaExperimentCfg: @@ -79,6 +80,28 @@ def _pi0_experiment_cfg(first_variant: str = "pi05") -> ArenaExperimentCfg: ) +def _mixed_experiment_cfg() -> ArenaExperimentCfg: + return ArenaExperimentCfg( + runs={ + "pi0_run": ArenaRunCfg( + name="pi0_run", + environment=PickAndPlaceMapleTableEnvironmentCfg(), + policy=Pi0RemotePolicyCfg(), + ), + "gr00t_run": ArenaRunCfg( + name="gr00t_run", + environment=PickAndPlaceMapleTableEnvironmentCfg(), + policy=Gr00tRemoteClosedloopPolicyCfg(policy_config_yaml_path=GR00T_CONFIG_YAML_PATH), + ), + "local": ArenaRunCfg( + name="local", + environment=PickAndPlaceMapleTableEnvironmentCfg(), + policy=ZeroActionPolicyCfg(), + ), + } + ) + + def _zero_action_experiment_cfg() -> ArenaExperimentCfg: return ArenaExperimentCfg( runs={ @@ -116,33 +139,19 @@ def _workflow_tasks(workflow: dict, group_index: int = 0) -> list[dict]: def _compose_submission( overrides: list[str] | None = None, experiment_cfg_path: Path = OPENPI_EXPERIMENT_CFG_PATH, - policy_server_name: str = "pi0", ) -> ArenaExperimentSubmissionCfg: - return build_arena_experiment_submission_cfg( - experiment_cfg_path, - policy_server_name, - overrides, - ) + return build_arena_experiment_submission_cfg(experiment_cfg_path, overrides) def _compose_and_submit( overrides: list[str], experiment_cfg_path: Path = OPENPI_EXPERIMENT_CFG_PATH, - policy_server_name: str = "pi0", ) -> int: - return submit_arena_experiment(_compose_submission(overrides, experiment_cfg_path, policy_server_name)) - + return submit_arena_experiment(_compose_submission(overrides, experiment_cfg_path)) -def test_declares_policy_server_name_and_workflow_mappings(): - """Keep policy-server selection and runtime workflow dispatch explicit.""" - assert POLICY_SERVER_TASK_CFG_BY_NAME == {"pi0": Pi0ServerTaskCfg} - assert POLICY_SERVER_WORKFLOW_BY_CONFIG_TYPE == {Pi0ServerTaskCfg: Pi0ArenaExperimentWorkflow} - assert Pi0ArenaExperimentWorkflow.task_cfg_type is ExperimentRunnerTaskCfg - assert Pi0ArenaExperimentWorkflow.server_task_cfg_type is Pi0ServerTaskCfg - -def test_explicit_experiment_and_policy_server_selector_compose_typed_defaults(): - """Compose an explicit Experiment path and built-in policy-server defaults.""" +def test_explicit_experiment_composes_typed_defaults(): + """Compose an explicit Experiment path with derived OSMO defaults.""" submission_cfg = _compose_submission() assert isinstance(submission_cfg.experiment_cfg, ArenaExperimentCfg) @@ -153,51 +162,21 @@ def test_explicit_experiment_and_policy_server_selector_compose_typed_defaults() assert submission_cfg.osmo.platform == "ovx-l40s" assert submission_cfg.experiment_runner == ExperimentRunnerTaskCfg() assert submission_cfg.experiment_runner.image == "nvcr.io/nvstaging/isaac-amr/isaaclab_arena:latest" - assert submission_cfg.policy_server == Pi0ServerTaskCfg() - assert submission_cfg.policy_server.client_ping_timeout_s == Pi0ServerTaskCfg.client_ping_timeout_s - - with pytest.raises(AssertionError, match="Unknown policy server 'unknown'.*pi0"): - _compose_submission(policy_server_name="unknown") - - with pytest.raises(AssertionError, match="policy_variant must be one of"): - _compose_submission(["policy_server.policy_variant=unknown"]) -def test_submitter_rejects_unregistered_policy_server_type(): - """Reject a typed policy server without a registered workflow implementation.""" - submission_cfg = ArenaExperimentSubmissionCfg( - experiment_cfg=_pi0_experiment_cfg(), - osmo=WorkflowCfg(dry_run=True), - policy_server=TaskCfg(), - ) - - with pytest.raises(AssertionError, match="No policy-server workflow.*TaskCfg"): - submit_arena_experiment(submission_cfg) - - -@pytest.mark.parametrize("config_path", ["osmo.not_a_field", "experiment_runner.not_a_field"]) +@pytest.mark.parametrize("config_path", ["osmo.not_a_field", "experiment_runner.not_a_field", "servers"]) def test_hydra_rejects_unknown_typed_config_fields(config_path): """Let the structured Hydra root reject fields outside their owning config.""" - with pytest.raises(ConfigCompositionException, match="not_a_field"): + with pytest.raises(ConfigCompositionException, match=config_path.split(".", maxsplit=1)[-1]): _compose_submission([f"{config_path}=true"]) -def test_policy_server_rejects_workflow_fields(): - """Validate server overrides against the selected task config type.""" - with pytest.raises(ConfigCompositionException, match="workflow_name"): - _compose_and_submit([ - "policy_server.workflow_name=experiment", - "osmo.dry_run=true", - ]) - - def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_experiment_output(): """Render one independent Run group per Run and collect their outputs into one Experiment output.""" source_experiment_cfg = _pi0_experiment_cfg() - workflow = Pi0ArenaExperimentWorkflow( + workflow = ArenaExperimentWorkflow( workflow_cfg=WorkflowCfg(workflow_name="pi0-experiment"), experiment_cfg=source_experiment_cfg, - server_task_cfg=Pi0ServerTaskCfg(), ) rendered_workflow = workflow.generate_workflow() @@ -237,9 +216,10 @@ def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_expe assert first_experiment["runs"]["first"]["policy"]["remote_host"] == Pi0ServerTask.host_token("policy-server-0") assert second_experiment["runs"]["second"]["policy"]["remote_host"] == Pi0ServerTask.host_token("policy-server-1") for run_name, experiment in (("first", first_experiment), ("second", second_experiment)): - policy = experiment["runs"][run_name]["policy"] - assert policy["remote_port"] == POLICY_SERVER_PORT - assert policy["ping_timeout"] == Pi0ServerTaskCfg.client_ping_timeout_s + assert experiment["runs"][run_name]["policy"]["remote_port"] == POLICY_SERVER_PORT + # Only the connection is rewritten: each Run keeps the ping timeout its policy asked for. + assert first_experiment["runs"]["first"]["policy"]["ping_timeout"] == 10 + assert second_experiment["runs"]["second"]["policy"]["ping_timeout"] == Pi0RemotePolicyCfg.ping_timeout assert "remote_host" not in local_experiment["runs"]["local"]["policy"] assert "remote_port" not in local_experiment["runs"]["local"]["policy"] assert source_experiment_cfg.runs["first"].policy.remote_host == "user-host" @@ -294,6 +274,79 @@ def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_expe assert rendered_workflow["workflow"]["resources"]["experiment-output"]["gpu"] == 0 +def test_mixed_pi0_variants_derive_per_run_server_checkpoints(): + """Each Run's dedicated pi0 server follows that Run's client policy_variant.""" + workflow = ArenaExperimentWorkflow( + workflow_cfg=WorkflowCfg(workflow_name="mixed-pi0-variants"), + experiment_cfg=_pi0_experiment_cfg(first_variant="pi0"), + ) + + groups = _workflow_groups(workflow.generate_workflow()) + first_server_command = _task_file(groups[0]["tasks"][1], "/tmp/entry.sh")["contents"] + second_server_command = _task_file(groups[1]["tasks"][1], "/tmp/entry.sh")["contents"] + assert "--policy.config=pi0_droid_jointpos_polaris" in first_server_command + assert "--policy.dir=gs://openpi-assets-simeval/pi0_droid_jointpos" in first_server_command + assert "--policy.config=pi05_droid_jointpos_polaris" in second_server_command + assert "--policy.dir=gs://openpi-assets-simeval/pi05_droid_jointpos" in second_server_command + + +def test_mixed_pi0_and_gr00t_experiment_fans_out_per_run_servers(): + """Derive one pi0 server and one GR00T server from a mixed Experiment; leave the local Run server-free.""" + workflow = ArenaExperimentWorkflow( + workflow_cfg=WorkflowCfg(workflow_name="mixed-experiment"), + experiment_cfg=_mixed_experiment_cfg(), + ) + + groups = _workflow_groups(workflow.generate_workflow()) + assert [group["name"] for group in groups] == [ + "arena-run-0", + "arena-run-1", + "arena-run-2", + "arena-experiment-output", + ] + pi0_tasks = groups[0]["tasks"] + gr00t_tasks = groups[1]["tasks"] + local_tasks = groups[2]["tasks"] + assert [task["name"] for task in pi0_tasks] == ["experiment-runner-0", "policy-server-0"] + assert [task["name"] for task in gr00t_tasks] == ["experiment-runner-1", "policy-server-1"] + assert [task["name"] for task in local_tasks] == ["experiment-runner-2"] + assert pi0_tasks[1]["image"] == Pi0ServerTaskCfg().image + assert gr00t_tasks[1]["image"] == Gr00tServerTaskCfg().image + + pi0_policy = _embedded_experiment(pi0_tasks[0])["runs"]["pi0_run"]["policy"] + assert pi0_policy["remote_host"] == Pi0ServerTask.host_token("policy-server-0") + assert pi0_policy["remote_port"] == POLICY_SERVER_PORT + gr00t_policy = _embedded_experiment(gr00t_tasks[0])["runs"]["gr00t_run"]["policy"] + assert gr00t_policy["remote_host"] == Gr00tServerTask.host_token("policy-server-1") + assert gr00t_policy["remote_port"] == POLICY_SERVER_PORT + assert "remote_host" not in _embedded_experiment(local_tasks[0])["runs"]["local"]["policy"] + + +def test_gr00t_cli_dry_run_renders_workflow(tmp_path, capsys): + """Compose and render a derived GR00T Experiment submission through the real CLI parser.""" + experiment_path = tmp_path / "gr00t_experiment.yaml" + experiment_path.write_text( + f"""runs: + gr00t_run: + environment: + type: pick_and_place_maple_table + policy: + type: isaaclab_arena_gr00t.policy.gr00t_remote_closedloop_policy.Gr00tRemoteClosedloopPolicy + policy_config_yaml_path: {GR00T_CONFIG_YAML_PATH} +""", + encoding="utf-8", + ) + + assert main(["--experiment_cfg", str(experiment_path), "--dry_run"]) == 0 + workflow = _rendered_workflow(capsys.readouterr().out) + tasks = _workflow_tasks(workflow) + assert [task["name"] for task in tasks] == ["experiment-runner-0", "policy-server-0"] + assert tasks[1]["image"] == Gr00tServerTaskCfg().image + served_policy = _embedded_experiment(tasks[0])["runs"]["gr00t_run"]["policy"] + assert served_policy["remote_host"] == Gr00tServerTask.host_token("policy-server-0") + assert served_policy["remote_port"] == POLICY_SERVER_PORT + + def test_embeds_effective_experiment_yaml(): """Embed the composed Experiment instead of staging its source file.""" experiment_runner_task = ExperimentRunnerTask( @@ -311,13 +364,24 @@ def test_embeds_effective_experiment_yaml(): assert embedded_experiment["runs"]["baseline"]["environment_builder"]["num_envs"] == 1 +def test_all_local_experiment_runs_standalone_without_servers(): + """An Experiment with no remote-policy Runs derives no servers and runs every Run standalone.""" + workflow = ArenaExperimentWorkflow( + workflow_cfg=WorkflowCfg(), + experiment_cfg=_zero_action_experiment_cfg(), + ) + + groups = _workflow_groups(workflow.generate_workflow()) + assert [group["name"] for group in groups] == ["arena-run-0", "arena-experiment-output"] + assert [task["name"] for task in groups[0]["tasks"]] == ["experiment-runner-0"] + + def test_submission_removes_temporary_workflow(monkeypatch): """Submit one temporary workflow and remove it afterwards.""" experiment_cfg = _pi0_experiment_cfg() - workflow = Pi0ArenaExperimentWorkflow( + workflow = ArenaExperimentWorkflow( workflow_cfg=WorkflowCfg(), experiment_cfg=experiment_cfg, - server_task_cfg=Pi0ServerTaskCfg(), ) captured_workflow_path = None @@ -360,16 +424,13 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys) "osmo.dry_run=true", "osmo.workflow_name=overridden-experiment", "experiment_runner.image=registry.example.com/evaluator:branch", - "policy_server.image=registry.example.com/openpi:overridden", - "policy_server.policy_config=overridden-pi0-config", - "policy_server.client_ping_timeout_s=600.0", "experiment_cfg.runs.openpi_maple_table.rollout_limit.num_episodes=4", "experiment_cfg.runs.openpi_maple_table.environment_builder.num_envs=2", "experiment_cfg.runs.openpi_maple_table.policy.ping_interval=33.0", "experiment_cfg.runs.openpi_maple_table.policy.ping_timeout=450.0", + "experiment_cfg.runs.openpi_maple_table.policy.policy_variant=pi0", ], experiment_path, - "pi0", ) assert return_code == 0 @@ -380,7 +441,7 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys) tasks = _workflow_tasks(workflow) assert [task["name"] for task in tasks] == ["experiment-runner-0", "policy-server-0"] assert tasks[0]["image"] == "registry.example.com/evaluator:branch" - assert tasks[1]["image"] == "registry.example.com/openpi:overridden" + assert tasks[1]["image"] == Pi0ServerTaskCfg().image experiment = _embedded_experiment(tasks[0]) policy = experiment["runs"]["openpi_maple_table"]["policy"] @@ -389,22 +450,21 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys) assert policy["ping_interval"] == 33.0 assert policy["remote_host"] == Pi0ServerTask.host_token("policy-server-0") assert policy["remote_port"] == POLICY_SERVER_PORT - assert policy["ping_timeout"] == 600.0 + assert policy["ping_timeout"] == 450.0 + assert policy["policy_variant"] == "pi0" assert "experiment_cfg.runs" not in _task_file(tasks[0], "/tmp/entry.sh")["contents"] server_command = _task_file(tasks[1], "/tmp/entry.sh")["contents"] - assert "--policy.config=overridden-pi0-config" in server_command - assert "--policy.dir=gs://openpi-assets-simeval/pi05_droid_jointpos" in server_command + assert "--policy.config=pi0_droid_jointpos_polaris" in server_command + assert "--policy.dir=gs://openpi-assets-simeval/pi0_droid_jointpos" in server_command def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tmp_path): """Keep every single-Run OSMO handoff compatible with the Experiment Runner loader.""" submission_cfg = _compose_submission() - assert isinstance(submission_cfg.policy_server, Pi0ServerTaskCfg) - workflow = Pi0ArenaExperimentWorkflow( + workflow = ArenaExperimentWorkflow( workflow_cfg=submission_cfg.osmo, experiment_cfg=submission_cfg.experiment_cfg, - server_task_cfg=submission_cfg.policy_server, task_cfg=submission_cfg.experiment_runner, ) rendered_workflow = workflow.generate_workflow() @@ -419,7 +479,7 @@ def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tm assert isinstance(run_cfg.policy, Pi0RemotePolicyCfg) assert run_cfg.policy.remote_host == Pi0ServerTask.host_token(f"policy-server-{index}") assert run_cfg.policy.remote_port == POLICY_SERVER_PORT - assert run_cfg.policy.ping_timeout == Pi0ServerTaskCfg.client_ping_timeout_s + assert run_cfg.policy.ping_timeout == Pi0RemotePolicyCfg.ping_timeout def test_embedded_graph_environment_experiment_composes_through_experiment_runner_loader(tmp_path): @@ -444,10 +504,9 @@ def test_embedded_graph_environment_experiment_composes_through_experiment_runne ["experiment_cfg.runs.graph_run.environment_builder.num_envs=4"], experiment_path, ) - workflow = Pi0ArenaExperimentWorkflow( + workflow = ArenaExperimentWorkflow( workflow_cfg=submission_cfg.osmo, experiment_cfg=submission_cfg.experiment_cfg, - server_task_cfg=submission_cfg.policy_server, task_cfg=submission_cfg.experiment_runner, ) rendered_workflow = workflow.generate_workflow() @@ -475,7 +534,7 @@ def test_embedded_graph_environment_experiment_composes_through_experiment_runne def test_submission_overrides_osmo_resources(monkeypatch): - """Apply scheduler overrides after the typed workflow defaults.""" + """Apply scheduler overrides after the derived workflow resource.""" submitted_command = None submitted_resources = None @@ -506,19 +565,15 @@ def capture_submission(command, **kwargs): assert submitted_resources["experiment-output"]["gpu"] == 0 -def test_cli_requires_experiment_cfg_path_and_policy_server(capsys): - """Require the Experiment config path and policy server at the CLI boundary.""" +def test_cli_requires_experiment_cfg_path(capsys): + """Require the Experiment config path at the CLI boundary.""" with pytest.raises(SystemExit, match="2"): main([]) assert "--experiment_cfg" in capsys.readouterr().err - with pytest.raises(SystemExit, match="2"): - main(["--experiment_cfg", str(OPENPI_EXPERIMENT_CFG_PATH)]) - assert "--policy_server" in capsys.readouterr().err - def test_cli_help_explains_paths_and_override_names(capsys): - """Describe the Experiment path, server selector, and typed override syntax.""" + """Describe the Experiment path and typed override syntax.""" with pytest.raises(SystemExit, match="0"): main(["--help"]) help_text = capsys.readouterr().out @@ -526,13 +581,9 @@ def test_cli_help_explains_paths_and_override_names(capsys): assert "--experiment_cfg PATH" in help_text assert "path to a typed Arena Experiment YAML configuration" in normalized_help_text assert "droid_pnp_srl_openpi_experiment.yaml" in help_text - assert "--policy_server {pi0}" in help_text assert "typed defaults < Experiment YAML < CLI overrides" in help_text assert "osmo.workflow_name=my-evaluation" in help_text assert "experiment_cfg.runs.droid_pnp_srl_openpi_billiard_hall.rollout_limit.num_episodes=4" in help_text - assert "Common overrides:" not in help_text - assert "Current defaults:" not in help_text - assert "experiment_cfg=" not in help_text def test_submission_rejects_legacy_json_experiment(tmp_path): @@ -541,7 +592,7 @@ def test_submission_rejects_legacy_json_experiment(tmp_path): experiment_path.write_text("{}", encoding="utf-8") with pytest.raises(AssertionError, match="requires a typed YAML Experiment Definition"): - build_arena_experiment_submission_cfg(experiment_path, "pi0") + build_arena_experiment_submission_cfg(experiment_path) def test_cli_accepts_arbitrary_paths_and_trailing_overrides(tmp_path, capsys): @@ -561,8 +612,6 @@ def test_cli_accepts_arbitrary_paths_and_trailing_overrides(tmp_path, capsys): main([ "--experiment_cfg", str(experiment_path), - "--policy_server", - "pi0", "osmo.dry_run=true", "osmo.workflow_name=path-based-submission", "experiment_runner.image=registry.example.com/evaluator:cli", @@ -577,47 +626,84 @@ def test_cli_accepts_arbitrary_paths_and_trailing_overrides(tmp_path, capsys): assert tasks[0]["image"] == "registry.example.com/evaluator:cli" +def test_dry_run_flag_renders_workflow_without_submitting(tmp_path, capsys): + """Render the workflow via the --dry_run flag instead of the osmo.dry_run override.""" + experiment_path = tmp_path / "experiment.yaml" + experiment_path.write_text( + """runs: + openpi: + environment: + type: pick_and_place_maple_table + policy: + type: isaaclab_arena_openpi.policy.pi0_remote_policy.Pi0RemotePolicy +""", + encoding="utf-8", + ) + + assert main(["--experiment_cfg", str(experiment_path), "--dry_run"]) == 0 + rendered = capsys.readouterr().out + assert "[dry-run] Rendered workflow YAML" in rendered + assert _rendered_workflow(rendered)["workflow"]["name"] == WorkflowCfg().workflow_name + + +def test_submission_cfg_to_str_lists_every_override_section(): + """Render the submission sections whose leaves are the valid Hydra override keys.""" + submission_cfg = _compose_submission() + + values = yaml.safe_load(submission_cfg_to_str(submission_cfg)) + + assert set(values) == {"osmo", "experiment_runner", "experiment_cfg"} + assert values["osmo"]["pool"] == WorkflowCfg().pool + assert values["experiment_runner"]["image"] == ExperimentRunnerTaskCfg().image + assert OPENPI_RUN_NAME in values["experiment_cfg"]["runs"] + + +def test_list_overrides_flag_prints_config_without_submitting(capsys): + """The --list_overrides flag composes and prints the config, then returns without submitting.""" + return_code = main([ + "--experiment_cfg", + str(OPENPI_EXPERIMENT_CFG_PATH), + "--list_overrides", + ]) + + assert return_code == 0 + values = yaml.safe_load(capsys.readouterr().out) + assert set(values) == {"osmo", "experiment_runner", "experiment_cfg"} + assert OPENPI_RUN_NAME in values["experiment_cfg"]["runs"] + + def test_experiment_path_is_relative_to_the_invocation_directory(tmp_path, monkeypatch): """Resolve a relative Experiment path from the caller's working directory.""" experiment_path = tmp_path / "experiment.yaml" experiment_path.write_text(OPENPI_EXPERIMENT_CFG_PATH.read_text(encoding="utf-8"), encoding="utf-8") monkeypatch.chdir(tmp_path) - submission_cfg = build_arena_experiment_submission_cfg("experiment.yaml", "pi0") + submission_cfg = build_arena_experiment_submission_cfg("experiment.yaml") assert isinstance(submission_cfg.experiment_cfg, ArenaExperimentCfg) - assert isinstance(submission_cfg.policy_server, Pi0ServerTaskCfg) - - -def test_structural_policy_override_is_checked_against_server_variant(): - """Check server compatibility against the effective structurally overridden policy.""" - with pytest.raises(AssertionError, match="pi0 server is configured for 'pi05'"): - _compose_and_submit([ - "osmo.dry_run=true", - f"experiment_cfg.runs.{OPENPI_RUN_NAME}.policy={{policy_variant:pi0}}", - ]) + assert isinstance(submission_cfg.experiment_cfg.runs[OPENPI_RUN_NAME].policy, Pi0RemotePolicyCfg) -def test_server_variant_cannot_relabel_known_pi05_model(): - """Reject a client-compatible label when the selected server model remains pi05.""" - with pytest.raises(AssertionError, match="policy_config.*serves variant 'pi05'.*policy_variant 'pi0'"): - _compose_and_submit([ - "osmo.dry_run=true", - f"experiment_cfg.runs.{OPENPI_RUN_NAME}.policy.policy_variant=pi0", - "policy_server.policy_variant=pi0", - ]) +def test_policy_variant_override_derives_matching_pi0_server(capsys): + """Derive each Run's pi0 server checkpoint from the effective client policy_variant.""" + return_code = _compose_and_submit([ + "osmo.dry_run=true", + f"experiment_cfg.runs.{OPENPI_RUN_NAME}.policy.policy_variant=pi0", + ]) + assert return_code == 0 + workflow = _rendered_workflow(capsys.readouterr().out) + server_command = _task_file(_workflow_tasks(workflow)[1], "/tmp/entry.sh")["contents"] + assert "--policy.config=pi0_droid_jointpos_polaris" in server_command + assert "--policy.dir=gs://openpi-assets-simeval/pi0_droid_jointpos" in server_command + served_policy = _embedded_experiment(_workflow_tasks(workflow)[0])["runs"][OPENPI_RUN_NAME]["policy"] + assert served_policy["policy_variant"] == "pi0" -def test_submitter_rejects_server_without_matching_run(): - """Reject an explicitly selected server that cannot serve any Experiment Run.""" - submission_cfg = ArenaExperimentSubmissionCfg( - experiment_cfg=_zero_action_experiment_cfg(), - osmo=WorkflowCfg(dry_run=True), - policy_server=Pi0ServerTaskCfg(), - ) - with pytest.raises(AssertionError, match="requires at least one Run using Pi0RemotePolicy"): - submit_arena_experiment(submission_cfg) +def test_pi0_server_cfg_rejects_mismatched_variant_and_config(): + """Reject a server config whose checkpoint fields do not match policy_variant.""" + with pytest.raises(AssertionError, match="policy_config.*serves variant 'pi05'.*policy_variant 'pi0'"): + Pi0ServerTaskCfg(policy_variant="pi0") def test_pi0_server_quotes_configurable_shell_values(): diff --git a/isaaclab_arena/utils/dicts.py b/isaaclab_arena/utils/dicts.py new file mode 100644 index 000000000..edb1e0509 --- /dev/null +++ b/isaaclab_arena/utils/dicts.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for working with plain dictionaries.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TypeVar + +KeyT = TypeVar("KeyT") +ValueT = TypeVar("ValueT") + + +def invert_dict(mapping: Mapping[KeyT, ValueT]) -> dict[ValueT, KeyT]: + """Return ``mapping`` with its keys and values swapped; the values must be unique. + + Args: + mapping: Mapping whose values are hashable and pairwise distinct. + + Returns: + A new dict mapping each value to the key it came from. + """ + inverted = {value: key for key, value in mapping.items()} + assert len(inverted) == len(mapping), f"Cannot invert a mapping with duplicate values: {mapping}" + return inverted diff --git a/isaaclab_arena_environments/robolab/experiment_configs/robolab_2_tasks_pi0_and_cosmos.yaml b/isaaclab_arena_environments/robolab/experiment_configs/robolab_2_tasks_pi0_and_cosmos.yaml new file mode 100644 index 000000000..5776c2645 --- /dev/null +++ b/isaaclab_arena_environments/robolab/experiment_configs/robolab_2_tasks_pi0_and_cosmos.yaml @@ -0,0 +1,68 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Runs two robolab tasks with both the OpenPI (pi0) and Cosmos policies (4 runs). On OSMO the +# inference server for each Run is derived from its policy type, so this launches pi0 servers for +# the *_pi0 Runs and Cosmos servers for the *_cosmos Runs. Both pi0 and Cosmos Runs use 5 parallel +# environments. +# All runs use the home_office_robolab HDR environment map as the dome-light background. +runs: + + banana_in_bowl_pi0: + environment: &banana_in_bowl_env + type: isaaclab_arena_environments/robolab/tasks/banana_in_bowl.yaml + enable_cameras: true + # This YAML anchor keeps the shared OpenPI policy configuration in one place. + policy: &openpi_policy + type: isaaclab_arena_openpi.policy.pi0_remote_policy.Pi0RemotePolicy + policy_variant: pi05 + policy_device: cuda:0 + remote_host: 127.0.0.1 + remote_port: 8000 + openpi_embodiment_adapter: droid + # This YAML anchor keeps the shared pi0 environment-builder settings in one place. + environment_builder: &pi0_builder + num_envs: 25 + variations: &hdr_variation + light: + hdr_image: + enabled: true + hdr_names: ["home_office_robolab"] + rollout_limit: + num_episodes: 100 + + banana_in_bowl_cosmos: + environment: *banana_in_bowl_env + # This YAML anchor keeps the shared Cosmos policy configuration in one place. + policy: &cosmos_policy + type: isaaclab_arena_cosmos.policy.cosmos_remote_policy.CosmosRemotePolicy + cosmos_embodiment_adapter: droid + policy_device: cuda:0 + remote_host: 127.0.0.1 + remote_port: 8000 + # This YAML anchor keeps the shared Cosmos environment-builder settings in one place. + environment_builder: &cosmos_builder + num_envs: 25 + variations: *hdr_variation + rollout_limit: + num_episodes: 100 + + banana_on_plate_pi0: + environment: &banana_on_plate_env + type: isaaclab_arena_environments/robolab/tasks/banana_on_plate.yaml + enable_cameras: true + policy: *openpi_policy + environment_builder: *pi0_builder + variations: *hdr_variation + rollout_limit: + num_episodes: 100 + + banana_on_plate_cosmos: + environment: *banana_on_plate_env + policy: *cosmos_policy + environment_builder: *cosmos_builder + variations: *hdr_variation + rollout_limit: + num_episodes: 100 diff --git a/isaaclab_arena_gr00t/policy/gr00t_remote_closedloop_policy.py b/isaaclab_arena_gr00t/policy/gr00t_remote_closedloop_policy.py index 7524d8326..343a6fa30 100644 --- a/isaaclab_arena_gr00t/policy/gr00t_remote_closedloop_policy.py +++ b/isaaclab_arena_gr00t/policy/gr00t_remote_closedloop_policy.py @@ -14,6 +14,7 @@ import gymnasium as gym import torch from dataclasses import dataclass +from enum import Enum from typing import Any from gr00t.policy.server_client import PolicyClient as Gr00tPolicyClient @@ -33,6 +34,20 @@ from isaaclab_arena_gr00t.utils.io_utils import create_config_from_yaml, load_gr00t_modality_config_from_file +class ActionSchedulerType(str, Enum): + """Action scheduler used to consume a policy's inference chunks.""" + + CHUNK = "chunk" + SYNCED_BATCH = "synced_batch" + + def get_scheduler_cls(self) -> type[ActionScheduler]: + """Return the action-scheduler class this type selects.""" + return { + ActionSchedulerType.CHUNK: ActionChunkScheduler, + ActionSchedulerType.SYNCED_BATCH: SyncedBatchActionScheduler, + }[self] + + # TODO(xinjieyao, 2026-04-27): Consider adding RemotePolicyCfg and deriving this config from it. @dataclass class Gr00tRemoteClosedloopPolicyCfg(Gr00tBasePolicyCfg): @@ -54,8 +69,8 @@ class Gr00tRemoteClosedloopPolicyCfg(Gr00tBasePolicyCfg): remote_api_token: str | None = None """Optional policy-server API token.""" - scheduler: str = "chunk" - """Action scheduler used to consume inference chunks: "chunk" or "synced_batch".""" + scheduler: ActionSchedulerType = ActionSchedulerType.CHUNK + """Action scheduler used to consume inference chunks.""" @register_policy @@ -71,12 +86,7 @@ class Gr00tRemoteClosedloopPolicy(PolicyBase[Gr00tRemoteClosedloopPolicyCfg]): def __init__(self, config: Gr00tRemoteClosedloopPolicyCfg): super().__init__(config) - action_scheduler_cls: type[ActionScheduler] - if config.scheduler == "synced_batch": - action_scheduler_cls = SyncedBatchActionScheduler - else: - assert config.scheduler == "chunk", f"Unknown action scheduler: {config.scheduler}" - action_scheduler_cls = ActionChunkScheduler + action_scheduler_cls = ActionSchedulerType(config.scheduler).get_scheduler_cls() # Policy config (for obs/action translation — no model loading) # TODO(xinjieyao, 2026-04-27): to be refactored diff --git a/osmo/submit_arena_experiment.py b/osmo/submit_arena_experiment.py index aa55886b8..b76fdd96f 100644 --- a/osmo/submit_arena_experiment.py +++ b/osmo/submit_arena_experiment.py @@ -9,6 +9,7 @@ import argparse import sys +import yaml from dataclasses import dataclass, field from pathlib import Path @@ -18,20 +19,13 @@ from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg from isaaclab_arena.evaluation.arena_experiment_config_loader import load_arena_experiment_from_config_file +from isaaclab_arena.hydra.typed_experiment_serializer import serialize_arena_experiment_to_yaml from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides -from osmo.tasks.base_task import TaskCfg from osmo.tasks.experiment_runner_task import ExperimentRunnerTaskCfg -from osmo.tasks.pi0_server_task import Pi0ServerTaskCfg -from osmo.workflows.arena_experiment_workflow import Pi0ArenaExperimentWorkflow +from osmo.workflows.arena_experiment_workflow import ArenaExperimentWorkflow from osmo.workflows.workflow import WorkflowCfg SUBMISSION_CONFIG_NAME = "osmo_arena_experiment_submission" -POLICY_SERVER_TASK_CFG_BY_NAME = { - "pi0": Pi0ServerTaskCfg, -} -POLICY_SERVER_WORKFLOW_BY_CONFIG_TYPE = { - Pi0ServerTaskCfg: Pi0ArenaExperimentWorkflow, -} @dataclass @@ -41,9 +35,6 @@ class ArenaExperimentSubmissionCfg: experiment_cfg: ArenaExperimentCfg """Evaluation semantics executed by ``experiment_runner.py``.""" - policy_server: TaskCfg - """Co-scheduled policy server used by the Experiment's remote policy clients.""" - osmo: WorkflowCfg = field(default_factory=WorkflowCfg) """OSMO scheduling, resource, and timeout configuration.""" @@ -55,37 +46,27 @@ def submit_arena_experiment(submission_cfg: ArenaExperimentSubmissionCfg) -> int """Build and submit the OSMO workflow described by ``submission_cfg``. Args: - submission_cfg: Composed Experiment, task, server, and OSMO configuration. + submission_cfg: Composed Experiment and OSMO configuration. Returns: The OSMO submission process status. """ - workflow_cfg = submission_cfg.osmo - experiment_runner_task_cfg = submission_cfg.experiment_runner - policy_server_task_cfg = submission_cfg.policy_server - workflow_cls = POLICY_SERVER_WORKFLOW_BY_CONFIG_TYPE.get(type(policy_server_task_cfg)) - assert ( - workflow_cls is not None - ), f"No policy-server workflow is registered for configuration type {type(policy_server_task_cfg).__name__}" - workflow = workflow_cls( - workflow_cfg=workflow_cfg, + workflow = ArenaExperimentWorkflow( + workflow_cfg=submission_cfg.osmo, experiment_cfg=submission_cfg.experiment_cfg, - server_task_cfg=policy_server_task_cfg, - task_cfg=experiment_runner_task_cfg, + task_cfg=submission_cfg.experiment_runner, ) return workflow.submit_workflow().returncode def build_arena_experiment_submission_cfg( experiment_cfg_path: str | Path, - policy_server_name: str, overrides: list[str] | None = None, ) -> ArenaExperimentSubmissionCfg: - """Load an Experiment, select its policy server, and apply typed overrides. + """Load an Experiment and apply typed submission overrides. Args: experiment_cfg_path: Arena Experiment configuration file. - policy_server_name: Built-in policy-server implementation name. overrides: Hydra field overrides rooted at the composed submission. Returns: @@ -97,18 +78,10 @@ def build_arena_experiment_submission_cfg( ".yml", }, f"OSMO Experiment submission requires a typed YAML Experiment Definition; got '{experiment_cfg_path}'" experiment_cfg = load_arena_experiment_from_config_file(experiment_cfg_path, device="cuda:0") - available_names = ", ".join(sorted(POLICY_SERVER_TASK_CFG_BY_NAME)) - assert ( - policy_server_name in POLICY_SERVER_TASK_CFG_BY_NAME - ), f"Unknown policy server '{policy_server_name}'. Available policy servers: {available_names}" - policy_server = POLICY_SERVER_TASK_CFG_BY_NAME[policy_server_name]() - base_submission = ArenaExperimentSubmissionCfg( - experiment_cfg=experiment_cfg, - policy_server=policy_server, - ) + base_submission = ArenaExperimentSubmissionCfg(experiment_cfg=experiment_cfg) - # The Experiment file and policy-server selector determine the concrete config types. - # Register that concrete root so Hydra validates every trailing override against it. + # The Experiment file determines the concrete config types. Register that concrete root so + # Hydra validates every trailing override against it. ConfigStore.instance().store(name=SUBMISSION_CONFIG_NAME, node=base_submission) with initialize(version_base=None, config_path=None): composed = compose(config_name=SUBMISSION_CONFIG_NAME, overrides=overrides or []) @@ -117,11 +90,25 @@ def build_arena_experiment_submission_cfg( return submission_cfg +def submission_cfg_to_str(submission_cfg: ArenaExperimentSubmissionCfg) -> str: + """Render the composed submission as YAML; every leaf is a valid Hydra KEY=VALUE override.""" + # osmo / experiment_runner are plain dataclasses; OmegaConf dumps them directly. + # experiment_cfg is polymorphic (policy.type, environment.type, …), so it needs the + # Experiment serializer to emit the same YAML shape Hydra overrides expect. + submission_values = { + "osmo": OmegaConf.to_container(OmegaConf.structured(submission_cfg.osmo), resolve=True, enum_to_str=True), + "experiment_runner": OmegaConf.to_container( + OmegaConf.structured(submission_cfg.experiment_runner), resolve=True, enum_to_str=True + ), + "experiment_cfg": yaml.safe_load(serialize_arena_experiment_to_yaml(submission_cfg.experiment_cfg)), + } + return yaml.safe_dump(submission_values, sort_keys=False) + + def _create_argument_parser() -> argparse.ArgumentParser: """Create the path-first submission command-line parser.""" - policy_server_choices = ",".join(POLICY_SERVER_TASK_CFG_BY_NAME) parser = argparse.ArgumentParser( - usage=f"%(prog)s [-h] --experiment_cfg PATH --policy_server {{{policy_server_choices}}} [OVERRIDE ...]", + usage="%(prog)s [-h] --experiment_cfg PATH [--dry_run] [--list_overrides] [OVERRIDE ...]", description="Submit a typed Arena Experiment as an OSMO workflow.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=r""" @@ -129,7 +116,6 @@ def _create_argument_parser() -> argparse.ArgumentParser: python -m osmo.submit_arena_experiment \ --experiment_cfg isaaclab_arena_environments/experiment_configs/droid_pnp_srl_openpi_experiment.yaml \ - --policy_server pi0 \ osmo.workflow_name=my-evaluation \ experiment_cfg.runs.droid_pnp_srl_openpi_billiard_hall.rollout_limit.num_episodes=4 @@ -147,10 +133,14 @@ def _create_argument_parser() -> argparse.ArgumentParser: help="path to a typed Arena Experiment YAML configuration", ) parser.add_argument( - "--policy_server", - required=True, - choices=POLICY_SERVER_TASK_CFG_BY_NAME, - help="co-scheduled policy-server implementation", + "--dry_run", + action="store_true", + help="render the workflow YAML and print it instead of submitting to OSMO", + ) + parser.add_argument( + "--list_overrides", + action="store_true", + help="print the composed submission configuration and exit; every leaf is a valid Hydra KEY=VALUE override", ) parser.allow_abbrev = False return parser @@ -158,16 +148,18 @@ def _create_argument_parser() -> argparse.ArgumentParser: def main(cli_args: list[str] | None = None) -> int: """Load the Experiment, apply overrides, and submit its OSMO workflow.""" - # Argparse resolves the Experiment path and server selector first; they determine - # the concrete configs Hydra receives for its remaining overrides. parser = _create_argument_parser() args, overrides = parser.parse_known_args(cli_args) assert_hydra_overrides(overrides, parser) + if args.dry_run: + overrides = [*overrides, "osmo.dry_run=true"] submission_cfg = build_arena_experiment_submission_cfg( experiment_cfg_path=args.experiment_cfg_path, - policy_server_name=args.policy_server, overrides=overrides, ) + if args.list_overrides: + print(submission_cfg_to_str(submission_cfg)) + return 0 return submit_arena_experiment(submission_cfg) diff --git a/osmo/tasks/base_task.py b/osmo/tasks/base_task.py index 0f479d299..a074d94e3 100644 --- a/osmo/tasks/base_task.py +++ b/osmo/tasks/base_task.py @@ -24,6 +24,9 @@ class TaskCfg: class BaseTask(ABC): """Abstract base task for an Isaac Lab Arena OSMO workflow.""" + task_cfg_type: type[TaskCfg] = TaskCfg + """Config dataclass this task consumes.""" + def __init__( self, task_cfg: TaskCfg | None = None, diff --git a/osmo/tasks/cosmos_server_task.py b/osmo/tasks/cosmos_server_task.py index 92282b5ea..7285c6b13 100644 --- a/osmo/tasks/cosmos_server_task.py +++ b/osmo/tasks/cosmos_server_task.py @@ -11,7 +11,10 @@ from dataclasses import dataclass from typing import Any -from osmo.tasks.base_task import BaseTask, TaskCfg +from isaaclab_arena_cosmos.policy.cosmos_remote_policy import CosmosRemotePolicy +from osmo.tasks.base_task import TaskCfg +from osmo.tasks.policy_server_task import PolicyServerTask +from osmo.workflows.server_task_registry import register_server_task from osmo.workflows.workflow_constants import POLICY_SERVER_PORT @@ -26,9 +29,13 @@ class CosmosServerTaskCfg(TaskCfg): """Checkpoint the server serves. Baked into the image at build time (see build_server_image.sh).""" -class CosmosServerTask(BaseTask): +@register_server_task +class CosmosServerTask(PolicyServerTask): """OSMO task that serves a Cosmos policy for an eval/policy-runner task to connect to.""" + policy_type = CosmosRemotePolicy + task_cfg_type = CosmosServerTaskCfg + def __init__( self, task_cfg: CosmosServerTaskCfg | None = None, @@ -36,7 +43,7 @@ def __init__( *, task_name: str, ) -> None: - super().__init__(task_name=task_name, task_cfg=task_cfg or CosmosServerTaskCfg(), lead=lead) + super().__init__(task_name=task_name, task_cfg=task_cfg or self.task_cfg_type(), lead=lead) def _get_image(self) -> str: return self.task_cfg.image diff --git a/osmo/tasks/dreamzero_server_task.py b/osmo/tasks/dreamzero_server_task.py index 2e9f1d868..8d17923d7 100644 --- a/osmo/tasks/dreamzero_server_task.py +++ b/osmo/tasks/dreamzero_server_task.py @@ -28,6 +28,8 @@ class DreamZeroServerTaskCfg(TaskCfg): class DreamZeroServerTask(BaseTask): """OSMO task that serves a DreamZero policy for a policy-runner task to connect to.""" + task_cfg_type = DreamZeroServerTaskCfg + def __init__( self, task_cfg: DreamZeroServerTaskCfg | None = None, @@ -35,7 +37,7 @@ def __init__( *, task_name: str, ) -> None: - super().__init__(task_name=task_name, task_cfg=task_cfg or DreamZeroServerTaskCfg(), lead=lead) + super().__init__(task_name=task_name, task_cfg=task_cfg or self.task_cfg_type(), lead=lead) def _get_image(self) -> str: return self.task_cfg.image diff --git a/osmo/tasks/experiment_runner_task.py b/osmo/tasks/experiment_runner_task.py index 1d9c34297..31e39ca43 100644 --- a/osmo/tasks/experiment_runner_task.py +++ b/osmo/tasks/experiment_runner_task.py @@ -33,6 +33,12 @@ class ExperimentRunnerTaskCfg(TaskCfg): image: str = DEFAULT_EXPERIMENT_RUNNER_IMAGE """Container image that runs the Arena Experiment.""" + record_camera_video: bool = True + """Record one mp4 per (env, camera, episode) from each Run's camera observations.""" + + record_viewport_video: bool = False + """Record a viewport video for each Run.""" + class ExperimentRunnerTask(BaseTask): """Lead OSMO task that runs every Run in one effective Arena Experiment.""" @@ -81,4 +87,8 @@ def _get_run_script(self) -> str: "none", "--enable_cameras", ] + if self.task_cfg.record_camera_video: + command.append("--record_camera_video") + if self.task_cfg.record_viewport_video: + command.append("--record_viewport_video") return f"set -euo pipefail\n{shlex.join(command)}\n" diff --git a/osmo/tasks/gr00t_server_task.py b/osmo/tasks/gr00t_server_task.py index 6b8f2197f..703209589 100644 --- a/osmo/tasks/gr00t_server_task.py +++ b/osmo/tasks/gr00t_server_task.py @@ -8,7 +8,10 @@ from dataclasses import dataclass from typing import Any -from osmo.tasks.base_task import BaseTask, TaskCfg +from isaaclab_arena_gr00t.policy.gr00t_remote_closedloop_policy import Gr00tRemoteClosedloopPolicy +from osmo.tasks.base_task import TaskCfg +from osmo.tasks.policy_server_task import PolicyServerTask +from osmo.workflows.server_task_registry import register_server_task from osmo.workflows.workflow_constants import POLICY_SERVER_PORT @@ -26,9 +29,13 @@ class Gr00tServerTaskCfg(TaskCfg): """Embodiment tag for the droid manipulation config (see droid_manip_gr00t_closedloop_config.yaml).""" -class Gr00tServerTask(BaseTask): +@register_server_task +class Gr00tServerTask(PolicyServerTask): """OSMO task that serves a GR00T policy for an eval/policy-runner task to connect to.""" + policy_type = Gr00tRemoteClosedloopPolicy + task_cfg_type = Gr00tServerTaskCfg + def __init__( self, task_cfg: Gr00tServerTaskCfg | None = None, @@ -36,7 +43,7 @@ def __init__( *, task_name: str, ) -> None: - super().__init__(task_name=task_name, task_cfg=task_cfg or Gr00tServerTaskCfg(), lead=lead) + super().__init__(task_name=task_name, task_cfg=task_cfg or self.task_cfg_type(), lead=lead) def _get_image(self) -> str: return self.task_cfg.image diff --git a/osmo/tasks/pi0_server_task.py b/osmo/tasks/pi0_server_task.py index effa96df5..d4fdd406e 100644 --- a/osmo/tasks/pi0_server_task.py +++ b/osmo/tasks/pi0_server_task.py @@ -9,7 +9,13 @@ from dataclasses import dataclass from typing import Any -from osmo.tasks.base_task import BaseTask, TaskCfg +from isaaclab_arena.policy.policy_base import PolicyCfg +from isaaclab_arena.utils.dicts import invert_dict +from isaaclab_arena_openpi.policy.pi0_remote_config import Pi0RemotePolicyCfg +from isaaclab_arena_openpi.policy.pi0_remote_policy import Pi0RemotePolicy +from osmo.tasks.base_task import TaskCfg +from osmo.tasks.policy_server_task import PolicyServerTask +from osmo.workflows.server_task_registry import register_server_task from osmo.workflows.workflow_constants import POLICY_SERVER_PORT OPENPI_APP_DIR = "/app" @@ -23,6 +29,8 @@ "gs://openpi-assets-simeval/pi0_droid_jointpos": "pi0", "gs://openpi-assets-simeval/pi05_droid_jointpos": "pi05", } +PI0_POLICY_CONFIG_BY_VARIANT = invert_dict(PI0_VARIANT_BY_POLICY_CONFIG) +PI0_POLICY_DIR_BY_VARIANT = invert_dict(PI0_VARIANT_BY_POLICY_DIR) @dataclass @@ -41,9 +49,6 @@ class Pi0ServerTaskCfg(TaskCfg): policy_dir: str = "gs://openpi-assets-simeval/pi05_droid_jointpos" """openpi checkpoint directory.""" - client_ping_timeout_s: float | None = 300.0 - """Seconds Arena clients wait for a pong while the OSMO server compiles its first inference.""" - def __post_init__(self) -> None: assert ( self.policy_variant in PI0_POLICY_VARIANTS @@ -59,10 +64,32 @@ def __post_init__(self) -> None: f"not policy_variant {self.policy_variant!r}" ) + @classmethod + def for_policy_variant(cls, policy_variant: str) -> Pi0ServerTaskCfg: + """Build the default server deployment for a client ``policy_variant``.""" + assert ( + policy_variant in PI0_POLICY_VARIANTS + ), f"pi0 server policy_variant must be one of {sorted(PI0_POLICY_VARIANTS)}, got {policy_variant!r}" + return cls( + policy_variant=policy_variant, + policy_config=PI0_POLICY_CONFIG_BY_VARIANT[policy_variant], + policy_dir=PI0_POLICY_DIR_BY_VARIANT[policy_variant], + ) + -class Pi0ServerTask(BaseTask): +@register_server_task +class Pi0ServerTask(PolicyServerTask): """OSMO task that serves a pi0 policy.""" + policy_type = Pi0RemotePolicy + task_cfg_type = Pi0ServerTaskCfg + + @classmethod + def task_cfg_for_policy(cls, policy_cfg: PolicyCfg) -> Pi0ServerTaskCfg: + """Derive the server checkpoint from the client policy variant.""" + assert isinstance(policy_cfg, Pi0RemotePolicyCfg) + return cls.task_cfg_type.for_policy_variant(policy_cfg.policy_variant) + def __init__( self, task_cfg: Pi0ServerTaskCfg | None = None, @@ -70,7 +97,7 @@ def __init__( *, task_name: str, ) -> None: - super().__init__(task_name=task_name, task_cfg=task_cfg or Pi0ServerTaskCfg(), lead=lead) + super().__init__(task_name=task_name, task_cfg=task_cfg or self.task_cfg_type(), lead=lead) def _get_image(self) -> str: return self.task_cfg.image diff --git a/osmo/tasks/policy_server_task.py b/osmo/tasks/policy_server_task.py new file mode 100644 index 000000000..7ffb638cb --- /dev/null +++ b/osmo/tasks/policy_server_task.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Base task for OSMO policy servers derived from Arena client policies.""" + +from isaaclab_arena.assets.registries import PolicyRegistry +from isaaclab_arena.policy.policy_base import PolicyBase, PolicyCfg +from osmo.tasks.base_task import BaseTask, TaskCfg + + +class PolicyServerTask(BaseTask): + """OSMO server task associated with one Arena client policy type.""" + + policy_type: type[PolicyBase] + """Client policy class served by this task.""" + + @classmethod + def task_cfg_for_policy(cls, policy_cfg: PolicyCfg) -> TaskCfg: + """Build this server's task config from a client policy config.""" + assert PolicyRegistry().get_policy_type_for_cfg(policy_cfg) is cls.policy_type + return cls.task_cfg_type() diff --git a/osmo/workflows/arena_experiment_workflow.py b/osmo/workflows/arena_experiment_workflow.py index b182131ca..1df10b95e 100644 --- a/osmo/workflows/arena_experiment_workflow.py +++ b/osmo/workflows/arena_experiment_workflow.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""OSMO workflows for evaluating complete Arena Experiments.""" +"""OSMO workflow for evaluating complete Arena Experiments.""" from __future__ import annotations @@ -12,47 +12,37 @@ from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg from isaaclab_arena.evaluation.arena_run import ArenaRunCfg -from isaaclab_arena_openpi.policy.pi0_remote_config import Pi0RemotePolicyCfg from osmo.tasks.base_task import BaseTask from osmo.tasks.collect_experiment_outputs_task import CollectExperimentOutputsTask from osmo.tasks.experiment_runner_task import ExperimentRunnerTask, ExperimentRunnerTaskCfg -from osmo.tasks.pi0_server_task import Pi0ServerTask, Pi0ServerTaskCfg +from osmo.tasks.policy_server_task import PolicyServerTask +from osmo.workflows.server_task_registry import ServerTaskRegistry from osmo.workflows.workflow import Workflow, WorkflowCfg from osmo.workflows.workflow_constants import POLICY_SERVER_PORT -class Pi0ArenaExperimentWorkflow(Workflow): - """Run every Arena Experiment Run in its own OSMO group.""" +class ArenaExperimentWorkflow(Workflow): + """Run every Arena Experiment Run in its own OSMO group, co-scheduling each Run's server.""" constructs_groups_directly = True task_cfg_type = ExperimentRunnerTaskCfg - server_task_cfg_type = Pi0ServerTaskCfg - """Configuration type used by this policy-server workflow.""" - experiment_output_resource_name = "experiment-output" def __init__( self, workflow_cfg: WorkflowCfg, experiment_cfg: ArenaExperimentCfg, - server_task_cfg: Pi0ServerTaskCfg, group_name: str = "arena", task_cfg: ExperimentRunnerTaskCfg | None = None, ) -> None: assert isinstance(experiment_cfg, ArenaExperimentCfg) self.experiment_cfg = deepcopy(experiment_cfg) - self.pi0_server_task_cfg = server_task_cfg super().__init__( workflow_cfg=workflow_cfg, task_cfg=task_cfg or ExperimentRunnerTaskCfg(), group_name=group_name, ) - # Every Pi0 Run gets a dedicated server task. Verify that all of those - # clients request the variant configured for the server deployment. - pi0_policy_variants_by_run = self._get_pi0_policy_variants_by_run() - self._assert_pi0_server_compatible(pi0_policy_variants_by_run) - def _get_group_dicts(self) -> list[dict[str, Any]]: """Create one independently scheduled group per Run, then collect their outputs into one Experiment output.""" run_group_dicts: list[dict[str, Any]] = [] @@ -77,22 +67,27 @@ def _create_run_group_dict( run_name: str, run_config: ArenaRunCfg, ) -> tuple[dict[str, Any], str]: - """Create one OSMO group that executes a single-Run Arena Experiment.""" + """Create one OSMO group that executes a single-Run Arena Experiment, plus its server if any.""" experiment_runner_task_name = f"experiment-runner-{run_index}" + # Snapshot this Run alone so the Experiment Runner task embeds a single-Run Experiment. single_run_experiment_config = ArenaExperimentCfg(runs={run_name: deepcopy(run_config)}) - pi0_policy_server_tasks: list[BaseTask] = [] + policy_server_tasks: list[BaseTask] = [] run_policy_config = single_run_experiment_config.runs[run_name].policy - if isinstance(run_policy_config, Pi0RemotePolicyCfg): - pi0_server_task_name = f"policy-server-{run_index}" - self._configure_pi0_remote_policy_for_server(run_policy_config, pi0_server_task_name) - pi0_policy_server_tasks.append( - Pi0ServerTask( - self.pi0_server_task_cfg, - lead=False, - task_name=pi0_server_task_name, - ) - ) + # Look up whether this client policy needs a co-scheduled OSMO inference server. + server_type: type[PolicyServerTask] | None = ServerTaskRegistry().get_server_type_for_policy_cfg( + run_policy_config + ) + if server_type is not None: + server_task_name = f"policy-server-{run_index}" + # Point the client policy at this group's dedicated server task host/port. + run_policy_config.remote_host = server_type.host_token(server_task_name) + run_policy_config.remote_port = POLICY_SERVER_PORT + # Build the server deployment config from the Run's client policy (e.g. pi0 variant). + server_cfg = server_type.task_cfg_for_policy(run_policy_config) + # Instantiate the OSMO server task associated with the client policy. + server_task = server_type(server_cfg, lead=False, task_name=server_task_name) + policy_server_tasks.append(server_task) # Construct this after connecting the policy because the task snapshots the Experiment. experiment_runner_task = ExperimentRunnerTask( @@ -102,7 +97,7 @@ def _create_run_group_dict( task_name=experiment_runner_task_name, published_output_url=None, ) - run_group_tasks = [experiment_runner_task, *pi0_policy_server_tasks] + run_group_tasks = [experiment_runner_task, *policy_server_tasks] run_group_dict = { "name": f"arena-run-{run_index}", @@ -135,37 +130,3 @@ def _create_resources_dict(self) -> dict[str, dict[str, Any]]: "default": run_task_resource, self.experiment_output_resource_name: experiment_output_task_resource, } - - def _get_pi0_policy_variants_by_run(self) -> dict[str, str]: - """Return effective pi0-remote Run variants needed for compatibility checks.""" - pi0_policy_variants_by_run = {} - for run_name, run_config in self.experiment_cfg.runs.items(): - if not isinstance(run_config.policy, Pi0RemotePolicyCfg): - continue - pi0_policy_variants_by_run[run_name] = run_config.policy.policy_variant - return pi0_policy_variants_by_run - - def _assert_pi0_server_compatible(self, pi0_policy_variants_by_run: dict[str, str]) -> None: - """Require Pi0RemotePolicy Runs whose variants match the deployed server.""" - assert pi0_policy_variants_by_run, "pi0 server requires at least one Run using Pi0RemotePolicy" - incompatible_policy_variants_by_run = { - run_name: policy_variant - for run_name, policy_variant in pi0_policy_variants_by_run.items() - if policy_variant != self.pi0_server_task_cfg.policy_variant - } - assert not incompatible_policy_variants_by_run, ( - f"pi0_remote Runs require variants {incompatible_policy_variants_by_run}, but the pi0 server is configured" - f" for '{self.pi0_server_task_cfg.policy_variant}'" - ) - - def _configure_pi0_remote_policy_for_server( - self, - pi0_remote_policy_config: Pi0RemotePolicyCfg, - pi0_server_task_name: str, - ) -> None: - """Configure a Pi0 remote policy to use its dedicated OSMO server task.""" - pi0_remote_policy_config.remote_host = Pi0ServerTask.host_token(pi0_server_task_name) - pi0_remote_policy_config.remote_port = POLICY_SERVER_PORT - # The first OSMO inference may compile longer than the policy's normal - # keepalive timeout. Use the timeout owned by this server deployment. - pi0_remote_policy_config.ping_timeout = self.pi0_server_task_cfg.client_ping_timeout_s diff --git a/osmo/workflows/server_task_registry.py b/osmo/workflows/server_task_registry.py new file mode 100644 index 000000000..de5660108 --- /dev/null +++ b/osmo/workflows/server_task_registry.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Look up the OSMO policy server associated with an Arena client policy.""" + +from isaaclab_arena.assets.registries import PolicyRegistry +from isaaclab_arena.policy.policy_base import PolicyBase, PolicyCfg +from isaaclab_arena.utils.singleton import SingletonMeta +from osmo.tasks.policy_server_task import PolicyServerTask + +_SERVER_TASKS_LOADED = False + + +class ServerTaskRegistry(metaclass=SingletonMeta): + """Map Arena client policy types to their OSMO server task types.""" + + def __init__(self) -> None: + self._server_types_by_policy_type: dict[type[PolicyBase], type[PolicyServerTask]] = {} + + def register(self, server_type: type[PolicyServerTask]) -> None: + """Register a policy server under the client policy type it serves.""" + assert ( + server_type.policy_type not in self._server_types_by_policy_type + ), f"Policy {server_type.policy_type.__name__} already has a policy server" + self._server_types_by_policy_type[server_type.policy_type] = server_type + + def has_server_for_policy_type(self, policy_type: type[PolicyBase]) -> bool: + """Return whether a policy server is already registered for ``policy_type``.""" + return policy_type in self._server_types_by_policy_type + + def get_server_type_for_policy_cfg(self, policy_cfg: PolicyCfg) -> type[PolicyServerTask] | None: + """Get the server task for a client policy config, or None if it needs no server.""" + ensure_server_tasks_registered() + policy_type = PolicyRegistry().get_policy_type_for_cfg(policy_cfg) + return self._server_types_by_policy_type.get(policy_type) + + def get_all_server_types(self) -> list[type[PolicyServerTask]]: + """Return every registered policy server task type.""" + ensure_server_tasks_registered() + return list(self._server_types_by_policy_type.values()) + + +def register_server_task(cls: type[PolicyServerTask]) -> type[PolicyServerTask]: + """Decorator registering a policy server with the ServerTaskRegistry.""" + registry = ServerTaskRegistry() + if registry.has_server_for_policy_type(cls.policy_type): + print(f"WARNING: Policy server for {cls.policy_type.__name__} is already registered. Doing nothing.") + else: + registry.register(cls) + return cls + + +def ensure_server_tasks_registered() -> None: + """Import policy-server modules so their ``@register_server_task`` decorators run.""" + global _SERVER_TASKS_LOADED + if _SERVER_TASKS_LOADED: + return + _SERVER_TASKS_LOADED = True + import osmo.tasks.cosmos_server_task # noqa: F401 + import osmo.tasks.gr00t_server_task # noqa: F401 + import osmo.tasks.pi0_server_task # noqa: F401