Distribute Arena Experiment Runs with exact outputs

isaaclab_arena/evaluation/experiment_runner.py

@@ -3,7 +3,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import os
 from pathlib import Path
 
 from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg
@@ -54,6 +53,19 @@ def _assert_camera_support_enabled(experiment_cfg: ArenaExperimentCfg, enable_ca
     )
 
 
+def _assert_exact_experiment_output_directory_is_available(experiment_output_directory: Path) -> None:
+    """Check that an exact Experiment output path is missing or empty."""
+    if experiment_output_directory.exists():
+        assert (
+            experiment_output_directory.is_dir()
+        ), f"Experiment output path exists but is not a directory: '{experiment_output_directory}'"
+        existing_experiment_output_paths = list(experiment_output_directory.iterdir())
+        assert len(existing_experiment_output_paths) == 0, (
+            f"Experiment output directory '{experiment_output_directory}' is not empty. Choose another directory,"
+            " clear it, or use --output_base_dir to create a timestamped Experiment directory."
+        )
+
+
 def main():
     args_cli, experiment_overrides = parse_experiment_runner_args()
     experiment_config_path = validate_experiment_config_path(args_cli.experiment_config)
@@ -87,9 +99,22 @@ def main():
         assert legacy_experiment_config is not None, "--chunk_size currently supports only legacy JSON Experiments"
 
         if len(legacy_experiment_config["jobs"]) > args_cli.chunk_size:
+            # TODO(alexmillane): Choose one timestamped Experiment output directory in the parent and pass that
+            # exact path to every legacy chunk worker. Each worker currently creates its own timestamped directory
+            # from --output_base_dir.
+            assert (
+                args_cli.experiment_output_directory is None
+            ), "--experiment_output_directory is not supported when --chunk_size dispatches multiple chunks"
             run_legacy_json_in_chunks(args_cli, legacy_experiment_config)
             return
 
+    if args_cli.experiment_output_directory is not None:
+        experiment_output_directory = args_cli.experiment_output_directory
+        _assert_exact_experiment_output_directory_is_available(experiment_output_directory)
+    else:
+        experiment_output_directory = Path(timestamped_run_dir(args_cli.output_base_dir))
+    experiment_output_directory.mkdir(parents=True, exist_ok=True)
+
     with SimulationAppContext(args_cli):
         experiment_cfg = load_arena_experiment_from_config_file(
             experiment_config_path,
@@ -101,19 +126,12 @@ def main():
 
         print(build_runs_info_table(experiment_cfg.runs.values(), []))
 
-        # One reverse-dated output directory for the Experiment, with one subdirectory
-        # per Run. Always date it so each invocation produces its own report directory.
-        # TODO(alexmillane): Currently each chunk produces its own output directory.
-        # We should use the same output directory for all chunks in the future.
-        experiment_output_dir = Path(timestamped_run_dir(args_cli.output_base_dir))
-
         if args_cli.record_viewport_video:
-            os.makedirs(experiment_output_dir, exist_ok=True)
-            print(f"[INFO] Video recording enabled. Videos will be saved to: {experiment_output_dir}")
+            print(f"[INFO] Video recording enabled. Videos will be saved to: {experiment_output_directory}")
 
         results = execute_experiment(
             experiment_cfg,
-            output_dir=experiment_output_dir,
+            output_dir=experiment_output_directory,
             record_viewport_video=args_cli.record_viewport_video,
             record_camera_video=args_cli.record_camera_video,
             continue_on_error=args_cli.continue_on_error,
@@ -126,7 +144,7 @@ def main():
         metrics_logger.print_metrics()
 
         # Write HTML report.
-        report_path = build_report(experiment_output_dir)
+        report_path = build_report(experiment_output_directory)
         if args_cli.serve_evaluation_report:
             serve_until_ctrl_c(report_path.parent, args_cli.evaluation_report_port, report_path.name)
 

isaaclab_arena/evaluation/experiment_runner_cli.py

@@ -4,11 +4,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import argparse
+from pathlib import Path
 
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
 from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides
 
 _DEFAULT_EXPERIMENT_CONFIG_PATH = "isaaclab_arena_environments/eval_jobs_configs/zero_action_jobs_config.json"
+_DEFAULT_EXPERIMENT_OUTPUT_BASE_DIRECTORY = "outputs"
 
 
 def add_experiment_runner_arguments(parser: argparse.ArgumentParser) -> None:
@@ -38,15 +40,32 @@ def add_experiment_runner_arguments(parser: argparse.ArgumentParser) -> None:
         default=False,
         help="Record one mp4 per (env, camera, episode) from obs['camera_obs'] for each Run.",
     )
-    parser.add_argument(
+    # Keep existing Experiment Runner commands backward compatible:
+    # --output_base_dir <base> writes to <base>/<timestamp>.
+    # OSMO workflow tasks use --experiment_output_directory <path> because each task
+    # must write directly to the exact {{output}} directory allocated by OSMO.
+    # TODO(cvolk): Replace these two path options with one path and an explicit
+    # timestamped-or-exact mode after existing --output_base_dir callers migrate.
+    output_directory_group = parser.add_mutually_exclusive_group()
+    output_directory_group.add_argument(
         "--output_base_dir",
         type=str,
-        default="outputs",
+        default=_DEFAULT_EXPERIMENT_OUTPUT_BASE_DIRECTORY,
         help=(
             "Base directory for evaluation outputs (videos, per-episode results, report); a"
             " reverse-dated Experiment subdirectory and per-Run subdirectory are added."
         ),
     )
+    output_directory_group.add_argument(
+        "--experiment_output_directory",
+        type=Path,
+        default=None,
+        help=(
+            "Exact directory that will contain this Experiment's report and one subdirectory per Run."
+            " The directory must be missing or empty. Managed execution can use this instead of a timestamped"
+            " directory."
+        ),
+    )
     parser.add_argument(
         "--serve_evaluation_report",
         action="store_true",

isaaclab_arena/tests/test_arena_experiment_config_loader.py

@@ -169,6 +169,43 @@ def test_experiment_runner_rejects_yaml_chunking_before_starting_simulation(monk
         experiment_runner.main()
 
 
+def test_experiment_runner_rejects_exact_output_for_multiple_legacy_chunks(monkeypatch, tmp_path):
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "experiment_runner.py",
+            "--experiment_config",
+            str(GETTING_STARTED_JSON_PATH),
+            "--chunk_size",
+            "1",
+            "--experiment_output_directory",
+            str(tmp_path / "exact-experiment-output"),
+        ],
+    )
+
+    with pytest.raises(AssertionError, match="not supported when --chunk_size dispatches multiple chunks"):
+        experiment_runner.main()
+
+
+def test_experiment_runner_rejects_nonempty_exact_output_before_starting_simulation(monkeypatch, tmp_path):
+    exact_experiment_output_directory = tmp_path / "existing-experiment-output"
+    exact_experiment_output_directory.mkdir()
+    (exact_experiment_output_directory / "existing-result.jsonl").write_text("{}\n", encoding="utf-8")
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "experiment_runner.py",
+            "--experiment_config",
+            str(GETTING_STARTED_YAML_PATH),
+            "--experiment_output_directory",
+            str(exact_experiment_output_directory),
+        ],
+    )
+
+    with pytest.raises(AssertionError, match="is not empty.*--output_base_dir"):
+        experiment_runner.main()
+
+
 def test_legacy_json_experiment_rejects_hydra_overrides():
     with pytest.raises(AssertionError, match="only for typed YAML"):
         load_arena_experiment_from_config_file(

isaaclab_arena/tests/test_experiment_runner.py

@@ -33,6 +33,44 @@ def test_experiment_runner_parses_native_hydra_overrides():
     ]
 
 
+def test_experiment_runner_parses_timestamped_base_or_exact_output_directory(tmp_path):
+    exact_experiment_output_directory = tmp_path / "exact-experiment-output"
+    timestamped_experiment_output_base_directory = tmp_path / "timestamped-experiment-outputs"
+
+    default_arguments, default_experiment_overrides = parse_experiment_runner_args([
+        "--experiment_config",
+        "experiment.yaml",
+    ])
+    assert default_arguments.output_base_dir == "outputs"
+    assert default_arguments.experiment_output_directory is None
+    assert default_experiment_overrides == []
+
+    timestamped_output_arguments, timestamped_output_experiment_overrides = parse_experiment_runner_args([
+        "--output_base_dir",
+        str(timestamped_experiment_output_base_directory),
+    ])
+    assert timestamped_output_arguments.output_base_dir == str(timestamped_experiment_output_base_directory)
+    assert timestamped_output_arguments.experiment_output_directory is None
+    assert timestamped_output_experiment_overrides == []
+
+    exact_output_arguments, exact_output_experiment_overrides = parse_experiment_runner_args([
+        "--experiment_config",
+        "experiment.yaml",
+        "--experiment_output_directory",
+        str(exact_experiment_output_directory),
+    ])
+    assert exact_output_arguments.experiment_output_directory == exact_experiment_output_directory
+    assert exact_output_experiment_overrides == []
+
+    with pytest.raises(SystemExit):
+        parse_experiment_runner_args([
+            "--output_base_dir",
+            str(tmp_path / "timestamped-outputs"),
+            "--experiment_output_directory",
+            str(exact_experiment_output_directory),
+        ])
+
+
 @pytest.mark.with_subprocess
 def test_experiment_runner_rejects_unknown_non_hydra_arguments():
     """Reject misspelled CLI flags in a fresh process."""
@@ -113,7 +151,7 @@ def test_experiment_runner_from_typed_yaml(tmp_path):
         str(experiment_config_path),
         config_option="--experiment_config",
         extra_args=[
-            "--output_base_dir",
+            "--experiment_output_directory",
             str(tmp_path / "output"),
             "runs.yaml_baseline.rollout_limit.num_steps=2",
         ],
@@ -123,6 +161,8 @@ def test_experiment_runner_from_typed_yaml(tmp_path):
     run_row = next(line for line in result.stdout.splitlines() if "yaml_baseline" in line and "pending" in line)
     run_cells = [cell.strip() for cell in run_row.split("|")[1:-1]]
     assert run_cells[4] == "2"
+    assert (tmp_path / "output/index.html").is_file()
+    assert (tmp_path / "output/yaml_baseline/episode_results_rebuild0.jsonl").is_file()
 
 
 @pytest.mark.with_subprocess

isaaclab_arena/tests/test_osmo_build_experiment_output.py

@@ -0,0 +1,97 @@
+# 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
+
+"""Verify building one Experiment output from exact Experiment Runner task outputs."""
+
+import json
+from pathlib import Path
+
+import pytest
+
+from osmo.scripts.build_experiment_output import (
+    build_experiment_output,
+    collect_run_outputs_into_experiment_output,
+    load_experiment_runner_output_directories_by_run_name,
+)
+
+
+def _write_run_output(run_output_directory: Path, run_name: str, success: bool) -> None:
+    run_output_directory.mkdir(parents=True)
+    episode_result = {
+        "job_name": run_name,
+        "env_id": 0,
+        "episode_in_env": 0,
+        "success": success,
+    }
+    (run_output_directory / "episode_results_rebuild0.jsonl").write_text(
+        json.dumps(episode_result) + "\n",
+        encoding="utf-8",
+    )
+
+
+def test_loads_experiment_runner_output_directories_as_paths(tmp_path):
+    experiment_runner_output_directories_file_path = tmp_path / "experiment-runner-output-directories.json"
+    experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+    experiment_runner_output_directories_file_path.write_text(
+        json.dumps({"first": str(experiment_runner_output_directory)}),
+        encoding="utf-8",
+    )
+
+    experiment_runner_output_directories_by_run_name = load_experiment_runner_output_directories_by_run_name(
+        experiment_runner_output_directories_file_path
+    )
+
+    assert experiment_runner_output_directories_by_run_name == {"first": experiment_runner_output_directory}
+
+
+def test_rejects_experiment_runner_output_without_the_requested_run(tmp_path):
+    experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+    (experiment_runner_output_directory / "another-run").mkdir(parents=True)
+
+    with pytest.raises(AssertionError, match="Expected Run output directory for Run 'first'.*first"):
+        collect_run_outputs_into_experiment_output(
+            {"first": experiment_runner_output_directory},
+            tmp_path / "experiment-output",
+        )
+
+
+def test_collects_run_outputs_without_building_report(tmp_path):
+    experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+    _write_run_output(experiment_runner_output_directory / "first", "first", True)
+    experiment_output_directory = tmp_path / "experiment-output"
+
+    collect_run_outputs_into_experiment_output(
+        {"first": experiment_runner_output_directory},
+        experiment_output_directory,
+    )
+
+    assert (experiment_output_directory / "first/episode_results_rebuild0.jsonl").is_file()
+    assert not (experiment_output_directory / "index.html").exists()
+
+
+def test_builds_experiment_output_from_separate_experiment_runner_outputs(tmp_path):
+    first_experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+    second_experiment_runner_output_directory = tmp_path / "experiment-runner-1-output"
+    first_run_output_directory = first_experiment_runner_output_directory / "first"
+    second_run_output_directory = second_experiment_runner_output_directory / "second"
+    _write_run_output(first_run_output_directory, "first", True)
+    _write_run_output(second_run_output_directory, "second", False)
+    experiment_output_directory = tmp_path / "experiment-output"
+
+    report_path = build_experiment_output(
+        {
+            "first": first_experiment_runner_output_directory,
+            "second": second_experiment_runner_output_directory,
+        },
+        experiment_output_directory,
+    )
+
+    assert report_path == experiment_output_directory / "index.html"
+    assert (experiment_output_directory / "first/episode_results_rebuild0.jsonl").is_file()
+    assert (experiment_output_directory / "second/episode_results_rebuild0.jsonl").is_file()
+    report_contents = report_path.read_text(encoding="utf-8")
+    assert "first" in report_contents
+    assert "second" in report_contents
+    assert "2 job(s)" in report_contents