Address review: rename flags, assemble command via a helper

- Rename --env_variations to --variations (and self.variations).
- Rename the env resolver to get_arena_env_token and expand its docstring
  with registered-name and YAML examples.
- Assemble the policy_runner command in _get_policy_runner_command,
  filtering empty segments so the command never has double spaces.

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

Author: alexmillane

osmo/tasks/policy_runner_task.py

@@ -41,7 +41,7 @@ def __init__(
         self.policy_runner_args = _normalize_args(self.task_args.policy_runner_args)
         self.arena_env = self.task_args.arena_env
         self.arena_env_args = _normalize_args(self.task_args.arena_env_args)
-        self.env_variations = _normalize_args(self.task_args.env_variations)
+        self.variations = _normalize_args(self.task_args.variations)
         self.image = image
 
     @staticmethod
@@ -63,7 +63,7 @@ def add_task_arguments(parser: argparse.ArgumentParser) -> None:
             help="Env-related arguments for the chosen Arena environment",
         )
         group.add_argument(
-            "--env_variations",
+            "--variations",
             default="",
             help="Hydra-style variation overrides appended to the env, e.g. 'light.hdr_image.enabled=true'",
         )
@@ -90,20 +90,33 @@ def _get_policy_args(self) -> list[str]:
         """
 
     def _get_run_script(self) -> str:
-        policy_args_str = " ".join(self._get_policy_args())
-        # Skip empty segments so the command has no stray double spaces.
-        env_spec = " ".join(part for part in (self._get_arena_env(), self.arena_env_args, self.env_variations) if part)
-        return (
-            "set -euxo pipefail\n"
-            f"{POLICY_RUNNER_COMMAND} "
-            f"{policy_args_str} "
-            f"--output_base_dir {OSMO_TASK_OUTPUT_DIR} "
-            f"{self.policy_runner_args} "
-            f"{env_spec}\n"
-        )
-
-    def _get_arena_env(self) -> str:
-        """Render the env source: a ``--env_graph_spec_yaml`` flag for a YAML path, else the example-env name."""
+        return f"set -euxo pipefail\n{self._get_policy_runner_command()}\n"
+
+    def _get_policy_runner_command(self) -> str:
+        """Assemble the policy_runner.py command, dropping empty segments so there are no double spaces."""
+        parts = [
+            POLICY_RUNNER_COMMAND,
+            *self._get_policy_args(),
+            "--output_base_dir",
+            OSMO_TASK_OUTPUT_DIR,
+            self.policy_runner_args,
+            self.get_arena_env_token(),
+            self.arena_env_args,
+            self.variations,
+        ]
+        return " ".join(part for part in parts if part)
+
+    def get_arena_env_token(self) -> str:
+        """Render the Arena env selector token as ``policy_runner.py`` expects it.
+
+        The ``--arena_env`` value chooses the environment source and is resolved as follows:
+
+        - A registered example-environment name is passed through unchanged, e.g.
+          ``kitchen_pick_and_place`` -> ``kitchen_pick_and_place``.
+        - A graph-spec YAML path (ending in ``.yaml``/``.yml``) is wrapped in the
+          ``--env_graph_spec_yaml`` flag, e.g.
+          ``robolab/mustard_raisin_box.yaml`` -> ``--env_graph_spec_yaml robolab/mustard_raisin_box.yaml``.
+        """
         if self.arena_env.endswith((".yaml", ".yml")):
             return f"--env_graph_spec_yaml {self.arena_env}"
         return self.arena_env