Address review: lean on base PolicyRunnerTask, align with pi0

- Slim Gr00tPolicyRunnerTask to just its policy args + gr00t-specific
  flags; rely on the base for the run script, outputs, task name, and
  image (drops the out-of-scope image override and the wait-for-server
  preamble for now).
- Align Gr00tServerTask with pi0_server_task: inline task name, drop the
  wait-for-server helper, and fold model/embodiment/port into constants.
- Centralize the policy-server port in workflow_constants (POLICY_SERVER_PORT
  = 8000) and use it for both the GR00T and pi0 remote tasks.
- Reverse the workflow task order to match the pi0 workflow.

Signed-off-by: alex <amillane@nvidia.com>

Author: alexmillane

osmo/tasks/gr00t_policy_runner_task.py

@@ -3,114 +3,52 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""GR00T policy-runner task for the Isaac Lab Arena OSMO workflow.
-
-Runs ``policy_runner.py`` with the GR00T remote closed-loop policy against the
-GR00T server that shares its OSMO group. Mirrors the eval side of the
-``test_gr00t_closedloop_e2e`` CI job in ``.github/workflows/ci.yml``.
-"""
+"""GR00T policy-runner task for the Isaac Lab Arena OSMO workflow."""
 
 import argparse
-from typing import Any
-
-from tasks.gr00t_server_task import GR00T_SERVER_HOST_TOKEN, get_wait_for_server_script
-from tasks.policy_runner_task import POLICY_RUNNER_COMMAND, PolicyRunnerTask
-from workflows.workflow_constants import EVAL_OUTPUT_SWIFT_URL, OSMO_TASK_OUTPUT_DIR
 
-# Arena image name
-DEFAULT_IMAGE = "nvcr.io/nvstaging/isaac-amr/isaaclab_arena:latest"
+from tasks.policy_runner_task import PolicyRunnerTask
+from workflows.workflow_constants import POLICY_SERVER_PORT
 
-# GR00T remote closed-loop policy and the policy closed-loop config.
-GR00T_POLICY_TYPE = "isaaclab_arena_gr00t.policy.gr00t_remote_closedloop_policy.Gr00tRemoteClosedloopPolicy"
 DEFAULT_POLICY_CONFIG = "isaaclab_arena_gr00t/policy/config/droid_manip_gr00t_closedloop_config.yaml"
 
-DEFAULT_POLICY_RUNNER_ARGS = "--num_episodes 2 --headless --enable_cameras --num_envs 4 --record_camera_video"
-
 
 class Gr00tPolicyRunnerTask(PolicyRunnerTask):
     """OSMO task that evaluates the GR00T remote policy via a connection to the GR00T server."""
 
     def __init__(
         self,
-        workflow_args: Any,
-        task_args: Any,
-        image: str = DEFAULT_IMAGE,
+        workflow_args: argparse.Namespace,
+        task_args: argparse.Namespace,
         lead: bool | None = None,
     ) -> None:
-        super().__init__(workflow_args, task_args, image=image, lead=lead)
-        self.image = getattr(task_args, "arena_image", image)
+        super().__init__(workflow_args=workflow_args, task_args=task_args, lead=lead)
         self.policy_config_yaml_path = task_args.policy_config_yaml_path
-
-        # Tasks in an OSMO group each get their own IP (no shared loopback), so the server is reached
-        # via the {{host:<task-name>}} token, which OSMO resolves to the server task's IP at runtime.
         self.remote_host = task_args.remote_host
-        self.remote_port = task_args.server_port
 
     @staticmethod
     def add_task_arguments(parser: argparse.ArgumentParser) -> None:
+        PolicyRunnerTask.add_task_arguments(parser)
         group = parser.add_argument_group("gr00t policy runner")
-        group.add_argument("--arena_image", default=DEFAULT_IMAGE, help="Override the Arena dev image")
         group.add_argument(
             "--policy_config_yaml_path", default=DEFAULT_POLICY_CONFIG, help="GR00T closed-loop config YAML"
         )
-        group.add_argument(
-            "--arena_env_args",
-            required=True,
-            help=(
-                "Arena example-environment name plus its env-related arguments, "
-                "e.g. 'kitchen_pick_and_place --object cracker_box'"
-            ),
-        )
+        # Tasks in an OSMO group each get their own IP, so the server is reached via the
+        # {{host:<task-name>}} token, which OSMO resolves to the server task's IP at runtime.
         group.add_argument(
             "--remote_host",
-            default=GR00T_SERVER_HOST_TOKEN,
-            help="GR00T server host (defaults to the {{host:gr00t_server}}) name",
+            default="{{host:gr00t_server}}",
+            help="GR00T server host (defaults to {{host:gr00t_server}})",
         )
-        group.add_argument(
-            "--policy_runner_args",
-            default=DEFAULT_POLICY_RUNNER_ARGS,
-            help=(
-                "Policy-runner related arguments, e.g. '--num_episodes, --headless, --enable_cameras, --num_envs,"
-                " --record_camera_video'"
-            ),
-        )
-
-    @staticmethod
-    def get_task_name() -> str:
-        return "gr00t_policy_runner"
-
-    def _get_outputs(self) -> list[dict[str, Any]]:
-        # Evaluation outputs (videos, per-episode results, report) are uploaded per run.
-        return [{"url": EVAL_OUTPUT_SWIFT_URL}]
 
     def _get_policy_args(self) -> list[str]:
         return [
             "--policy_type",
-            GR00T_POLICY_TYPE,
+            "isaaclab_arena_gr00t.policy.gr00t_remote_closedloop_policy.Gr00tRemoteClosedloopPolicy",
             "--policy_config_yaml_path",
             self.policy_config_yaml_path,
             "--remote_host",
             self.remote_host,
             "--remote_port",
-            str(self.remote_port),
+            str(POLICY_SERVER_PORT),
         ]
-
-    def _get_run_script(self) -> str:
-        # Override the base runner: block on the GR00T server coming up before launching the eval.
-        policy_args_str = " ".join(self._get_policy_args())
-        return (
-            "set -euxo pipefail\n"
-            "ldconfig\n"
-            "nvidia-smi\n"
-            "cd /workspaces/isaaclab_arena\n"
-            "[ -e submodules/IsaacLab/_isaac_sim ] || ln -s /isaac-sim/ submodules/IsaacLab/_isaac_sim\n"
-            "\n"
-            f"{get_wait_for_server_script(self.remote_host, self.remote_port)}"
-            "\n"
-            f"{POLICY_RUNNER_COMMAND} "
-            f"{policy_args_str} "
-            # Write evaluation outputs to the OSMO task output mount (uploaded to EVAL_OUTPUT_SWIFT_URL).
-            f"--output_base_dir {OSMO_TASK_OUTPUT_DIR} "
-            f"{self.policy_runner_args} "
-            f"{self.arena_env_args}\n"
-        )

osmo/tasks/gr00t_server_task.py

@@ -5,39 +5,17 @@
 
 """GR00T inference-server task for OSMO eval workflows, used by the GR00T policy-runner task."""
 
-import argparse
 from typing import Any
 
 from tasks.base_task import BaseTask
+from workflows.workflow_constants import POLICY_SERVER_PORT
 
 # GR00T server image (containing the droid checkpoints).
 DEFAULT_IMAGE = "nvcr.io/nvstaging/isaac-amr/gr00t_1_6_droid"
 # Droid checkpoint baked into the GR00T server image.
-DEFAULT_MODEL_PATH = "/workspace/pretrained_ckpts/GR00T-N1.6-DROID"
+MODEL_PATH = "/workspace/pretrained_ckpts/GR00T-N1.6-DROID"
 # Embodiment tag for the droid manipulation config (see droid_manip_gr00t_closedloop_config.yaml).
-DEFAULT_EMBODIMENT_TAG = "OXE_DROID"
-DEFAULT_SERVER_PORT = 5555
-
-# OSMO task name for the server. The eval runner reaches it via the {{host:<name>}} token,
-# which OSMO resolves to the runtime IP.
-GR00T_SERVER_TASK_NAME = "gr00t_server"
-GR00T_SERVER_HOST_TOKEN = "{{host:" + GR00T_SERVER_TASK_NAME + "}}"
-
-# Run in the Arena eval image (not the server image): polls the server until it accepts requests.
-WAIT_FOR_SERVER_COMMAND = "/isaac-sim/python.sh -u -m isaaclab_arena_gr00t.utils.wait_for_gr00t_server"
-
-
-def get_wait_for_server_script(host: str, port: int) -> str:
-    """Return the bash snippet that blocks until the GR00T server accepts requests."""
-    return (
-        "# Wait for the GR00T server task to come up before starting the eval.\n"
-        f"{WAIT_FOR_SERVER_COMMAND} \\\n"
-        f"  --host {host} \\\n"
-        f"  --port {port} \\\n"
-        "  --timeout-sec 1200 \\\n"
-        "  --poll-interval-sec 15 \\\n"
-        "  --request-timeout-ms 5000\n"
-    )
+EMBODIMENT_TAG = "OXE_DROID"
 
 
 class Gr00tServerTask(BaseTask):
@@ -51,27 +29,14 @@ def __init__(
         lead: bool | None = None,
     ) -> None:
         super().__init__(workflow_args=workflow_args, task_args=task_args, lead=lead)
-        self.docker_image = getattr(task_args, "gr00t_server_image", image)
-        self.model_path = task_args.gr00t_model_path
-        self.embodiment_tag = task_args.gr00t_embodiment_tag
-        self.server_port = task_args.server_port
-
-    @staticmethod
-    def add_task_arguments(parser: argparse.ArgumentParser) -> None:
-        group = parser.add_argument_group("gr00t server")
-        group.add_argument("--gr00t_server_image", default=DEFAULT_IMAGE, help="Override the GR00T server image")
-        group.add_argument("--gr00t_model_path", default=DEFAULT_MODEL_PATH, help="Model path for the GR00T policy")
-        group.add_argument(
-            "--gr00t_embodiment_tag", default=DEFAULT_EMBODIMENT_TAG, help="Embodiment tag for the GR00T policy"
-        )
-        group.add_argument("--server_port", type=int, default=DEFAULT_SERVER_PORT, help="GR00T server port")
+        self.image = image
 
     @staticmethod
     def get_task_name() -> str:
-        return GR00T_SERVER_TASK_NAME
+        return "gr00t_server"
 
     def _get_image(self) -> str:
-        return self.docker_image
+        return self.image
 
     def _get_inputs(self) -> list[dict[str, Any]]:
         return []
@@ -85,8 +50,8 @@ def _get_run_script(self) -> str:
             "nvidia-smi\n"
             "cd /workspace\n"
             "exec uv run python gr00t/eval/run_gr00t_server.py \\\n"
-            f"  --model_path={self.model_path} \\\n"
-            f"  --embodiment_tag={self.embodiment_tag} \\\n"
+            f"  --model_path={MODEL_PATH} \\\n"
+            f"  --embodiment_tag={EMBODIMENT_TAG} \\\n"
             "  --host=0.0.0.0 \\\n"
-            f"  --port={self.server_port}\n"
+            f"  --port={POLICY_SERVER_PORT}\n"
         )

osmo/tasks/pi0_remote_policy_runner_task.py

@@ -6,6 +6,7 @@
 """Remote pi0 policy-runner task for the Isaac Lab Arena OSMO workflow."""
 
 from tasks.policy_runner_task import PolicyRunnerTask
+from workflows.workflow_constants import POLICY_SERVER_PORT
 
 
 class Pi0RemotePolicyRunnerTask(PolicyRunnerTask):
@@ -18,7 +19,7 @@ def _get_policy_args(self) -> list[str]:
             "--remote_host",
             "{{host:policy_server}}",
             "--remote_port",
-            "8000",
+            str(POLICY_SERVER_PORT),
             # Raised from the default: on OSMO the first inference timed out while the
             # server was still compiling kernels, which dropped the connection.
             "--ping_timeout",

osmo/workflows/gr00t_policy_runner_workflow.py

@@ -20,6 +20,6 @@
 class Gr00tPolicyRunnerWorkflow(Workflow):
     """Two-task workflow: a GR00T server plus the lead policy-runner eval task."""
 
-    task_cls_list = [Gr00tServerTask, Gr00tPolicyRunnerTask]
     # The policy runner is the lead: it drives completion and the server runs until the lead finishes.
-    lead_list = [False, True]
+    task_cls_list = [Gr00tPolicyRunnerTask, Gr00tServerTask]
+    lead_list = [True, False]

osmo/workflows/workflow_constants.py

@@ -20,6 +20,5 @@
 DATASET_HTTPS_URL = f"{HTTPS_URL_PREFIX}/{DATASETS_PATH}/{{{{workflow_id}}}}"
 DATASET_SWIFT_URL = f"{SWIFT_URL_PREFIX}/{DATASETS_PATH}/{{{{workflow_id}}}}"
 
-# Evaluation output dataset, written per run under the existing isaaclab_arena container.
-EVAL_OUTPUT_PATH = "AUTH_team-isaac/isaaclab_arena/datasets"
-EVAL_OUTPUT_SWIFT_URL = f"{SWIFT_URL_PREFIX}/{EVAL_OUTPUT_PATH}/{{{{workflow_id}}}}"
+# Port a policy server binds to and its remote policy-runner client connects to.
+POLICY_SERVER_PORT = 8000