Agent params

src/cloudai/cli/handlers.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,11 +21,12 @@
 import signal
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Callable, List, Optional
+from typing import Any, Callable, List, Optional
 from unittest.mock import Mock
 
 import toml
 import yaml
+from pydantic import ValidationError
 
 from cloudai.core import (
     BaseInstaller,
@@ -40,6 +41,11 @@
     TestParser,
     TestScenario,
 )
+from cloudai.models.agent_config import (
+    BayesianOptimizationConfig,
+    GeneticAlgorithmConfig,
+    MultiArmedBanditConfig,
+)
 from cloudai.models.scenario import ReportConfig
 from cloudai.models.workload import TestDefinition
 from cloudai.parser import HOOK_ROOT
@@ -145,7 +151,19 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int:
             continue
 
         env = CloudAIGymEnv(test_run=test_run, runner=runner.runner)
-        agent = agent_class(env)
+
+        try:
+            agent_overrides = validate_agent_overrides(agent_type, test_run.test.agent_config)
+        except ValidationError as e:
+            logging.error(f"Invalid agent_config for agent '{agent_type}':")
+            for error in e.errors():
+                field = ".".join(str(loc) for loc in error["loc"])
+                logging.error(f"  - {field}: {error['msg']}")
+            err = 1
+            continue
+
+        agent = agent_class(env, **agent_overrides) if agent_overrides else agent_class(env)
+
         for step in range(agent.max_steps):
             result = agent.select_action()
             if result is None:
@@ -166,6 +184,31 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int:
     return err
 
 
+def validate_agent_overrides(agent_type: str, agent_config: Optional[dict[str, Any]]) -> dict[str, Any]:
+    """Validate and process agent configuration overrides."""
+    if not agent_config:
+        return {}
+
+    config_class_map = {
+        "ga": GeneticAlgorithmConfig,
+        "bo": BayesianOptimizationConfig,
+        "mab": MultiArmedBanditConfig,
+    }
+
+    config_class = config_class_map.get(agent_type)
+    if not config_class:
+        logging.debug(f"No config validation available for agent type '{agent_type}', using defaults.")
+        return {}
+
+    validated_config = config_class.model_validate(agent_config)
+    agent_kwargs = validated_config.model_dump(exclude_none=True)
+
+    if agent_kwargs:
+        logging.info(f"Applying agent config overrides for '{agent_type}': {agent_kwargs}")
+
+    return agent_kwargs
+
+
 def generate_reports(system: System, test_scenario: TestScenario, result_dir: Path) -> None:
     registry = Registry()
 

src/cloudai/models/agent_config.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2024-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.
+
+from abc import ABC
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AgentConfig(BaseModel, ABC):
+    """Base configuration for agent overrides."""
+
+    model_config = ConfigDict(extra="forbid")
+    random_seed: Optional[int] = Field(default=None, description="Random seed for reproducibility")
+
+
+class GeneticAlgorithmConfig(AgentConfig):
+    """Configuration overrides for Genetic Algorithm agent."""
+
+    population_size: Optional[int] = Field(default=None, ge=2, description="Population size for the genetic algorithm")
+    n_offsprings: Optional[int] = Field(default=None, ge=1, description="Number of offsprings per generation")
+    crossover_prob: Optional[float] = Field(default=None, ge=0.0, le=1.0, description="Crossover probability")
+    mutation_prob: Optional[float] = Field(default=None, ge=0.0, le=1.0, description="Mutation probability")
+
+
+class BayesianOptimizationConfig(AgentConfig):
+    """Configuration overrides for Bayesian Optimization agent."""
+
+    sobol_num_trials: Optional[int] = Field(default=None, ge=1, description="Number of SOBOL initialization trials")
+    botorch_num_trials: Optional[int] = Field(
+        default=None, description="Number of BoTorch trials (-1 for unlimited until max_steps)"
+    )
+
+
+class MultiArmedBanditConfig(AgentConfig):
+    """Configuration overrides for Multi-Armed Bandit agent."""
+
+    algorithm: Optional[str] = Field(
+        default=None,
+        description="MAB algorithm: ucb1, ts (thompson_sampling), epsilon_greedy, softmax, or random",
+    )
+    algorithm_params: Optional[dict[str, Any]] = Field(
+        default=None, description="Algorithm-specific parameters (e.g., alpha for UCB1, epsilon for epsilon_greedy)"
+    )
+    seed_parameters: Optional[dict[str, Any]] = Field(
+        default=None, description="Initial seed configuration to evaluate first"
+    )
+    max_arms: Optional[int] = Field(default=None, ge=1, description="Maximum number of arms in the action space")
+    warm_start_size: Optional[int] = Field(
+        default=None, ge=0, description="Number of arms to randomly explore initially"
+    )
+    epsilon_override: Optional[float] = Field(
+        default=None, ge=0.0, le=1.0, description="Epsilon value for exploration (overrides algorithm epsilon)"
+    )
+    max_explore_steps: Optional[int] = Field(
+        default=None, ge=0, description="Maximum steps for epsilon exploration (None for unlimited)"
+    )
+    prefer_unseen_random: Optional[bool] = Field(
+        default=None, description="Prefer unseen arms during random exploration (epsilon)"
+    )

src/cloudai/models/workload.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -107,6 +107,7 @@ class TestDefinition(BaseModel, ABC):
     agent_steps: int = 1
     agent_metrics: list[str] = Field(default=["default"])
     agent_reward_function: str = "inverse"
+    agent_config: Optional[dict[str, Any]] = None
 
     @property
     def cmd_args_dict(self) -> Dict[str, Union[str, List[str]]]: