ui and inference changes

deploy/grafana/provisioning/dashboards/json/antiatropos-overview.json

@@ -76,8 +76,8 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_reward",
-          "legendFormat": "reward",
+          "expr": "scalar(avg(last_over_time(antiatropos_reward{mode=\"simulated\"}[1m])))",
+          "legendFormat": "reward (simulated)",
           "range": true,
           "refId": "A"
         }
@@ -143,8 +143,8 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_total_queue_backlog",
-          "legendFormat": "queue backlog",
+          "expr": "scalar(avg(last_over_time(antiatropos_total_queue_backlog{mode=\"simulated\"}[1m])))",
+          "legendFormat": "queue backlog (simulated)",
           "range": true,
           "refId": "A"
         }
@@ -210,8 +210,8 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_average_latency_norm",
-          "legendFormat": "latency",
+          "expr": "scalar(avg(last_over_time(antiatropos_average_latency_norm{mode=\"simulated\"}[1m])))",
+          "legendFormat": "latency (simulated)",
           "range": true,
           "refId": "A"
         }
@@ -277,8 +277,8 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_lyapunov_energy",
-          "legendFormat": "lyapunov energy",
+          "expr": "scalar(avg(last_over_time(antiatropos_lyapunov_energy{mode=\"simulated\"}[1m])))",
+          "legendFormat": "lyapunov energy (simulated)",
           "range": true,
           "refId": "A"
         }
@@ -369,15 +369,15 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_reward",
-          "legendFormat": "reward {{task_id}}",
+          "expr": "antiatropos_reward{mode=\"simulated\"}",
+          "legendFormat": "reward {{task_id}} ({{mode}})",
           "range": true,
           "refId": "A"
         },
         {
           "editorMode": "code",
-          "expr": "antiatropos_lyapunov_energy",
-          "legendFormat": "lyapunov {{task_id}}",
+          "expr": "antiatropos_lyapunov_energy{mode=\"simulated\"}",
+          "legendFormat": "lyapunov {{task_id}} ({{mode}})",
           "range": true,
           "refId": "B"
         }
@@ -468,15 +468,15 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "antiatropos_total_queue_backlog",
-          "legendFormat": "queue {{task_id}}",
+          "expr": "antiatropos_total_queue_backlog{mode=\"simulated\"}",
+          "legendFormat": "queue {{task_id}} ({{mode}})",
           "range": true,
           "refId": "A"
         },
         {
           "editorMode": "code",
-          "expr": "antiatropos_average_latency_norm",
-          "legendFormat": "latency {{task_id}}",
+          "expr": "antiatropos_average_latency_norm{mode=\"simulated\"}",
+          "legendFormat": "latency {{task_id}} ({{mode}})",
           "range": true,
           "refId": "B"
         }
@@ -535,15 +535,15 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "rate(antiatropos_steps_total[1m])",
-          "legendFormat": "steps/sec {{task_id}}",
+          "expr": "sum by (task_id, mode) (rate(antiatropos_steps_total{mode=\"simulated\"}[1m]))",
+          "legendFormat": "steps/sec {{task_id}} ({{mode}})",
           "range": true,
           "refId": "A"
         },
         {
           "editorMode": "code",
-          "expr": "rate(antiatropos_actions_total[1m])",
-          "legendFormat": "actions/sec {{action_type}}",
+          "expr": "sum by (task_id, mode, action_type) (rate(antiatropos_actions_total{mode=\"simulated\"}[1m]))",
+          "legendFormat": "actions/sec {{action_type}} ({{task_id}}, {{mode}})",
           "range": true,
           "refId": "B"
         }
@@ -602,14 +602,14 @@
       "targets": [
         {
           "editorMode": "code",
-          "expr": "rate(antiatropos_executor_errors_total[5m])",
-          "legendFormat": "executor errors {{error_code}}",
+          "expr": "sum by (mode, error_code) (rate(antiatropos_executor_errors_total{mode=\"simulated\"}[5m]))",
+          "legendFormat": "executor errors {{error_code}} ({{mode}})",
           "range": true,
           "refId": "A"
         },
         {
           "editorMode": "code",
-          "expr": "histogram_quantile(0.95, sum(rate(antiatropos_executor_latency_ms_bucket[5m])) by (le, mode))",
+          "expr": "histogram_quantile(0.95, sum(rate(antiatropos_executor_latency_ms_bucket{mode=\"simulated\"}[5m])) by (le, mode))",
           "legendFormat": "p95 executor latency {{mode}}",
           "range": true,
           "refId": "B"
@@ -637,6 +637,6 @@
   "timezone": "browser",
   "title": "AntiAtropos Overview",
   "uid": "antiatropos-overview",
-  "version": 1,
+  "version": 2,
   "weekStart": ""
 }

inference.py

@@ -1,50 +1,57 @@
 import asyncio
 import json
 import os
+import random
 import textwrap
-from typing import List, Optional
+import time
+from contextlib import asynccontextmanager
+from typing import Dict, List
+from urllib.parse import urlparse
 
 from dotenv import load_dotenv
-
-load_dotenv()  # Load variables from .env file
-
 from openai import AsyncOpenAI
 
 from AntiAtropos.client import AntiAtroposEnv
 from AntiAtropos.grader import EpisodeGrader
 from AntiAtropos.models import ActionType, SREAction
 
+load_dotenv()
 
 API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.1-70b-versatile")
+MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.1-8b-instant")
 API_KEY = (
-    os.getenv("HF_TOKEN")
+    os.getenv("GROQ_API_KEY")      # prioritize Groq key since we default to groq API
     or os.getenv("OPENAI_API_KEY")
     or os.getenv("API_KEY")
-    or os.getenv("GROQ_API_KEY")
+    or os.getenv("HF_TOKEN")
 )
-LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
-ENV_URL = os.getenv("ANTIATROPOS_ENV_URL", "http://127.0.0.1:8000")
-TASK_NAME = os.getenv("ANTIATROPOS_TASK", "task-3")
-BENCHMARK = os.getenv("ANTIATROPOS_BENCHMARK", "antiatropos")
+
+ENV_URL = os.getenv("ANTIATROPOS_ENV_URL", "https://pranavkk-antiatropos.hf.space")
 ENV_MODE = os.getenv("ANTIATROPOS_MODE", "simulated")
-MAX_STEPS = int(os.getenv("ANTIATROPOS_MAX_STEPS", "35"))
-TEMPERATURE = float(os.getenv("ANTIATROPOS_TEMPERATURE", "0.05"))
+TASKS = ["task-1", "task-2", "task-3"]
+
+TOTAL_BUDGET_SECONDS = 1080  # 18-minute limit
+MIN_TASK_BUDGET_SECONDS = 60
+MAX_STEPS_PER_TASK = 60       # 60 steps = ~5 minutes at this rate
+MESSAGE_TIMEOUT_S = 300
+MODEL_TIMEOUT_S = 25
+
+TEMPERATURE = float(os.getenv("ANTIATROPOS_TEMPERATURE", "0.0"))
 MAX_TOKENS = int(os.getenv("ANTIATROPOS_MAX_TOKENS", "180"))
+SEED = int(os.getenv("ANTIATROPOS_SEED", "42"))
 SUCCESS_SCORE_THRESHOLD = float(os.getenv("ANTIATROPOS_SUCCESS_THRESHOLD", "0.55"))
 
+TASK_BRIEFS: Dict[str, str] = {
+    "task-1": "Traffic increases linearly. Scale proactively to keep latency low and cost efficient.",
+    "task-2": "A node fails randomly. Detect quickly and recover with reroute/scale actions.",
+    "task-3": "Protect VIP node-0 under surges. Keep VIP healthy without invalid actions.",
+}
+
 SYSTEM_PROMPT = textwrap.dedent(
     """
     You are an autonomous SRE controller managing a five-node microservice cluster.
 
-    Objectives:
-    - minimize Lyapunov energy and queue growth
-    - keep normalized average latency at or below 0.20
-    - avoid invalid actions, especially SHED_LOAD on node-0, node-1, and node-2
-    - scale proactively because SCALE_UP takes 5 ticks to take effect
-    - protect the VIP gateway node-0
-
-    Output exactly one JSON object:
+    Return exactly one JSON object:
     {
       "action_type": "SCALE_UP" | "SCALE_DOWN" | "REROUTE_TRAFFIC" | "SHED_LOAD" | "NO_OP",
       "target_node_id": "node-0" | "node-1" | "node-2" | "node-3" | "node-4",
@@ -54,57 +61,94 @@ SYSTEM_PROMPT = textwrap.dedent(
 ).strip()
 
 
-def log_start(task: str, env: str, model: str) -> None:
-    print(f"[START] task={task} env={env} model={model}", flush=True)
+def _seed_everything(seed: int) -> None:
+    random.seed(seed)
+    try:
+        import numpy as np
 
+        np.random.seed(seed)
+    except Exception:
+        pass
 
-def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
-    error_val = error if error else "null"
-    done_val = str(done).lower()
-    print(
-        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
-        flush=True,
-    )
 
+def _task_seed(base_seed: int, task_id: str) -> int:
+    offsets = {"task-1": 0, "task-2": 1, "task-3": 2}
+    return int(base_seed + offsets.get(task_id, 0))
 
-def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    rewards_str = ",".join(f"{reward:.2f}" for reward in rewards)
-    print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}",
-        flush=True,
-    )
 
+def _hf_web_fallback_url(base_url: str) -> str:
+    parsed = urlparse(base_url)
+    host = parsed.netloc.lower()
+    path = parsed.path.rstrip("/")
+    if host.endswith(".hf.space") and path == "":
+        return base_url.rstrip("/") + "/web"
+    return base_url
 
-def build_user_prompt(step: int, obs: dict, history: List[str]) -> str:
-    history_block = "\n".join(history[-4:]) if history else "None"
+
+@asynccontextmanager
+async def open_env_with_ws_fallback(base_url: str, message_timeout_s: int):
+    try:
+        async with AntiAtroposEnv(base_url, message_timeout_s=message_timeout_s) as env:
+            yield env
+            return
+    except ConnectionError as e:
+        fallback_url = _hf_web_fallback_url(base_url)
+        if fallback_url == base_url or "404" not in str(e):
+            raise
+        print(f"[connect] ws 404 on {base_url}; retrying with {fallback_url}", flush=True)
+        async with AntiAtroposEnv(fallback_url, message_timeout_s=message_timeout_s) as env:
+            yield env
+
+
+def build_user_prompt(task_id: str, step: int, obs: dict, history: List[str]) -> str:
+    recent = "\n".join(history[-4:]) if history else "None"
+    brief = TASK_BRIEFS.get(task_id, "Maintain SLA, stability, and efficient cost.")
     return textwrap.dedent(
         f"""
+        Task: {task_id}
+        Objective: {brief}
         Step: {step}
+
         Current state:
         {json.dumps(obs, separators=(",", ":"))}
 
         Recent decisions:
-        {history_block}
+        {recent}
 
         Choose the next SRE action.
         """
     ).strip()
 
 
-def compact_action(action: SREAction) -> str:
-    payload = {
-        "action_type": action.action_type.value,
-        "target_node_id": action.target_node_id,
-        "parameter": round(float(action.parameter), 4),
+def observation_for_model(obs) -> dict:
+    return {
+        "task_id": obs.task_id,
+        "mode": getattr(obs.mode, "value", str(obs.mode)),
+        "step": obs.step,
+        "max_steps": obs.max_steps,
+        "lyapunov_energy": obs.lyapunov_energy,
+        "average_latency_ms": obs.average_latency_ms,
+        "error_rate": obs.error_rate,
+        "total_queue_backlog": obs.total_queue_backlog,
+        "sla_violations": obs.sla_violations,
+        "invalid_action_count": obs.invalid_action_count,
+        "nodes": [
+            {
+                "node_id": node.node_id,
+                "status": getattr(node.status, "value", str(node.status)),
+                "is_vip": node.is_vip,
+                "queue_depth": node.queue_depth,
+                "latency_ms": node.latency_ms,
+                "incoming_request_rate": node.incoming_request_rate,
+                "cpu_utilization": node.cpu_utilization,
+            }
+            for node in obs.nodes
+        ],
     }
-    return json.dumps(payload, separators=(",", ":"))
 
 
-def extract_json_object(text: str) -> dict:
+def _extract_json_object(text: str) -> dict:
     stripped = text.strip()
-    if not stripped:
-        raise ValueError("empty model response")
-
     start = stripped.find("{")
     end = stripped.rfind("}")
     if start == -1 or end == -1 or end < start:
@@ -112,7 +156,7 @@ def extract_json_object(text: str) -> dict:
     return json.loads(stripped[start : end + 1])
 
 
-def parse_action(payload: dict) -> SREAction:
+def _parse_action(payload: dict) -> SREAction:
     action_type = str(payload.get("action_type", "NO_OP")).upper()
     target_node_id = str(payload.get("target_node_id", "node-0"))
     parameter = float(payload.get("parameter", 0.0))
@@ -123,139 +167,168 @@ def parse_action(payload: dict) -> SREAction:
     )
 
 
-async def get_model_action(
-    client: AsyncOpenAI,
-    step: int,
-    obs: dict,
-    history: List[str],
-) -> SREAction:
-    user_prompt = build_user_prompt(step, obs, history)
+async def get_model_action(client: AsyncOpenAI, task_id: str, step: int, obs: dict, history: List[str]) -> SREAction:
+    prompt = build_user_prompt(task_id=task_id, step=step, obs=obs, history=history)
     try:
         completion = await client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
+                {"role": "user", "content": prompt},
             ],
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
             response_format={"type": "json_object"},
+            timeout=MODEL_TIMEOUT_S,
+            seed=SEED,
         )
         content = completion.choices[0].message.content or ""
-        return parse_action(extract_json_object(content))
-    except Exception:
-        return SREAction(
-            action_type=ActionType.NO_OP,
-            target_node_id="node-0",
-            parameter=0.0,
-        )
+        return _parse_action(_extract_json_object(content))
+    except Exception as e:
+        print(f"[LLM_ERROR] task={task_id} step={step} error={type(e).__name__}: {e}", flush=True)
+        return SREAction(action_type=ActionType.NO_OP, target_node_id="node-0", parameter=0.0)
 
 
-def observation_for_model(obs) -> dict:
-    return {
-        "task_id": obs.task_id,
-        "mode": getattr(obs.mode, "value", str(obs.mode)),
-        "step": obs.step,
-        "max_steps": obs.max_steps,
-        "lyapunov_energy": obs.lyapunov_energy,
-        "average_latency_ms": obs.average_latency_ms,
-        "error_rate": obs.error_rate,
-        "total_queue_backlog": obs.total_queue_backlog,
-        "sla_violations": obs.sla_violations,
-        "invalid_action_count": obs.invalid_action_count,
-        "nodes": [
-            {
-                "node_id": node.node_id,
-                "status": getattr(node.status, "value", str(node.status)),
-                "is_vip": node.is_vip,
-                "queue_depth": node.queue_depth,
-                "latency_ms": node.latency_ms,
-                "incoming_request_rate": node.incoming_request_rate,
-                "cpu_utilization": node.cpu_utilization,
-            }
-            for node in obs.nodes
-        ],
+def _compact_action(action: SREAction) -> str:
+    payload = {
+        "action_type": action.action_type.value,
+        "target_node_id": action.target_node_id,
+        "parameter": round(float(action.parameter), 4),
     }
+    return json.dumps(payload, separators=(",", ":"))
 
 
-async def run_episode() -> None:
+async def run_single_task(env: AntiAtroposEnv, client: AsyncOpenAI, task_id: str, deadline: float) -> dict:
+    start = time.monotonic()
+    task_seed = _task_seed(SEED, task_id)
+    result = await env.reset(task_id=task_id, mode=ENV_MODE, seed=task_seed)
+
+    grader = EpisodeGrader(task_id=task_id)
+    grader.record(result.observation)
+    history: List[str] = []
     rewards: List[float] = []
     steps_taken = 0
-    success = False
-    score = 0.0
-    env = None
-    client = None
+    timed_out = False
+
+    for step in range(1, MAX_STEPS_PER_TASK + 1):
+        if time.monotonic() >= deadline:
+            timed_out = True
+            break
+        if result.done:
+            break
+
+        action = await get_model_action(
+            client=client,
+            task_id=task_id,
+            step=step,
+            obs=observation_for_model(result.observation),
+            history=history,
+        )
+        result = await env.step(action)
+        grader.record(result.observation)
 
-    log_start(TASK_NAME, BENCHMARK, MODEL_NAME)
+        reward = float(result.reward or 0.0)
+        rewards.append(reward)
+        steps_taken = step
+        ack = getattr(result.observation, "action_ack_status", "")
+        action_str = _compact_action(action)
+        history.append(f"step={step} action={action_str} reward={reward:.4f} ack={ack or 'null'}")
+
+        error = ack if ack.startswith(("Rejected:", "Error:")) else None
+        print(
+            f"[STEP] task={task_id} step={step} action={action_str} reward={reward:.4f} done={str(result.done).lower()} error={error or 'null'}",
+            flush=True,
+        )
 
-    try:
-        if not API_KEY:
-            raise RuntimeError("missing API key")
+    grade = grader.score()
+    score = max(0.0, min(1.0, float(grade.composite)))
+    elapsed = time.monotonic() - start
+    success = score >= SUCCESS_SCORE_THRESHOLD and not timed_out
+    print(
+        f"[TASK_END] task={task_id} success={str(success).lower()} score={score:.4f} "
+        f"steps={steps_taken} elapsed_s={elapsed:.1f} timed_out={str(timed_out).lower()} seed={task_seed}",
+        flush=True,
+    )
+    return {
+        "task_id": task_id,
+        "success": success,
+        "score": score,
+        "steps": steps_taken,
+        "elapsed_seconds": elapsed,
+        "timed_out": timed_out,
+        "grade_summary": grade.summary(),
+        "rewards": rewards,
+    }
 
-        client = AsyncOpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-        if LOCAL_IMAGE_NAME:
-            env = await AntiAtroposEnv.from_docker_image(LOCAL_IMAGE_NAME)
-        else:
-            env = AntiAtroposEnv(base_url=ENV_URL)
-            await env.__aenter__()
 
-        grader = EpisodeGrader(task_id=TASK_NAME)
-        history: List[str] = []
+async def run_all_tasks() -> None:
+    _seed_everything(SEED)
+    tasks = [task for task in TASKS if task in {"task-1", "task-2", "task-3"}]
+    if not tasks:
+        raise RuntimeError("ANTIATROPOS_TASKS must include at least one of: task-1,task-2,task-3")
+    if not API_KEY:
+        raise RuntimeError("Missing API key (HF_TOKEN/OPENAI_API_KEY/API_KEY/GROQ_API_KEY).")
 
-        result = await env.reset(task_id=TASK_NAME, mode=ENV_MODE)
-        grader.record(result.observation)
+    print(
+        f"[START] tasks={','.join(tasks)} env={ENV_URL} mode={ENV_MODE} model={MODEL_NAME} "
+        f"budget_s={TOTAL_BUDGET_SECONDS} seed={SEED}",
+        flush=True,
+    )
 
-        for step in range(1, MAX_STEPS + 1):
-            if result.done:
-                break
-
-            obs = result.observation
-            action = await get_model_action(
-                client=client,
-                step=step,
-                obs=observation_for_model(obs),
-                history=history,
-            )
-            result = await env.step(action)
-            grader.record(result.observation)
-
-            reward = float(result.reward or 0.0)
-            rewards.append(reward)
-            steps_taken = step
-
-            ack_status = getattr(result.observation, "action_ack_status", "")
-            error = ack_status if ack_status.startswith(("Rejected:", "Error:")) else None
-            action_str = compact_action(action)
-            log_step(step=step, action=action_str, reward=reward, done=result.done, error=error)
-
-            history.append(
-                f"step={step} action={action_str} reward={reward:.2f} ack={ack_status or 'null'}"
-            )
-
-            if result.done:
-                break
-
-        score = max(0.0, min(1.0, grader.score().composite))
-        success = score >= SUCCESS_SCORE_THRESHOLD
-    except Exception as e:
-        print(f"[CRITICAL ERROR] Episode failed to initialize: {e}")
-        success = False
+    start = time.monotonic()
+    deadline = start + TOTAL_BUDGET_SECONDS
+    reports: List[dict] = []
+
+    client = AsyncOpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    try:
+        async with open_env_with_ws_fallback(ENV_URL, MESSAGE_TIMEOUT_S) as env:
+            for idx, task_id in enumerate(tasks):
+                now = time.monotonic()
+                if now >= deadline:
+                    print(f"[BUDGET] stopping before {task_id}; time budget exhausted", flush=True)
+                    break
+
+                remaining_tasks = len(tasks) - idx
+                remaining_seconds = max(0.0, deadline - now)
+                allocated_seconds = max(
+                    float(MIN_TASK_BUDGET_SECONDS),
+                    remaining_seconds / float(remaining_tasks),
+                )
+                task_deadline = min(deadline, now + allocated_seconds)
+                print(
+                    f"[BUDGET] task={task_id} allocated_s={allocated_seconds:.1f} "
+                    f"remaining_s={remaining_seconds:.1f} remaining_tasks={remaining_tasks}",
+                    flush=True,
+                )
+
+                report = await run_single_task(
+                    env=env,
+                    client=client,
+                    task_id=task_id,
+                    deadline=task_deadline,
+                )
+                reports.append(report)
     finally:
-        if client is not None:
-            try:
-                await client.close()
-            except Exception:
-                pass
-        if env is not None:
-            try:
-                await env.close()
-            except Exception:
-                pass
-        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+        await client.close()
+
+    total_elapsed = time.monotonic() - start
+    completed_scores = [r["score"] for r in reports]
+    aggregate_score = sum(completed_scores) / len(completed_scores) if completed_scores else 0.0
+    aggregate_score = max(0.0, min(1.0, aggregate_score))
+    all_success = len(reports) == len(tasks) and all(r["success"] for r in reports)
+
+    for report in reports:
+        print(f"[GRADE] {report['grade_summary']}", flush=True)
+
+    print(
+        f"[END] success={str(all_success).lower()} completed_tasks={len(reports)}/{len(tasks)} "
+        f"aggregate_score={aggregate_score:.4f} elapsed_s={total_elapsed:.1f}",
+        flush=True,
+    )
 
 
 def main() -> None:
-    asyncio.run(run_episode())
+    asyncio.run(run_all_tasks())
 
 
 if __name__ == "__main__":

server/AntiAtropos_environment.py

@@ -79,7 +79,7 @@ class AntiAtroposEnvironment(Environment):
             self._reward_output_mode = "normalized"
         self._last_metric_time: float = 0.0
 
-    def reset(self, task_id: str = "task-1", mode: str = "simulated") -> ClusterObservation:
+    def reset(self, task_id: str = "task-1", mode: str = "simulated", seed: int | None = None) -> ClusterObservation:
         """
         Start a fresh episode with a specific task profile and mode.
         """
@@ -110,7 +110,7 @@ class AntiAtroposEnvironment(Environment):
             # self._telemetry = PrometheusClient(url=os.getenv("PROMETHEUS_URL"))
             pass
 
-        self._sim.reset(task_id=task_id)
+        self._sim.reset(task_id=task_id, seed=seed)
         
         # If in hybrid mode, immediately pull a baseline
         if self._mode in [EnvironmentMode.HYBRID, EnvironmentMode.LIVE]:
@@ -396,4 +396,3 @@ class AntiAtroposEnvironment(Environment):
             reward=0.0,
         )
 
-

simulator.py

@@ -133,6 +133,7 @@ class ClusterSimulator:
         # Default to non-deterministic RNG seeding so fresh simulator instances
         # do not replay identical domain-randomization sequences.
         # Pass an explicit seed for reproducible experiments.
+        self._seed: Optional[int] = seed
         self._rng = random.Random(seed)
         self._tick_count: int = 0
         self._failed_node_id: Optional[str] = None
@@ -179,8 +180,12 @@ class ClusterSimulator:
             for i in range(self._n_nodes)
         ]
 
-    def reset(self, task_id: str = "task-1") -> None:
+    def reset(self, task_id: str = "task-1", seed: Optional[int] = None) -> None:
         """Restart the simulator for a fresh episode."""
+        if seed is not None:
+            self._seed = seed
+            # Reinitialize RNG so episode generation is reproducible for a given seed.
+            self._rng = random.Random(seed)
         self._task_id = task_id
         self._tick_count = 0
         self._failed_node_id = None
@@ -482,4 +487,3 @@ class ClusterSimulator:
         # We'll skip it if we just reconciled to keep the blended values, OR refine it.
         # For now, let's just make sure statuses are updated based on new queue depths.
         self._update_statuses()
-