DIME / agents /random_agent.py
Naseer-010's picture
added reward function normalization, hidden eval set
15673b8
Raw
History Blame Contribute Delete
2.77 kB
"""Seeded random DIME baseline."""
from __future__ import annotations
import random
from typing import Any
from agents.base_agent import BaseAgent
from benchmark.utils import observation_to_dict
from server.models import InfraAction
class RandomAgent(BaseAgent):
"""Uniformly sample valid DIME management actions with deterministic seeding."""
def __init__(
self,
seed: int = 0,
action_probabilities: dict[str, float] | None = None,
) -> None:
self._base_seed = seed
self._rng = random.Random(seed)
self._action_probabilities = action_probabilities or {
"restart_node": 1.0,
"reroute_traffic": 1.0,
"throttle": 1.0,
"scale_up": 1.0,
"query_logs": 1.0,
"no_op": 1.0,
}
def reset(self, seed: int | None = None, task_id: str | None = None) -> None:
self._rng = random.Random(self._base_seed if seed is None else seed)
def act(self, observation: Any) -> InfraAction:
obs = observation_to_dict(observation)
node_count = max(1, len(obs.get("cpu_loads", []) or [0]))
actions = list(self._action_probabilities.keys())
weights = [max(0.0, float(self._action_probabilities[action])) for action in actions]
if not actions or sum(weights) <= 0.0:
actions = ["no_op"]
weights = [1.0]
action_type = self._rng.choices(actions, weights=weights, k=1)[0]
if action_type == "restart_node":
failed = list(obs.get("failed_nodes", []) or [])
target = int(self._rng.choice(failed)) if failed else self._rng.randrange(node_count)
return InfraAction(action_type="restart_node", target=target)
if action_type == "reroute_traffic" and node_count > 1:
src = self._rng.randrange(node_count)
dst_choices = [idx for idx in range(node_count) if idx != src]
return InfraAction(
action_type="reroute_traffic",
from_node=src,
to_node=self._rng.choice(dst_choices),
)
if action_type == "throttle":
return InfraAction(action_type="throttle", rate=self._rng.choice([0.3, 0.5, 0.7, 0.9]))
if action_type == "scale_up":
return InfraAction(action_type="scale_up")
if action_type == "query_logs":
telemetry = obs.get("telemetry_status", {}) or {}
timed_out = [int(idx) for idx, status in telemetry.items() if status == "timeout"]
target = self._rng.choice(timed_out) if timed_out else self._rng.randrange(node_count)
return InfraAction(action_type="query_logs", target=int(target))
return InfraAction(action_type="no_op")