feat(configurator): env_params as first-class trial-identity field

Make env_params a first-class part of CloudAIGymEnv trial identity so the
trajectory cache keys on (action, env_params) rather than action alone,
fixing the domain-randomization correctness bug where the same action under
a different env_params sample returned a stale reward.

- Cache key now includes env_params; cache-key tests pin the contract
  (formerly the TDD-red specs of NVIDIA#900, folded in here).
- Keep env.csv and trajectory.csv 1:1 step-aligned: a single TrajectoryEntry
  sinks both files coherently, including on constraint failure.
- Reject env_params on non-DSE jobs; reject non-finite / negative weights.
- Add cache-hit + declared-env_params integration coverage.

Folds the test-only PR NVIDIA#900 (cache-key TDD) into this PR so the stack has no
permanently-red standalone PR.

Author: rutayan-nv

src/cloudai/_core/test_scenario.py

@@ -97,6 +97,7 @@ class TestRun:
     reports: Set[Type[ReportGenerationStrategy]] = field(default_factory=set)
     extra_srun_args: str | None = None
     num_nodes_explicit: bool = False
+    current_env_params: dict[str, Any] = field(default_factory=dict)
 
     def __hash__(self) -> int:
         return hash(self.name + self.test.name + str(self.iterations) + str(self.current_iteration))

src/cloudai/cli/handlers.py

@@ -39,6 +39,7 @@
     System,
     TestParser,
     TestScenario,
+    TestScenarioParsingError,
 )
 from cloudai.models.scenario import ReportConfig
 from cloudai.models.workload import TestDefinition
@@ -297,12 +298,35 @@ def _check_installation(
     return result
 
 
+def validate_dse_env_params(test_scenario: TestScenario) -> None:
+    """
+    Reject prepped configs that declare env_params on a non-DSE test run.
+
+    env_params are sampled only during DSE (by CloudAIGymEnv); on a non-DSE run they would be
+    silently ignored. is_dse_job is a property of the fully prepped config, so this is validated
+    here rather than at parse time.
+    """
+    offenders = [tr.name for tr in test_scenario.test_runs if tr.test.env_params and not tr.is_dse_job]
+    if offenders:
+        raise TestScenarioParsingError(
+            f"Tests {offenders} declare env_params but are not DSE jobs. env_params are sampled only during "
+            "DSE (by CloudAIGymEnv); add a sweep (a list-valued cmd_args/extra_env_vars entry or num_nodes) "
+            "or remove env_params."
+        )
+
+
 def handle_dry_run_and_run(args: argparse.Namespace) -> int:
     setup_result = _setup_system_and_scenario(args)
     if setup_result is None:
         return 1
     system, test_scenario, tests = setup_result
 
+    try:
+        validate_dse_env_params(test_scenario)
+    except TestScenarioParsingError as e:
+        logging.error(str(e))
+        return 1
+
     if not _handle_single_sbatch(args, system):
         return 1
 
@@ -491,7 +515,8 @@ def verify_test_scenarios(
             tests = Parser.parse_tests(test_tomls, system)
             hook_tests = Parser.parse_tests(hook_test_tomls, system)
             hooks = Parser.parse_hooks(hook_tomls, system, {t.name: t for t in hook_tests})
-            Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks)
+            scenario = Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks)
+            validate_dse_env_params(scenario)
         except Exception:
             nfailed += 1
 

src/cloudai/configurator/cloudai_gym.py

@@ -19,13 +19,14 @@
 import dataclasses
 import logging
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from cloudai.core import METRIC_ERROR, BaseRunner, Registry, TestRun
 from cloudai.util.lazy_imports import lazy
 
 from .base_agent import RewardOverrides
 from .base_gym import BaseGym
+from .env_params import CsvSink, EnvParamsObserver, EnvParamsSink, StepObserver
 
 
 @dataclasses.dataclass(frozen=True)
@@ -36,6 +37,7 @@ class TrajectoryEntry:
     action: dict[str, Any]
     reward: float
     observation: list
+    env_params: dict[str, Any] = dataclasses.field(default_factory=dict)
 
 
 class CloudAIGymEnv(BaseGym):
@@ -61,8 +63,28 @@ def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrid
         self.max_steps = test_run.test.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
         self.trajectory: dict[int, list[TrajectoryEntry]] = {}
+        self._env_sink: EnvParamsSink | None = None
+        self.observers: List[StepObserver] = self._build_observers()
         super().__init__()
 
+    def _build_observers(self) -> List[StepObserver]:
+        """
+        Construct the per-step observers implied by the TestDefinition.
+
+        Workloads opt in to env_params via a TOML ``[env_params.<name>]`` block;
+        an empty mapping yields no observers and zero overhead.
+        """
+        observers: List[StepObserver] = []
+        if self.test_run.test.env_params:
+            seed = int((self.test_run.test.agent_config or {}).get("random_seed", 0))
+            self._env_sink = CsvSink(self._env_csv_path())
+            observers.append(EnvParamsObserver(self.test_run.test.env_params, seed))
+        return observers
+
+    def _env_csv_path(self) -> Path:
+        """``env.csv`` lives alongside ``trajectory.csv`` so a plain ``merge`` joins them."""
+        return self.trajectory_file_path.parent / "env.csv"
+
     def define_action_space(self) -> Dict[str, list[Any]]:
         return self.test_run.param_space
 
@@ -121,6 +143,9 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
         self.test_run.increment_step()
         self.test_run = self.test_run.apply_params_set(action)
 
+        for observer in self.observers:
+            observer.before_step(self.test_run)
+
         cached_result = self.get_cached_trajectory_result(action)
         if cached_result is not None:
             logging.info(
@@ -134,8 +159,11 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
                     action=action,
                     reward=cached_result.reward,
                     observation=cached_result.observation,
+                    env_params=dict(self.test_run.current_env_params),
                 )
             )
+            for observer in self.observers:
+                observer.after_step(self.test_run, cached_result.observation, cached_result.reward)
             return cached_result.observation, cached_result.reward, False, {}
 
         if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
@@ -162,6 +190,11 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
             self.test_run.step = new_tr.step
             self.test_run.output_path = new_tr.output_path
 
+        # Rebuilding/replacing test_run above can drop the per-trial env_params sample;
+        # restore it so the trajectory entry, the cache key, and env.csv all record the
+        # params the trial actually ran with.
+        self.test_run.current_env_params = new_tr.current_env_params
+
         observation = self.get_observation(action)
         reward = self.compute_reward(observation)
 
@@ -171,9 +204,13 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
                 action=action,
                 reward=reward,
                 observation=observation,
+                env_params=dict(self.test_run.current_env_params),
             )
         )
 
+        for observer in self.observers:
+            observer.after_step(self.test_run, observation, reward)
+
         return observation, reward, False, {}
 
     def render(self, mode: str = "human"):
@@ -230,7 +267,14 @@ def get_observation(self, action: Any) -> list:
         return observation
 
     def write_trajectory(self, entry: TrajectoryEntry):
-        """Append the trajectory to the CSV file and to the local attribute."""
+        """
+        Append the entry to the in-memory cache and trajectory.csv (plus env.csv when declared).
+
+        ``trajectory.csv`` and the ``env.csv`` projection are sunk from the same
+        ``TrajectoryEntry`` here, so a trial that never produces an entry (e.g. a
+        constraint failure returns before this call) lands in neither file and the
+        two stay 1:1 step-aligned.
+        """
         self.current_trajectory.append(entry)
 
         file_exists = self.trajectory_file_path.exists()
@@ -243,6 +287,9 @@ def write_trajectory(self, entry: TrajectoryEntry):
                 writer.writerow(["step", "action", "reward", "observation"])
             writer.writerow([entry.step, entry.action, entry.reward, entry.observation])
 
+        if self._env_sink is not None:
+            self._env_sink.write(entry.step, entry.env_params)
+
     @property
     def trajectory_file_path(self) -> Path:
         return self.runner.scenario_root / self.test_run.name / f"{self.test_run.current_iteration}" / "trajectory.csv"
@@ -252,8 +299,21 @@ def current_trajectory(self) -> list[TrajectoryEntry]:
         return self.trajectory.setdefault(self.test_run.current_iteration, [])
 
     def get_cached_trajectory_result(self, action: Any) -> TrajectoryEntry | None:
+        """
+        Return a cached entry only when the full trial identity matches.
+
+        Trial identity is ``(action, env_params)``: env-randomized parameters
+        change the workload's behaviour, so a trial repeating the same action
+        under a different ``env_params`` sample must miss and re-run. Empty
+        env_params on both sides is the back-compat path for workloads that
+        do not declare any ``[env_params.*]`` block.
+        """
+        current_env_params = getattr(self.test_run, "current_env_params", {}) or {}
         for entry in self.current_trajectory:
-            if self._values_match_exact(entry.action, action):
+            if not self._values_match_exact(entry.action, action):
+                continue
+            entry_env = getattr(entry, "env_params", {}) or {}
+            if self._values_match_exact(entry_env, current_env_params):
                 return entry
 
         return None

src/cloudai/configurator/env_params.py

@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Domain-randomization primitives for CloudAI DSE.
+
+An env-randomized parameter is a workload knob whose value the environment
+samples per trial (categorical, optional weights). It is sibling to
+``cmd_args`` on a ``TestDefinition`` and does not enter the agent's action
+space; the policy learns a robust mapping under that variation.
+
+This module owns the data schema (``EnvParamSpec``), the deterministic
+sampler (``EnvParamsSampler``), the persistence interface
+(``EnvParamsSink`` + ``CsvSink``) and the per-step observer
+(``EnvParamsObserver``). ``CloudAIGymEnv`` consumes these directly so the
+artifacts (``env.csv``) and the cache key align 1:1 with ``trajectory.csv``
+regardless of agent (PPO, BO, GA, MAB) or workload.
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+import random
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
+
+
+class EnvParamSpec(BaseModel):
+    """Specification of one env-randomized parameter (categorical)."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    values: List[Any] = Field(
+        min_length=2,
+        description="Candidate values; a single-valued parameter is just a fixed cmd_args entry.",
+    )
+    weights: Optional[List[float]] = Field(
+        default=None,
+        description="Optional probability weights aligned with values; uniform if omitted.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_weights(self) -> Self:
+        if self.weights is None:
+            return self
+        if len(self.weights) != len(self.values):
+            raise ValueError(
+                f"env_params weights length {len(self.weights)} does not match values length {len(self.values)}"
+            )
+        for w in self.weights:
+            if not math.isfinite(w) or w < 0:
+                raise ValueError(f"env_params weights must be finite and non-negative; got {w}")
+        if sum(self.weights) <= 0:
+            raise ValueError("env_params weights must have a positive sum")
+        return self
+
+
+class EnvParamsSampler:
+    """
+    Per-trial categorical sampler.
+
+    Determinism contract: ``sample(t)`` returns the same dict on every call
+    (across processes) for the same ``(seed, env_params, t)``.
+
+    Independence contract: each parameter uses an RNG seeded by
+    ``f"{seed}:{name}:{trial}"`` so adding or removing an unrelated
+    parameter does not perturb existing parameters' draw sequences.
+    """
+
+    def __init__(self, env_params: Dict[str, EnvParamSpec], seed: int) -> None:
+        self._env_params = env_params
+        self._seed = seed
+
+    def sample(self, trial: int) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        for name, spec in self._env_params.items():
+            rng = random.Random(f"{self._seed}:{name}:{trial}")
+            if spec.weights is not None:
+                out[name] = rng.choices(spec.values, weights=spec.weights, k=1)[0]
+            else:
+                out[name] = rng.choice(spec.values)
+        return out
+
+
+@runtime_checkable
+class EnvParamsSink(Protocol):
+    """Persist one trial's env_params sample; empty samples must be no-ops."""
+
+    def write(self, step: int, sample: Dict[str, Any]) -> None: ...
+
+
+class CsvSink:
+    """
+    Append per-trial env_params samples to a step-aligned CSV.
+
+    The CSV mirrors how ``trajectory.csv`` serialises its ``action`` column
+    (one row per env.step(), sample dict stringified in a single cell) so the
+    two files align 1:1 on ``step`` and a plain ``merge`` joins them.
+    """
+
+    def __init__(self, path: Path) -> None:
+        self._path = path
+
+    def write(self, step: int, sample: Dict[str, Any]) -> None:
+        if step < 1:
+            raise ValueError(f"step must be a positive trial index (cloudai DSE is 1-based); got {step}")
+        if not sample:
+            return
+        new_file = not self._path.exists()
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        with self._path.open("a", newline="") as f:
+            writer = csv.writer(f)
+            if new_file:
+                writer.writerow(("step", "env"))
+            writer.writerow([step, sample])
+
+
+@runtime_checkable
+class StepObserver(Protocol):
+    """
+    Hook fired by ``CloudAIGymEnv.step()`` around each trial.
+
+    ``before_step`` runs before the cache lookup and before any workload
+    execution. ``after_step`` runs after the trajectory row is written.
+    """
+
+    def before_step(self, test_run: Any) -> None: ...
+
+    def after_step(self, test_run: Any, observation: list, reward: float) -> None: ...
+
+
+class EnvParamsObserver:
+    """
+    Sample env_params per step and stash them for the cache and the workload.
+
+    Pre-step: samples ``test_run.test.env_params`` for ``test_run.step`` and
+    stashes the result on ``test_run.current_env_params`` so the cache key and
+    the workload's substitution both see it. Persistence is owned by
+    ``CloudAIGymEnv.write_trajectory``, which sinks ``trajectory.csv`` and the
+    ``env.csv`` projection from the single ``TrajectoryEntry`` - so a trial that
+    writes no trajectory row writes no env.csv row either, keeping the two files
+    1:1 step-aligned. Post-step: no-op.
+    """
+
+    def __init__(self, env_params: Dict[str, EnvParamSpec], seed: int) -> None:
+        self._sampler = EnvParamsSampler(env_params, seed=seed)
+
+    def before_step(self, test_run: Any) -> None:
+        test_run.current_env_params = self._sampler.sample(test_run.step)
+
+    def after_step(self, test_run: Any, observation: list, reward: float) -> None:
+        del test_run, observation, reward  # persistence handled by CloudAIGymEnv.write_trajectory