Update mems_adaptive_inspection_env_curriculum_v5_SOFT_RESET_STABLE.py

mems_adaptive_inspection_env_curriculum_v5_SOFT_RESET_STABLE.py

@@ -173,15 +173,11 @@ class InspectionConfig:
     miss_penalty: float = 200.0
     max_steps: Optional[int] = 2200
     # Belief-lite parameters
-    prior_belief: float = 0.12           # Higher baseline (was 0.03)
+    prior_belief: float = 0.12
     belief_update_radius: int = 3
     belief_increase_rate: float = 0.3
-    belief_decrease_rate: float = 0.05   # Slower decay (was 0.2)
-    belief_floor: float = 0.005          # CHANGED: was 0.02 — must allow belief to drop below
-                                          # rational termination threshold of ~0.01
-                                          # (P < inspection_cost / defect_catch_value = 1/99 ≈ 0.01)
-                                          # A floor of 0.02 mathematically prevented EV from going negative,
-                                          # making rational termination impossible.
+    belief_decrease_rate: float = 0.05
+    belief_floor: float = 0.005
     real_data_path: Optional[str] = None
     preload_wafers: int = 0
     image_size: int = 64
@@ -193,7 +189,7 @@ class InspectionConfig:
     # Default False — preserves original behavior for validation and baseline training
     economic_randomization: bool = False
     clean_episode_ratio: float = 0.7   # 70% clean, 30% varied when randomization enabled
-    cost_range: tuple = (0.8, 1.2)     # Narrow range — avoids distribution shock
+    cost_range: tuple = (0.8, 1.2)     # Narrow range to avoid distribution shock
     budget_range: tuple = (0.9, 1.1)   # ±10% budget variation
     
     def __post_init__(self):
@@ -238,11 +234,10 @@ class ResolutionAgnosticInspectionEnv(gym.Env):
         self.current_budget = self.config.inspection_budget
         self.wafer_radius_mm = self.config.wafer_diameter_mm / 2.0
         
-        # Soft reset mode for production optimization
         self.soft_reset_enabled = self.config.soft_reset
         
         # EARLY TERMINATION: The terminate action is the last discrete action (N²)
-        # It must remain UNMASKED in your policy network's action mask.
+        # It must remain unmasked in your policy network's action mask.
         # The spatial action_mask (grid_size × grid_size) covers only inspection cells;
         # terminate lives outside that spatial mask and must be explicitly kept available.
         self.terminate_action = self.config.grid_size * self.config.grid_size
@@ -275,7 +270,6 @@ class ResolutionAgnosticInspectionEnv(gym.Env):
         self._init_persistent_tensors()
     
     def _init_persistent_tensors(self):
-        """Pre-allocate tensors that will be reused in soft reset mode"""
         self.wafer_map = np.zeros((self.current_grid_size, self.current_grid_size), dtype=np.int8)
         self.belief_map = np.full((self.current_grid_size, self.current_grid_size), 
                                    self.config.prior_belief, dtype=np.float32)
@@ -300,7 +294,6 @@ class ResolutionAgnosticInspectionEnv(gym.Env):
             self._init_belief_gpu_cache()
 
     def _init_belief_gpu_cache(self):
-        """Persistent GPU buffers for belief updates (no per-step allocation)."""
         grid_size = self.current_grid_size
 
         self._belief_gpu = torch.zeros(
@@ -352,7 +345,6 @@ class ResolutionAgnosticInspectionEnv(gym.Env):
         # action_space: grid_size² spatial actions + 1 terminate action
         # NOTE for policy implementors: when building your action mask tensor, do:
         #   full_mask = torch.cat([spatial_mask.flatten(), torch.ones(1, dtype=torch.bool)])
-        # The terminate action must NEVER be masked out.
         self.action_space = spaces.Discrete(self.config.grid_size * self.config.grid_size + 1)
 
     def _check_reset_invariants(self):