🌱 [MRG] [exp] Fix seed

Merge pull request #46 from MLSysOps/fix-seed

Author: YuanmingLeee

exps/abr_sim.py

@@ -322,14 +322,13 @@ def train_agent(
     memory,
     max_messages,
     summary_every: int = 100,
-    base_seed: int | None = None,
 ):
     ep = 0
     global_step = 0
     episode_reward = 0.0
     step_logs = []
     # Single-stream seeding: seed only the first reset, then let RNG continue
-    obs, info = env.reset(seed=base_seed)
+    obs, info = env.reset()
 
     system_prompt = f"{obs['task_description']}\n\n{memory.as_context_block()}\n\n"
     chat_history = [{"role": "system", "content": system_prompt}]
@@ -543,12 +542,14 @@ def run_llm_agent(
         max_episode_steps=1000_000,
         reward_scale=reward_scale,
         dummy_mode=not real_env,
+        seed=seed,
     )
     test_env = infragym.make(
         'abr_sim',
         max_episode_steps=test_steps,
         reward_scale=reward_scale,
         dummy_mode=not real_env,
+        seed=seed + 10_000 if seed is not None else None,  # Stable test seed
     )
 
     # Environment summary table
@@ -568,6 +569,7 @@ def run_llm_agent(
             "test_steps": test_steps,
             "summary_every": summary_every,
             "reward_scale": reward_scale,
+            "seed": seed,
             "real_env": real_env,
             "model_name": model_name,
             "thinking_mode": thinking_mode,
@@ -590,7 +592,6 @@ def run_llm_agent(
         memory=memory,
         max_messages=max_messages,
         summary_every=summary_every,  # summarize every 100 steps
-        base_seed=seed,               # seed only the first reset
     ):
         # ---- Structured Episode Summary (printed + stored) ----
         ep_summary = result["ep_summary"]
@@ -679,6 +680,7 @@ def run_rule_policy(
         max_episode_steps=test_steps,
         reward_scale=reward_scale,
         dummy_mode=not real_env,
+        seed=seed,
     )
 
     # Environment summary table
@@ -702,7 +704,7 @@ def run_rule_policy(
     )
     console.print(Panel(f"Trace Space: {trace_space.location}", title="Tracing", border_style="blue"))
 
-    obs, info = env.reset(seed=seed)
+    obs, info = env.reset()
     ep = 0  # No episodes in rule policy, just a single test run
     test_steps = 0
     test_total_reward = 0.0
@@ -806,6 +808,7 @@ def run_llm(
         max_episode_steps=test_steps,
         reward_scale=reward_scale,
         dummy_mode=not real_env,
+        seed=seed,
     )
 
     # Environment summary table
@@ -824,14 +827,15 @@ def run_llm(
             "test_steps": test_steps,
             "reward_scale": reward_scale,
             "real_env": real_env,
+            "seed": seed,
             "model_name": model_name,
             "thinking_mode": thinking_mode,
             "max_messages": max_messages,
         }
     )
     console.print(Panel(f"Trace Space: {trace_space.location}", title="Tracing", border_style="blue"))
 
-    obs, info = env.reset(seed=seed)
+    obs, info = env.reset()
     ep = 0  # No episodes in rule policy, just a single test run
     test_steps = 1
     test_total_reward = 0.0

exps/load_balance_sim.py

@@ -412,6 +412,7 @@ def run_llm_agent(
         num_servers=num_servers,
         max_episode_steps=max_steps,
         reward_scale=reward_scale,
+        seed=seed,
     )
 
     # Environment summary
@@ -446,7 +447,7 @@ def run_llm_agent(
     for ep in range(episodes):
         console.rule(f"Episode {ep + 1}")
         # Single-stream seeding: only the FIRST reset gets the seed
-        obs, info = env.reset(seed=seed if ep == 0 else None)
+        obs, info = env.reset()
 
         # Build system prompt with task + memory
         system_prompt = (
@@ -603,6 +604,7 @@ def run_rule_policy(
         max_episode_steps=test_steps,
         reward_scale=reward_scale,
         num_servers=num_servers,
+        seed=seed,
     )
 
     # Environment summary
@@ -628,7 +630,7 @@ def run_rule_policy(
     )
     console.print(Panel(f"Trace Space: {trace_space.location}", title="Tracing", border_style="blue"))
 
-    obs, info = env.reset(seed=seed)
+    obs, info = env.reset()
     ep = 0  # No episodes in rule policy, just a single test run
     test_steps = 0
     test_total_reward = 0.0
@@ -785,6 +787,7 @@ def run_llm(
         max_episode_steps=test_steps,
         reward_scale=reward_scale,
         num_servers=num_servers,
+        seed=seed,
     )
 
     # Environment summary
@@ -814,7 +817,7 @@ def run_llm(
     # Tracing space for this run
     console.print(Panel(f"Trace Space: {trace_space.location}", title="Tracing", border_style="blue"))
 
-    obs, info = env.reset(seed=seed)
+    obs, info = env.reset()
     ep = 0
     chat_history = [{"role": "system", "content": obs["task_description"]}]
     steps = 1

infragym/abr_sim/abr_llm_gym.py

@@ -32,9 +32,6 @@ def __init__(
         seed: Optional[int] = 42,
     ):
         super().__init__()
-        self.seed = seed
-        if self.seed is not None:
-            np.random.seed(self.seed)
 
         self.max_episode_steps = max_episode_steps
         self.enable_llm_friendly_obs = enable_llm_friendly_obs
@@ -57,11 +54,11 @@ def __init__(
         self.past_download_times = deque(maxlen=self.past_chunk_window)
 
         # Simulated network conditions (simplified for demo)
-        self.network_conditions = self._generate_network_traces()
+        self.network_conditions: np.ndarray = None  # Generated / loaded on reset
         self.current_network_idx = 0
 
         # Video chunk sizes (simplified)
-        self.chunk_sizes = self._generate_chunk_sizes()
+        self.chunk_sizes: np.ndarray = None  # Generated / loaded on reset
 
         # Setup spaces
         self.action_space = gym.spaces.Discrete(len(self.bitrate_levels))
@@ -87,7 +84,7 @@ def __init__(
         Each chunk is 4 seconds of video content.
         """
 
-        self.reset()
+        self.reset(seed=seed)
 
     def _generate_network_traces(self) -> np.ndarray:
         """Generate simplified network bandwidth traces for simulation."""
@@ -97,7 +94,7 @@ def _generate_network_traces(self) -> np.ndarray:
             sigma = 0.6  # multiplicative variability
 
             mu = np.log(median_mbps)
-            x = np.random.normal(mu, sigma, size=2000)
+            x = self.np_random.normal(mu, sigma, size=2000)
             bw = np.exp(x)
             # AR(1) smoothing
             for t in range(1, 2000):
@@ -107,7 +104,7 @@ def _generate_network_traces(self) -> np.ndarray:
             return traces
         else:
             all_traces = np.load(TRACE_PATH)
-            traces = all_traces[np.random.choice(len(all_traces))]
+            traces = all_traces[self.np_random.choice(len(all_traces))]
             return traces[1]
 
     def _generate_chunk_sizes(self) -> np.ndarray:
@@ -119,7 +116,7 @@ def _generate_chunk_sizes(self) -> np.ndarray:
                 # Approximate: bitrate * chunk_duration / 8 (bytes)
                 base_size = bitrate * self.chunk_duration * 1e6 / 8
                 # Add some variation
-                variations = np.random.normal(1.0, 0.1, 100)
+                variations = self.np_random.normal(1.0, 0.1, 100)
                 sizes = base_size * variations
                 chunk_sizes.append(sizes)
             return np.array(chunk_sizes)
@@ -202,6 +199,12 @@ def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None) -> T
         self.past_download_times.clear()
         self.current_network_idx = 0
 
+        # Generate network conditions and chunk sizes
+        if self.network_conditions is None:
+            self.network_conditions = self._generate_network_traces()
+        if self.chunk_sizes is None:
+            self.chunk_sizes = self._generate_chunk_sizes()
+
         # Initialize past observations
         for _ in range(self.past_chunk_window):
             self.past_throughputs.append(0.0)

infragym/load_balance/load_balance_llm_gym.py

@@ -35,43 +35,31 @@ def __init__(
         arrival_rate: Optional[float] = None,
     ):
         super().__init__()
-        self.seed = seed
-        if self.seed is not None:
-            np.random.seed(self.seed)
+        self.utilization_target = utilization_target
+        self.auto_scale_arrivals = auto_scale_arrivals
+        self._arrival_rate_arg = arrival_rate
 
         self.num_servers = num_servers
         self.max_episode_steps = max_episode_steps
         self.enable_llm_friendly_obs = enable_llm_friendly_obs
         self.reward_scale = reward_scale
 
         # Load balancing parameters
-        self.service_rates = self._generate_service_rates()
+        self.service_rates = None
+        self.arrival_rate = None
         self.job_size_range = (1, 10)  # Job sizes in arbitrary units (smaller for faster processing)
 
-        # compute E[size] for your job distribution (Pareto a=2, xm=1, truncated at b)
-        def expected_job_size(xm=1.0, b=10.0) -> float:
-            # closed form for α=2, truncated at b: E[X | X≤b] = 2*xm / (1 + xm/b)
-            return 2.0 * xm / (1.0 + xm / b)
-
-        if auto_scale_arrivals and arrival_rate is None:
-            e_s = expected_job_size(*self.job_size_range)
-            cap = float(sum(self.service_rates))  # size units / time
-            self.arrival_rate = max(1e-6, utilization_target * cap / e_s)
-        else:
-            # Jobs per time unit (higher rate for more jobs)
-            self.arrival_rate = arrival_rate if arrival_rate is not None else 2.0
-
         # Environment state
         self.current_step = 0
         self.current_time = 0.0
-        self.servers = self._initialize_servers()
+        self.servers = []  # Init later in reset
         self.job_queue = deque()
         self.finished_jobs = []
         self.total_waiting_time = 0.0
         self.total_processing_time = 0.0
 
         # Job generation
-        self.next_job_arrival = self._generate_next_arrival()
+        self.next_job_arrival = None  # Set in reset
 
         # Setup spaces
         self.action_space = gym.spaces.Discrete(num_servers)
@@ -97,7 +85,7 @@ def expected_job_size(xm=1.0, b=10.0) -> float:
         Each server has different processing capabilities (service rates).
         """
 
-        self.reset()
+        self.reset(seed=seed)
 
     def _generate_service_rates(self) -> List[float]:
         """Generate service rates for servers (jobs per time unit)."""
@@ -107,7 +95,7 @@ def _generate_service_rates(self) -> List[float]:
         for i in range(self.num_servers):
             # Add some variation based on server ID
             variation = 0.3 * np.sin(i * np.pi / self.num_servers)
-            rate = base_rate + variation + np.random.normal(0, 0.2)
+            rate = base_rate + variation + self.np_random.normal(0, 0.2)
             rates.append(max(0.5, rate))  # Ensure minimum rate
         return rates
 
@@ -129,11 +117,11 @@ def _initialize_servers(self) -> List[Dict]:
 
     def _generate_next_arrival(self) -> float:
         """Generate next job arrival time using exponential distribution."""
-        return self.current_time + np.random.exponential(1.0 / self.arrival_rate)
+        return self.current_time + self.np_random.exponential(1.0 / self.arrival_rate)
 
     def _generate_job(self) -> Dict:
         """Generate a new job with size and arrival time."""
-        size = np.random.pareto(2.0) + 1  # Pareto distribution for job sizes
+        size = self.np_random.pareto(2.0) + 1  # Pareto distribution for job sizes
         size = min(size, self.job_size_range[1])  # Cap maximum size
 
         return {
@@ -249,6 +237,21 @@ def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None) -> T
         """Reset the environment."""
         super().reset(seed=seed)
 
+        # Generate service rates ONCE (first episode) after seeding → fixed across episodes
+        if self.service_rates is None:
+            self.service_rates = self._generate_service_rates()
+
+        # arrival rate can be recomputed each reset (deterministic given service_rates)
+        def expected_job_size(xm=1.0, b=10.0) -> float:
+            return 2.0 * xm / (1.0 + xm / b)
+
+        if self.auto_scale_arrivals and self._arrival_rate_arg is None:
+            e_s = expected_job_size(1.0, 10.0)
+            cap = float(sum(self.service_rates))
+            self.arrival_rate = max(1e-6, self.utilization_target * cap / e_s)
+        else:
+            self.arrival_rate = self._arrival_rate_arg if self._arrival_rate_arg is not None else 2.0
+
         self.current_step = 0
         self.current_time = 0.0
         self.servers = self._initialize_servers()