Enable multi-policy experiments on OSMO

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

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

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

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,19 +90,32 @@ 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"""
 Example:
 
   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,27 +133,33 @@ 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
 
 
 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)
 
 

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,

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,17 +29,21 @@ 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,
         lead: bool | None = None,
         *,
         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

osmo/tasks/dreamzero_server_task.py

@@ -28,14 +28,16 @@ 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,
         lead: bool | None = None,
         *,
         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