Update fair_adversarial_validation_framework.py

fair_adversarial_validation_framework.py

@@ -14,11 +14,11 @@ from collections import defaultdict
 import matplotlib.pyplot as plt
 import seaborn as sns
 
-# Add this import for the wrapper
 from gru_env_wrappers import GRUStateManager
 
+
 # ============================================================================
-# TEST CASE DEFINITIONS
+# DATA CLASSES
 # ============================================================================
 
 @dataclass
@@ -32,7 +32,8 @@ class FairTestCase:
     env_config_modifier: Optional[Any] = None
     expected_behavior: str = ""
     pass_threshold: float = 0.80
-    
+
+
 @dataclass
 class FairTestResult:
     """Results from a fair test case"""
@@ -51,89 +52,51 @@ class FairTestResult:
 
 
 # ============================================================================
-# FAIR PERTURBATION GENERATORS
+# PERTURBATION GENERATORS
 # ============================================================================
 
 class FairPerturbations:
     """Realistic perturbations that respect model assumptions"""
-    
+
     @staticmethod
     def realistic_sensor_noise(observation: Dict, noise_std: float = 0.01) -> Dict:
-        """
-        Add small Gaussian noise simulating realistic sensor imperfections.
-        
-        noise_std=0.01 represents ~1% measurement uncertainty
-        This is what you'd see from real metrology equipment.
-        """
         obs = observation.copy()
         belief_map = obs['belief_map'].copy()
-        
-        # Only add noise to non-zero regions (actual wafer)
         wafer_mask = obs.get('wafer_map', np.ones_like(belief_map)) > 0
-        
         noise = np.random.normal(0, noise_std, belief_map.shape)
         belief_map = belief_map + (noise * wafer_mask)
-        
-        # Preserve probability semantics
         obs['belief_map'] = np.clip(belief_map, 0.0, 1.0)
         return obs
-    
+
     @staticmethod
     def calibration_drift(observation: Dict, drift_factor: float = 0.05) -> Dict:
-        """
-        Simulate systematic calibration drift (e.g., tool aging).
-        
-        drift_factor=0.05 means beliefs are systematically 5% off.
-        This represents gradual tool degradation.
-        """
         obs = observation.copy()
         belief_map = obs['belief_map'].copy()
-        
-        # Systematic scaling (not random)
         drift = 1.0 + np.random.uniform(-drift_factor, drift_factor)
         belief_map = belief_map * drift
-        
         obs['belief_map'] = np.clip(belief_map, 0.0, 1.0)
         return obs
-    
+
     @staticmethod
     def local_degradation(observation: Dict, affected_ratio: float = 0.1) -> Dict:
-        """
-        Simulate localized tool degradation affecting part of wafer.
-        
-        affected_ratio=0.1 means 10% of wafer has degraded sensing.
-        This represents edge-of-wafer effects or local contamination.
-        """
         obs = observation.copy()
         belief_map = obs['belief_map'].copy()
         H, W = belief_map.shape
-        
-        # Create localized degradation zone (e.g., one quadrant)
         if np.random.random() < 0.5:
-            # Edge degradation
             margin = int(H * 0.1)
-            belief_map[:margin, :] *= 0.8  # 20% reduced sensitivity
+            belief_map[:margin, :] *= 0.8
             belief_map[-margin:, :] *= 0.8
         else:
-            # Quadrant degradation
             belief_map[:H//2, :W//2] *= 0.85
-        
         obs['belief_map'] = np.clip(belief_map, 0.0, 1.0)
         return obs
-    
+
     @staticmethod
     def quantization_noise(observation: Dict, bits: int = 8) -> Dict:
-        """
-        Simulate ADC quantization (realistic for real sensors).
-        
-        bits=8 means 256 discrete levels (standard ADC resolution).
-        """
         obs = observation.copy()
         belief_map = obs['belief_map'].copy()
-        
         levels = 2 ** bits
         quantized = np.round(belief_map * levels) / levels
-        
         obs['belief_map'] = quantized
         return obs
 
@@ -144,68 +107,64 @@ class FairPerturbations:
 
 class FairEnvModifiers:
     """Realistic environment modifications"""
-    
+
     @staticmethod
     def fab_variation_defect_rate(base_rate: float = 0.03, variation: float = 0.3):
-        """
-        Simulate normal fab variation in defect rate.
-        
-        variation=0.3 means ±30% from baseline
-        Example: 3% baseline → 2.1% to 3.9% range
-        """
         return base_rate * (1.0 + np.random.uniform(-variation, variation))
-    
+
     @staticmethod
     def budget_efficiency_test(base_budget: int, efficiency: float = 0.8):
-        """
-        Test with reduced budget (simulating faster throughput requirement).
-        
-        efficiency=0.8 means 80% of normal budget (20% faster needed)
-        """
         return int(base_budget * efficiency)
-    
+
     @staticmethod
     def cost_pressure(base_cost: float, multiplier: float = 1.5):
-        """
-        Simulate cost pressure (inspection became more expensive).
-        
-        multiplier=1.5 means 50% cost increase
-        """
         return base_cost * multiplier
 
 
 # ============================================================================
-# FAIR ADVERSARIAL TEST SUITE
+# FAIR ADVERSARIAL TEST SUITE (with resume support)
 # ============================================================================
 
 class FairAdversarialTestSuite:
-    """Fair, realistic adversarial validation"""
-    
-    def __init__(self, model, env_factory, output_dir: str = "./fair_adversarial_results"):
-        """
-        Args:
-            model: Trained SB3 model
-            env_factory: Function that creates fresh environment (critical!)
-            output_dir: Where to save results
-        """
+    """Fair, realistic adversarial validation with resume capability"""
+
+    def __init__(self, model, env_factory, output_dir: str = "./fair_adversarial_results", start_test: int = 1):
         self.model = model
-        self.env_factory = env_factory  # Function, not instance!
+        self.env_factory = env_factory
         self.output_dir = Path(output_dir)
         self.output_dir.mkdir(exist_ok=True, parents=True)
-        
+        self.start_test = max(1, start_test)
+
         self.test_cases = self._define_fair_tests()
-        self.results = []
-    
+        self.results: List[FairTestResult] = self._load_progress()
+
+        print(f"✅ FairAdversarialTestSuite initialized with {len(self.test_cases)} tests. "
+              f"Starting from test #{self.start_test}")
+
+    def _load_progress(self) -> List[FairTestResult]:
+        progress_file = self.output_dir / "validation_progress.json"
+        if progress_file.exists():
+            try:
+                with open(progress_file) as f:
+                    data = json.load(f)
+                last = data.get("last_completed_test", 0)
+                print(f"✅ Found previous progress. Last completed test: {last}")
+            except Exception as e:
+                print(f"Warning: Could not read progress file: {e}")
+        return []
+
+    def _save_progress(self, last_completed: int):
+        progress_file = self.output_dir / "validation_progress.json"
+        with open(progress_file, 'w') as f:
+            json.dump({
+                "last_completed_test": last_completed,
+                "timestamp": time.time()
+            }, f, indent=2)
+
     def _define_fair_tests(self) -> List[FairTestCase]:
-        """Define fair, realistic test cases"""
-        
         tests = []
-        
-        # =================================================================
-        # PRODUCTION LEVEL: What you'd see in normal fab operation
-        # Expected: 90-98% performance
-        # =================================================================
-        
+
+        # ==================== PRODUCTION ====================
         tests.extend([
             FairTestCase(
                 name="baseline_clean",
@@ -231,10 +190,7 @@ class FairAdversarialTestSuite:
                 description="±10% defect rate variation",
                 category="distribution_robustness",
                 difficulty="production",
-                env_config_modifier=lambda config: {
-                    **config,
-                    'prior_belief': 0.1
-                },
+                env_config_modifier=lambda config: {**config, 'prior_belief': 0.1},
                 expected_behavior="Handle normal fab variation",
                 pass_threshold=0.90
             ),
@@ -260,12 +216,8 @@ class FairAdversarialTestSuite:
                 pass_threshold=0.90
             ),
         ])
-        
-        # =================================================================
-        # STRESS LEVEL: Challenging but realistic scenarios
-        # Expected: 75-90% performance
-        # =================================================================
-        
+
+        # ==================== STRESS ====================
         tests.extend([
             FairTestCase(
                 name="sensor_noise_2pct",
@@ -281,10 +233,7 @@ class FairAdversarialTestSuite:
                 description="±30% defect rate variation",
                 category="distribution_robustness",
                 difficulty="stress",
-                env_config_modifier=lambda config: {
-                    **config,
-                    'prior_belief': 0.1
-                },
+                env_config_modifier=lambda config: {**config, 'prior_belief': 0.1},
                 expected_behavior="Adapt to significant fab shifts",
                 pass_threshold=0.75
             ),
@@ -331,12 +280,8 @@ class FairAdversarialTestSuite:
                 pass_threshold=0.85
             ),
         ])
-        
-        # =================================================================
-        # EXTREME LEVEL: Edge cases and breaking points
-        # Expected: 60-75% performance (degradation is acceptable)
-        # =================================================================
-        
+
+        # ==================== EXTREME ====================
         tests.extend([
             FairTestCase(
                 name="sensor_noise_5pct",
@@ -352,10 +297,7 @@ class FairAdversarialTestSuite:
                 description="2x normal defect rate",
                 category="distribution_robustness",
                 difficulty="extreme",
-                env_config_modifier=lambda config: {
-                    **config,
-                    'prior_belief': 0.06
-                },
+                env_config_modifier=lambda config: {**config, 'prior_belief': 0.06},
                 expected_behavior="Adapt to crisis scenario",
                 pass_threshold=0.60
             ),
@@ -386,236 +328,177 @@ class FairAdversarialTestSuite:
                 pass_threshold=0.60
             ),
         ])
-        
+
         return tests
-    
+
     def run_test_case(self, test_case: FairTestCase, num_episodes: int = 30) -> FairTestResult:
-        """Run a single fair test case"""
-        
         print(f"\n{'='*80}")
         print(f"Running: {test_case.name}")
         print(f"Category: {test_case.category} | Difficulty: {test_case.difficulty}")
         print(f"Description: {test_case.description}")
         print(f"Pass Threshold: {test_case.pass_threshold:.2%}")
         print(f"{'='*80}")
-        
-        catch_rates = []
-        rewards = []
-        step_counts = []
-        
+
+        catch_rates, rewards, step_counts = [], [], []
+
         for episode in range(num_episodes):
             if episode == 0 or (episode + 1) % 10 == 0:
                 print(f" Starting episode {episode + 1}/{num_episodes}...")
-            # CREATE FRESH ENVIRONMENT (critical!)
+
             env = self.env_factory()
-            
-            # Apply GRUStateManager if not already in factory (defensive)
             if not isinstance(env, GRUStateManager):
                 env = GRUStateManager(env, policy=self.model.policy)
-            
-            # Apply environment config modifications if specified
+
             if test_case.env_config_modifier:
-                # Get modified config
                 base_config = {
-                    'grid_size': env.unwrapped.config.grid_size,
-                    'inspection_budget': env.unwrapped.config.inspection_budget,
-                    'inspection_cost': env.unwrapped.config.inspection_cost,
-                    'prior_belief': env.unwrapped.config.prior_belief,
+                    'grid_size': getattr(env.unwrapped.config, 'grid_size', 64),
+                    'inspection_budget': getattr(env.unwrapped.config, 'inspection_budget', 3000),
+                    'inspection_cost': getattr(env.unwrapped.config, 'inspection_cost', 1.0),
+                    'prior_belief': getattr(env.unwrapped.config, 'prior_belief', 0.1),
                 }
-                modified_config = test_case.env_config_modifier(base_config)
-                
-                # Apply modifications
-                for key, value in modified_config.items():
+                modified = test_case.env_config_modifier(base_config)
+                for key, value in modified.items():
                     if hasattr(env.unwrapped.config, key):
                         setattr(env.unwrapped.config, key, value)
                     if key == 'inspection_budget':
                         env.unwrapped.current_budget = value
-            
-            # Reset environment (wrapper handles GRU reset internally)
+
             obs, info = env.reset()
             last_catch_rate = 0.0
-            
-            episode_reward = 0
+            episode_reward = 0.0
             steps = 0
             done = False
             MAX_STEPS = 5000
+
             while not done and steps < MAX_STEPS:
-                # Apply perturbation if specified
                 if test_case.perturbation_fn:
                     obs = test_case.perturbation_fn(obs)
-    
-                # Predict (GRUStateManager handles GRU state internally)
+
                 action, _ = self.model.predict(obs, deterministic=True)
-    
                 if steps == 0:
                     print(f" [DEBUG] First action: {action}, type: {type(action)}")
-    
-                # Step environment
+
                 obs, reward, terminated, truncated, info = env.step(action)
-    
                 episode_reward += reward
                 steps += 1
                 last_catch_rate = info.get('catch_rate', last_catch_rate)
                 done = terminated or truncated
-    
+
                 if steps % 100 == 0:
                     print(f" [Step {steps}] budget_left={info.get('remaining_budget', '?')}")
-            
+
             if steps >= MAX_STEPS:
                 print(f" ⚠️ Episode hit max steps ({MAX_STEPS})")
+
             if episode == 0 or (episode + 1) % 10 == 0:
                 print(f" ✓ Episode {episode + 1} complete: {steps} steps, reward={episode_reward:.2f}")
-            
-            # Record results - read from info BEFORE soft_reset clears detected_defects
-            catch_rate = last_catch_rate  # Use tracked value, not post-reset info
-            
-            catch_rates.append(catch_rate)
+
+            catch_rates.append(last_catch_rate)
             rewards.append(episode_reward)
             step_counts.append(steps)
             env.close()
-            
+
             if (episode + 1) % 10 == 0:
-                print(f" Episode {episode+1}/{num_episodes} - "
-                      f"Catch Rate: {catch_rate:.3f}, Reward: {episode_reward:.1f}")
-        
+                print(f" Episode {episode+1}/{num_episodes} - Catch Rate: {last_catch_rate:.3f}, Reward: {episode_reward:.1f}")
+
         # Compute statistics
-        avg_catch_rate = np.mean(catch_rates)
-        std_catch_rate = np.std(catch_rates)
-        min_catch_rate = np.min(catch_rates)
-        max_catch_rate = np.max(catch_rates)
-        avg_reward = np.mean(rewards)
-        avg_steps = np.mean(step_counts)
-        
-        pass_status = avg_catch_rate >= test_case.pass_threshold
-        
+        avg_catch = float(np.mean(catch_rates))
+        std_catch = float(np.std(catch_rates))
+        min_catch = float(np.min(catch_rates))
+        max_catch = float(np.max(catch_rates))
+        avg_reward = float(np.mean(rewards))
+        avg_steps = float(np.mean(step_counts))
+
+        pass_status = avg_catch >= test_case.pass_threshold
+
         result = FairTestResult(
             test_name=test_case.name,
             category=test_case.category,
             difficulty=test_case.difficulty,
-            catch_rate=avg_catch_rate,
+            catch_rate=avg_catch,
             avg_reward=avg_reward,
             avg_steps=avg_steps,
             pass_status=pass_status,
-            std_catch_rate=std_catch_rate,
-            min_catch_rate=min_catch_rate,
-            max_catch_rate=max_catch_rate,
+            std_catch_rate=std_catch,
+            min_catch_rate=min_catch,
+            max_catch_rate=max_catch,
             pass_threshold=test_case.pass_threshold
         )
-        
+
         status = "✅ PASS" if pass_status else "❌ FAIL"
-        print(f"\n{status} - Catch Rate: {avg_catch_rate:.3f} "
-              f"(threshold: {test_case.pass_threshold:.3f})")
-        print(f"Stats: μ={avg_catch_rate:.3f}, σ={std_catch_rate:.3f}, "
-              f"min={min_catch_rate:.3f}, max={max_catch_rate:.3f}")
-        
+        print(f"\n{status} - Catch Rate: {avg_catch:.3f} (threshold: {test_case.pass_threshold:.3f})")
+        print(f"Stats: μ={avg_catch:.3f}, σ={std_catch:.3f}, min={min_catch:.3f}, max={max_catch:.3f}")
+
         return result
-    
+
     def run_all_tests(self, num_episodes_per_test: int = 30):
-        """Run all fair test cases"""
-        
         print(f"\n{'#'*80}")
         print(f"FAIR ADVERSARIAL VALIDATION TEST SUITE")
         print(f"Total Tests: {len(self.test_cases)}")
         print(f"Episodes per Test: {num_episodes_per_test}")
+        print(f"Starting from Test #{self.start_test}")
         print(f"{'#'*80}\n")
-        
+
         start_time = time.time()
-        
+
         for i, test_case in enumerate(self.test_cases, 1):
+            if i < self.start_test:
+                print(f"⏭️  Skipping Test {i}/{len(self.test_cases)}: {test_case.name}")
+                continue
+
             print(f"\n[Test {i}/{len(self.test_cases)}]")
             result = self.run_test_case(test_case, num_episodes_per_test)
             self.results.append(result)
-        
+
+            self._save_progress(i)   # Save progress after every completed test
+
         elapsed = time.time() - start_time
-        
         print(f"\n{'#'*80}")
         print(f"FAIR ADVERSARIAL VALIDATION COMPLETE")
         print(f"Total Time: {elapsed/60:.1f} minutes")
         print(f"{'#'*80}\n")
-        
+
         self._generate_summary()
         self._save_results()
         self._generate_visualizations()
-    
+
     def _generate_summary(self):
-        """Generate comprehensive summary"""
-        
         print(f"\n{'='*80}")
         print("FAIR ADVERSARIAL VALIDATION SUMMARY")
         print(f"{'='*80}\n")
-        
-        # Overall statistics
+
+        if not self.results:
+            print("No results to summarize.")
+            return
+
         total = len(self.results)
         passed = sum(1 for r in self.results if r.pass_status)
-        pass_rate = passed / total if total > 0 else 0
+        pass_rate = passed / total
         avg_catch = np.mean([r.catch_rate for r in self.results])
-        
+
         print(f"Overall Pass Rate: {pass_rate:.1%} ({passed}/{total})")
-        print(f"Average Catch Rate: {avg_catch:.3f}")
-        print()
-        
-        # By difficulty level
-        print("Performance by Difficulty:")
-        print("-" * 80)
-        
+        print(f"Average Catch Rate: {avg_catch:.3f}\n")
+
         for difficulty in ["production", "stress", "extreme"]:
             diff_results = [r for r in self.results if r.difficulty == difficulty]
             if not diff_results:
                 continue
-            
-            diff_passed = sum(1 for r in diff_results if r.pass_status)
-            diff_total = len(diff_results)
-            diff_pass_rate = diff_passed / diff_total
-            diff_avg_catch = np.mean([r.catch_rate for r in diff_results])
-            
-            status = "✅" if diff_pass_rate >= 0.8 else "⚠️" if diff_pass_rate >= 0.5 else "❌"
-            print(f"{status} {difficulty.upper():12s} | "
-                  f"Pass: {diff_pass_rate:5.1%} ({diff_passed}/{diff_total}) | "
-                  f"Avg Catch: {diff_avg_catch:.3f}")
-        
-        print()
-        
-        # By category
-        print("Performance by Category:")
-        print("-" * 80)
-        
-        categories = defaultdict(list)
-        for r in self.results:
-            categories[r.category].append(r)
-        
-        for category, results in sorted(categories.items()):
-            cat_avg = np.mean([r.catch_rate for r in results])
-            cat_passed = sum(1 for r in results if r.pass_status)
-            cat_total = len(results)
-            
-            print(f"{category:25s} | Avg Catch: {cat_avg:.3f} | "
-                  f"Pass: {cat_passed}/{cat_total}")
-        
-        print()
-        
-        # Failed tests
-        failed = [r for r in self.results if not r.pass_status]
-        if failed:
-            print("Failed Tests:")
-            print("-" * 80)
-            for r in failed:
-                print(f"❌ {r.test_name:30s} | "
-                      f"Catch: {r.catch_rate:.3f} (need {r.pass_threshold:.3f}) | "
-                      f"{r.difficulty}")
-        else:
-            print("✅ All tests passed!")
-        
+            d_passed = sum(1 for r in diff_results if r.pass_status)
+            d_total = len(diff_results)
+            d_rate = d_passed / d_total
+            d_avg = np.mean([r.catch_rate for r in diff_results])
+            status = "✅" if d_rate >= 0.8 else "⚠️" if d_rate >= 0.5 else "❌"
+            print(f"{status} {difficulty.upper():12s} | Pass: {d_rate:5.1%} ({d_passed}/{d_total}) | Avg Catch: {d_avg:.3f}")
+
         print(f"\n{'='*80}\n")
-    
+
     def _save_results(self):
-        """Save results to JSON"""
-        
         results_dict = {
             'summary': {
                 'total_tests': len(self.results),
                 'passed_tests': sum(1 for r in self.results if r.pass_status),
-                'pass_rate': sum(1 for r in self.results if r.pass_status) / len(self.results),
-                'avg_catch_rate': float(np.mean([r.catch_rate for r in self.results])),
+                'pass_rate': sum(1 for r in self.results if r.pass_status) / len(self.results) if self.results else 0,
+                'avg_catch_rate': float(np.mean([r.catch_rate for r in self.results])) if self.results else 0,
                 'timestamp': time.time()
             },
             'test_results': [
@@ -635,158 +518,49 @@ class FairAdversarialTestSuite:
                 for r in self.results
             ]
         }
-        
+
         output_file = self.output_dir / 'fair_adversarial_results.json'
         with open(output_file, 'w') as f:
-            json.dump(results_dict, f, indent=2, default=lambda o: bool(o) if hasattr(o, "item") else o)
-        
+            json.dump(results_dict, f, indent=2)
+
         print(f"✅ Results saved to: {output_file}")
-    
+
     def _generate_visualizations(self):
-        """Generate visualization plots"""
-        
         if not self.results:
             return
-        
+
         sns.set_style("whitegrid")
-        
-        # 1. Performance by Difficulty
+
+        # Performance by Difficulty
         fig, ax = plt.subplots(figsize=(10, 6))
-        
         difficulties = ['production', 'stress', 'extreme']
-        diff_data = {d: [] for d in difficulties}
-        
-        for r in self.results:
-            if r.difficulty in diff_data:
-                diff_data[r.difficulty].append(r.catch_rate)
-        
-        positions = []
-        data_to_plot = []
-        labels = []
-        
-        for i, diff in enumerate(difficulties):
-            if diff_data[diff]:
-                positions.append(i)
-                data_to_plot.append(diff_data[diff])
-                labels.append(diff.capitalize())
-        
-        bp = ax.boxplot(data_to_plot, positions=positions, labels=labels,
-                        patch_artist=True, widths=0.6)
-        
-        # Color boxes
-        colors = ['lightgreen', 'orange', 'lightcoral']
-        for patch, color in zip(bp['boxes'], colors[:len(bp['boxes'])]):
-            patch.set_facecolor(color)
-        
-        ax.set_ylabel('Catch Rate', fontsize=12)
-        ax.set_xlabel('Difficulty Level', fontsize=12)
-        ax.set_title('Fair Adversarial Testing - Performance by Difficulty', 
-                    fontsize=14, fontweight='bold')
-        ax.set_ylim([0, 1.05])
-        ax.grid(True, alpha=0.3)
-        
-        plt.tight_layout()
-        plt.savefig(self.output_dir / 'performance_by_difficulty.png', dpi=300)
-        plt.close()
-        
-        # 2. Individual Test Results
-        fig, ax = plt.subplots(figsize=(12, 10))
-        
-        test_names = [r.test_name for r in self.results]
-        catch_rates = [r.catch_rate for r in self.results]
-        thresholds = [r.pass_threshold for r in self.results]
-        pass_statuses = [r.pass_status for r in self.results]
-        
-        y_pos = np.arange(len(test_names))
-        
-        # Plot bars
-        colors = ['green' if p else 'red' for p in pass_statuses]
-        bars = ax.barh(y_pos, catch_rates, color=colors, alpha=0.6)
-        
-        # Plot thresholds
-        ax.scatter(thresholds, y_pos, color='blue', marker='|', s=200, 
-                  linewidths=3, label='Pass Threshold', zorder=3)
-        
-        ax.set_yticks(y_pos)
-        ax.set_yticklabels(test_names, fontsize=9)
-        ax.set_xlabel('Catch Rate', fontsize=12)
-        ax.set_title('Fair Adversarial Testing - Individual Results', 
-                    fontsize=14, fontweight='bold')
-        ax.set_xlim([0, 1.05])
-        ax.legend()
-        ax.grid(True, alpha=0.3, axis='x')
-        
-        plt.tight_layout()
-        plt.savefig(self.output_dir / 'individual_test_results.png', dpi=300, bbox_inches='tight')
-        plt.close()
-        
-        # 3. Category Performance
-        fig, ax = plt.subplots(figsize=(12, 6))
-        
-        categories = defaultdict(list)
-        for r in self.results:
-            categories[r.category].append(r.catch_rate)
-        
-        cat_names = list(categories.keys())
-        cat_means = [np.mean(rates) for rates in categories.values()]
-        cat_stds = [np.std(rates) for rates in categories.values()]
-        
-        bars = ax.bar(cat_names, cat_means, yerr=cat_stds, capsize=5, alpha=0.7)
-        
-        # Color based on performance
-        for bar, mean in zip(bars, cat_means):
-            if mean >= 0.85:
-                bar.set_color('green')
-            elif mean >= 0.70:
-                bar.set_color('orange')
-            else:
-                bar.set_color('red')
-        
-        ax.set_ylabel('Average Catch Rate', fontsize=12)
-        ax.set_xlabel('Category', fontsize=12)
-        ax.set_title('Fair Adversarial Testing - Performance by Category', 
-                    fontsize=14, fontweight='bold')
-        ax.set_ylim([0, 1.05])
-        plt.xticks(rotation=45, ha='right')
-        ax.grid(True, alpha=0.3, axis='y')
-        
-        plt.tight_layout()
-        plt.savefig(self.output_dir / 'performance_by_category.png', dpi=300)
-        plt.close()
-        
+        diff_data = {d: [r.catch_rate for r in self.results if r.difficulty == d] for d in difficulties}
+        data_to_plot = [diff_data[d] for d in difficulties if diff_data[d]]
+        labels = [d.capitalize() for d in difficulties if diff_data[d]]
+
+        if data_to_plot:
+            bp = ax.boxplot(data_to_plot, labels=labels, patch_artist=True)
+            colors = ['lightgreen', 'orange', 'lightcoral']
+            for patch, color in zip(bp['boxes'], colors[:len(bp['boxes'])]):
+                patch.set_facecolor(color)
+            ax.set_ylabel('Catch Rate')
+            ax.set_title('Performance by Difficulty')
+            ax.set_ylim([0, 1.05])
+            plt.tight_layout()
+            plt.savefig(self.output_dir / 'performance_by_difficulty.png', dpi=300)
+            plt.close()
+
         print(f"✅ Visualizations saved to: {self.output_dir}")
 
 
 # ============================================================================
-# MAIN EXECUTION
+# MAIN (for direct testing)
 # ============================================================================
 
 def main():
-    """Example usage"""
-    
-    print("""
-    ╔════════════════════════════════════════════════════════════════╗
-    ║                                                                 ║
-    ║  FAIR ADVERSARIAL VALIDATION FRAMEWORK                         ║
-    ║                                                                 ║
-    ║  Tests realistic, production-relevant scenarios                ║
-    ║  Provides interpretable, actionable results                    ║
-    ║                                                                 ║
-    ╚════════════════════════════════════════════════════════════════╝
-    """)
-    
-    print("\nThis framework tests:")
-    print("  ✅ Realistic sensor noise (not random corruption)")
-    print("  ✅ Normal fab variation (not 3x jumps)")
-    print("  ✅ Efficiency improvements (not crisis scenarios)")
-    print("  ✅ Production-relevant perturbations")
-    print()
-    print("Expected performance ranges:")
-    print("  • Production tests: 90-98% catch rate")
-    print("  • Stress tests: 75-90% catch rate")
-    print("  • Extreme tests: 60-75% catch rate")
-    print()
+    print("FAIR ADVERSARIAL VALIDATION FRAMEWORK")
+    print("Use via: python run_fair_adversarial_validation.py --start_test X")
 
 
 if __name__ == "__main__":
-    main()
+    main()