/Learning/environment/envirnomentrobustness.py
[INFO] Number of classes provided: 2
[INFO] Initial Model Accuracy: 0.4700
[INFO] Starting environment...
[INFO] Step 1: Applying Attack 'deepfool' and Defense 'gaussian_augmentation'
DeepFool: 100%|█████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 12.84it/s]
[ERROR] Error during step execution: operands could not be broadcast together with shapes (64,) (32,) 
import gymnasium as gym
import numpy as np
import torch
from torch import nn
from typing import Tuple, Dict, Any, Optional, Callable
from art.attacks.evasion import FastGradientMethod, ProjectedGradientDescent, DeepFool
from art.defences.preprocessor import FeatureSqueezing, SpatialSmoothing
from art.estimators.classification import PyTorchClassifier
import logging
from collections import deque
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import tkinter as tk
from tkinter import ttk, scrolledtext
from art.defences.preprocessor import GaussianAugmentation

class RobustModelEnvironment(gym.Env):
    """
    Environment for training a robust model against various adversarial attacks.
    """

    metadata = {'render.modes': ['human']}

    def __init__(
        self,
        model: nn.Module,
        dataset: torch.utils.data.Dataset,
        batch_size: int = 32,
        device: str = 'cuda' if torch.cuda.is_available() else 'cpu',
        memory_efficient: bool = True,
        max_steps: int = 1000,
        num_classes: Optional[int] = None  # Optional parameter to specify number of classes
    ) -> None:
        """
        Initialize the RobustModelEnvironment.

        Args:
            model (nn.Module): The PyTorch model to evaluate and train.
            dataset (torch.utils.data.Dataset): The dataset for training and evaluation.
            batch_size (int, optional): Batch size for data loading. Defaults to 32.
            device (str, optional): Device to run computations on. Defaults to 'cuda' if available.
            memory_efficient (bool, optional): Flag for memory-efficient operations. Defaults to True.
            max_steps (int, optional): Maximum number of steps per episode. Defaults to 1000.
            num_classes (Optional[int], optional): Number of distinct classes. If not provided, deduced from dataset.
        """
        super(RobustModelEnvironment, self).__init__()
        self.model = model.to(device)
        self.dataset = dataset
        self.batch_size = batch_size
        self.device = device
        self.memory_efficient = memory_efficient
        self.max_steps = max_steps
        self.steps = 0

        # Initialize logger first
        self.logger = self._setup_logger()

        self.initial_checkpoint = self._get_model_checkpoint()

        if num_classes is not None:
            self.num_classes = num_classes
            self.logger.info(f"Number of classes provided: {self.num_classes}")
        else:
            self.num_classes = self._get_num_classes()

        self.classifier = self._setup_classifier()
        self.original_accuracy = self._evaluate_accuracy()

        self.attacks = self._setup_attacks()
        self.defenses = self._setup_defenses()

        self.action_space = gym.spaces.Discrete(len(self.attacks) * len(self.defenses))
        self.observation_space = self._setup_observation_space()

        self.performance_history = deque(maxlen=100)
        self.performance_history.append(self.original_accuracy)

    def _setup_logger(self) -> logging.Logger:
        """
        Set up the logger for the environment.

        Returns:
            logging.Logger: Configured logger.
        """
        logger = logging.getLogger('RobustModelEnvironment')
        logger.setLevel(logging.INFO)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('[%(levelname)s] %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        return logger

    def _get_num_classes(self) -> int:
        """
        Get the number of classes from the dataset.

        Returns:
            int: Number of distinct classes.
        """
        try:
            # Attempt to extract labels
            if isinstance(self.dataset, TensorDataset):
                # TensorDataset stores data as tuples of tensors (inputs, labels)
                _, labels = self.dataset.tensors
                labels = labels.tolist()
            else:
                # General case: iterate through dataset to extract labels
                labels = [label for _, label in self.dataset]

            # Ensure labels are integers
            labels = [label if isinstance(label, int) else label.item() for label in labels]

            num_classes = len(set(labels))
            self.logger.info(f"Detected number of classes: {num_classes}")
            return num_classes
        except Exception as e:
            self.logger.error(f"Failed to get number of classes: {e}")
            raise

    def _setup_classifier(self) -> PyTorchClassifier:
        """
        Set up the PyTorch classifier using ART.

        Returns:
            PyTorchClassifier: Configured classifier.
        """
        try:
            return PyTorchClassifier(
                model=self.model,
                loss=nn.CrossEntropyLoss(),
                optimizer=optim.Adam(self.model.parameters()),
                input_shape=self._get_input_shape(),
                nb_classes=self.num_classes,
                device_type='cuda' if self.device == 'cuda' else 'cpu'
            )
        except Exception as e:
            self.logger.error(f"Failed to set up classifier: {e}")
            raise

    def _get_input_shape(self) -> Tuple[int, ...]:
        """
        Get the input shape from the dataset.

        Returns:
            Tuple[int, ...]: Shape of a single input sample.
        """
        try:
            sample_input, _ = self.dataset[0]
            return sample_input.shape
        except Exception as e:
            self.logger.error(f"Failed to get input shape: {e}")
            raise

    def _setup_attacks(self) -> Dict[str, Callable[[], Any]]:
        """
        Set up the attack methods.

        Returns:
            Dict[str, Callable[[], Any]]: Dictionary of attack names to their constructors.
        """
        try:
            return {
                'fgsm': lambda: FastGradientMethod(estimator=self.classifier, eps=0.1),
                'pgd': lambda: ProjectedGradientDescent(estimator=self.classifier, eps=0.1, max_iter=10),
                'deepfool': lambda: DeepFool(classifier=self.classifier, max_iter=50,batch_size=32)
            }
        except Exception as e:
            self.logger.error(f"Failed to set up attacks: {e}")
            raise

    def _setup_defenses(self) -> Dict[str, Callable[[], Any]]:
        """
        Set up the defense methods.
    
        Returns:
            Dict[str, Callable[[], Any]]: Dictionary of defense names to their constructors.
        """
        try:
            return {
                'feature_squeezing': lambda: FeatureSqueezing(bit_depth=5, clip_values=(0.0, 1.0)),
                'gaussian_augmentation': lambda: GaussianAugmentation( clip_values=(0.0, 1.0))
            }
        except Exception as e:
            self.logger.error(f"Failed to set up defenses: {e}")
            raise

    def _setup_observation_space(self) -> gym.spaces.Dict:
        """
        Set up the observation space.

        Returns:
            gym.spaces.Dict: Observation space definition.
        """
        return gym.spaces.Dict({
            'model_state': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32),
            'attack_impact': gym.spaces.Box(low=0.0, high=1.0, shape=(1,), dtype=np.float32),
            'current_accuracy': gym.spaces.Box(low=0.0, high=1.0, shape=(1,), dtype=np.float32),
            'step': gym.spaces.Box(low=0, high=self.max_steps, shape=(1,), dtype=np.int32)
        })

    def step(self, action: int) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
        """
        Execute one step in the environment.
    
        Args:
            action (int): Combined index for attack and defense.
    
        Returns:
            Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
                - observation
                - reward
                - done
                - truncated
                - info
        """
        self.steps += 1
        attack_idx = action // len(self.defenses)
        defense_idx = action % len(self.defenses)
    
        try:
            attack_name = list(self.attacks.keys())[attack_idx]
            defense_name = list(self.defenses.keys())[defense_idx]
            attack = self.attacks[attack_name]()
            defense = self.defenses[defense_name]()
    
            self.logger.info(f"Step {self.steps}: Applying Attack '{attack_name}' and Defense '{defense_name}'")
    
            data_loader = DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True)
            x_test, y_test = next(iter(data_loader))
            x_test, y_test = x_test.to(self.device), y_test.to(self.device)
    
            x_test_np = x_test.cpu().numpy()
            y_test_np = y_test.cpu().numpy()
    
            # Generate adversarial examples
            x_adv = attack.generate(x=x_test_np)
    
            # Manually compute attack accuracy
            y_pred_adv = self.classifier.predict(x_adv)
            self.logger.debug(f"y_pred_adv before argmax: {y_pred_adv.shape}, sample: {y_pred_adv[:2]}")
            if y_pred_adv.ndim > 1 and y_pred_adv.shape[1] > 1:
                y_pred_adv = np.argmax(y_pred_adv, axis=1)
                self.logger.debug(f"y_pred_adv after argmax: {y_pred_adv.shape}, sample: {y_pred_adv[:2]}")
            attack_accuracy = np.mean(y_pred_adv == y_test_np)
            attack_impact = self.original_accuracy - attack_accuracy
    
            # Apply defense
            x_defended = defense(x_adv)[0]
            x_defended_tensor = torch.from_numpy(x_defended).to(self.device)
    
            # Compute new accuracy manually
            y_pred_defended = self.classifier.predict(x_defended)
            self.logger.debug(f"y_pred_defended before argmax: {y_pred_defended.shape}, sample: {y_pred_defended[:2]}")
            if y_pred_defended.ndim > 1 and y_pred_defended.shape[1] > 1:
                y_pred_defended = np.argmax(y_pred_defended, axis=1)
                self.logger.debug(f"y_pred_defended after argmax: {y_pred_defended.shape}, sample: {y_pred_defended[:2]}")
            new_accuracy = np.mean(y_pred_defended == y_test_np)
    
            self.logger.info(f"Attack Accuracy: {attack_accuracy:.4f}, New Accuracy after Defense: {new_accuracy:.4f}")
    
            reward = self._calculate_reward(new_accuracy, attack_impact)
            done = self.steps >= self.max_steps
            observation = self._get_observation(new_accuracy, attack_impact)
    
            self.performance_history.append(new_accuracy)
    
            # Simulate model weight degradation
            self._degrade_model_weights(attack_impact)
    
            return observation, reward, done, False, {
                'attack_name': attack_name,
                'defense_name': defense_name,
                'attack_accuracy': attack_accuracy,
                'new_accuracy': new_accuracy
            }
    
        except StopIteration:
            self.logger.warning("End of DataLoader reached. Resetting environment.")
            return self.reset()
        except IndexError:
            self.logger.error("Invalid action index.")
            return self._get_observation(self.original_accuracy, 0.0), 0.0, True, False, {}
        except Exception as e:
            self.logger.error(f"Error during step execution: {e}")
            return self._get_observation(self.original_accuracy, 0.0), 0.0, True, False, {}

    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[Dict[str, Any]] = None
    ) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
        """
        Reset the environment to its initial state.

        Args:
            seed (Optional[int], optional): Random seed. Defaults to None.
            options (Optional[Dict[str, Any]], optional): Additional options. Defaults to None.

        Returns:
            Tuple[Dict[str, np.ndarray], Dict[str, Any]]: Initial observation and info.
        """
        super().reset(seed=seed)
        self.steps = 0
        self._load_model_checkpoint(self.initial_checkpoint)
        self.performance_history = deque(maxlen=100)
        self.performance_history.append(self.original_accuracy)
        observation = self._get_observation(self.original_accuracy, 0.0)
        self.logger.info("Environment has been reset.")
        return observation, {}

    def _calculate_reward(self, new_accuracy: float, attack_impact: float) -> float:
        """
        Calculate the reward based on accuracy improvement and attack impact.

        Args:
            new_accuracy (float): Current model accuracy post-defense.
            attack_impact (float): Impact of the attack on the model.

        Returns:
            float: Calculated reward.
        """
        robustness_improvement = new_accuracy - self.original_accuracy
        reward = robustness_improvement * 100 + (1 - attack_impact) * 50
        self.logger.debug(f"Calculated Reward: {reward}")
        return reward

    def _get_model_checkpoint(self) -> Dict[str, torch.Tensor]:
        """
        Get a checkpoint of the model's current state.

        Returns:
            Dict[str, torch.Tensor]: Dictionary of model parameters.
        """
        return {name: param.data.clone() for name, param in self.model.named_parameters()}

    def _update_model_checkpoint(self) -> None:
        """
        Update the initial checkpoint with the current model state.
        """
        self.initial_checkpoint = self._get_model_checkpoint()

    def _load_model_checkpoint(self, checkpoint: Dict[str, torch.Tensor]) -> None:
        """
        Load a model checkpoint.

        Args:
            checkpoint (Dict[str, torch.Tensor]): Model parameters to load.
        """
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                if name in checkpoint:
                    param.copy_(checkpoint[name])
        self.logger.info("Model checkpoint loaded.")

    def _evaluate_accuracy(self) -> float:
        """
        Evaluate the current accuracy of the model.

        Returns:
            float: Model accuracy.
        """
        self.model.eval()
        correct = 0
        total = 0
        try:
            data_loader = DataLoader(self.dataset, batch_size=self.batch_size, shuffle=False)
            with torch.no_grad():
                for inputs, labels in data_loader:
                    inputs, labels = inputs.to(self.device), labels.to(self.device)
                    outputs = self.model(inputs)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            accuracy = correct / total
            self.logger.info(f"Initial Model Accuracy: {accuracy:.4f}")
            return accuracy
        except Exception as e:
            self.logger.error(f"Failed to evaluate accuracy: {e}")
            return 0.0

    def _compute_accuracy_tensor(self, outputs: torch.Tensor, labels: torch.Tensor) -> float:
        """
        Compute accuracy given model outputs and labels.

        Args:
            outputs (torch.Tensor): Model outputs.
            labels (torch.Tensor): Ground truth labels.

        Returns:
            float: Computed accuracy.
        """
        try:
            _, predicted = torch.max(outputs.data, 1)
            correct = (predicted == labels).sum().item()
            accuracy = correct / labels.size(0)
            return accuracy
        except Exception as e:
            self.logger.error(f"Failed to compute accuracy: {e}")
            return 0.0

    def _get_observation(self, accuracy: float, attack_impact: float) -> Dict[str, np.ndarray]:
        """
        Get the current observation of the environment.

        Args:
            accuracy (float): Current model accuracy.
            attack_impact (float): Impact of the latest attack.

        Returns:
            Dict[str, np.ndarray]: Observation dictionary.
        """
        model_state_sum = sum(p.sum().item() for p in self.model.parameters())
        return {
            'model_state': np.array([model_state_sum], dtype=np.float32),
            'attack_impact': np.array([attack_impact], dtype=np.float32),
            'current_accuracy': np.array([accuracy], dtype=np.float32),
            'step': np.array([self.steps], dtype=np.int32)
        }

    def _degrade_model_weights(self, attack_impact: float) -> None:
        """
        Simulate model weight degradation based on attack impact.

        Args:
            attack_impact (float): Impact of the attack on the model.
        """
        try:
            degradation_factor = attack_impact * 0.01  # Scale factor for degradation
            with torch.no_grad():
                for param in self.model.parameters():
                    noise = torch.randn(param.size()).to(self.device) * degradation_factor
                    param.add_(noise)
            self.logger.debug(f"Model weights degraded by factor: {degradation_factor:.4f}")
        except Exception as e:
            self.logger.error(f"Failed to degrade model weights: {e}")

    def render(self, mode: str = 'human') -> None:
        """
        Render the current state of the environment.

        Args:
            mode (str, optional): Mode to render with. Defaults to 'human'.
        """
        if mode == 'human':
            current_accuracy = self.performance_history[-1]
            self.logger.info(f"Step: {self.steps}, Current Accuracy: {current_accuracy:.4f}")

    def close(self) -> None:
        """
        Clean up resources used by the environment.
        """
        self.logger.info("Closing environment.")
        pass


class EnvironmentGUI:
    """
    Advanced GUI for visualizing the RobustModelEnvironment.
    """

    def __init__(self, env: RobustModelEnvironment) -> None:
        """
        Initialize the GUI with the given environment.

        Args:
            env (RobustModelEnvironment): The environment to visualize.
        """
        self.env = env
        self.root = tk.Tk()
        self.root.title("Robust Model Environment Visualization")
        self._setup_gui_elements()
        self._update_gui()

    def _setup_gui_elements(self) -> None:
        """
        Set up GUI components.
        """
        # Frame for Metrics
        metrics_frame = ttk.LabelFrame(self.root, text="Metrics")
        metrics_frame.grid(column=0, row=0, padx=10, pady=10, sticky="nsew")

        # Metrics Labels
        self.step_var = tk.StringVar(value="Step: 0")
        self.accuracy_var = tk.StringVar(value="Accuracy: 0.0")
        self.attack_impact_var = tk.StringVar(value="Attack Impact: 0.0")

        ttk.Label(metrics_frame, textvariable=self.step_var).grid(column=0, row=0, sticky="w", padx=5, pady=5)
        ttk.Label(metrics_frame, textvariable=self.accuracy_var).grid(column=0, row=1, sticky="w", padx=5, pady=5)
        ttk.Label(metrics_frame, textvariable=self.attack_impact_var).grid(column=0, row=2, sticky="w", padx=5, pady=5)

        # Frame for Controls
        controls_frame = ttk.LabelFrame(self.root, text="Controls")
        controls_frame.grid(column=0, row=1, padx=10, pady=10, sticky="nsew")

        # Control Buttons
        self.start_button = ttk.Button(controls_frame, text="Start", command=self.start_environment)
        self.start_button.grid(column=0, row=0, padx=5, pady=5)

        self.pause_button = ttk.Button(controls_frame, text="Pause", command=self.pause_environment)
        self.pause_button.grid(column=1, row=0, padx=5, pady=5)

        self.reset_button = ttk.Button(controls_frame, text="Reset", command=self.reset_environment)
        self.reset_button.grid(column=2, row=0, padx=5, pady=5)

        # Frame for Logs
        logs_frame = ttk.LabelFrame(self.root, text="Logs")
        logs_frame.grid(column=0, row=2, padx=10, pady=10, sticky="nsew")

        # Scrolled Text for Logs
        self.log_text = scrolledtext.ScrolledText(logs_frame, width=80, height=20, state='disabled')
        self.log_text.grid(column=0, row=0, padx=5, pady=5)

        # Configure Grid Weights
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(2, weight=1)
        metrics_frame.columnconfigure(0, weight=1)
        controls_frame.columnconfigure((0, 1, 2), weight=1)
        logs_frame.columnconfigure(0, weight=1)

        # Redirect logger to GUI
        self._redirect_logging()

    def _redirect_logging(self) -> None:
        """
        Redirect logging output to the GUI log text widget.
        """
        class TextHandler(logging.Handler):
            def __init__(self, text_widget: scrolledtext.ScrolledText) -> None:
                super().__init__()
                self.text_widget = text_widget

            def emit(self, record: logging.LogRecord) -> None:
                msg = self.format(record) + '\n'
                self.text_widget.configure(state='normal')
                self.text_widget.insert(tk.END, msg)
                self.text_widget.configure(state='disabled')
                self.text_widget.yview(tk.END)

        handler = TextHandler(self.log_text)
        handler.setFormatter(logging.Formatter('[%(levelname)s] %(message)s'))
        self.env.logger.addHandler(handler)

    def start_environment(self) -> None:
        """
        Start the environment's training loop.
        """
        if not getattr(self, 'running', False):
            self.env.logger.info("Starting environment...")
            self.running = True
            self._run_steps()

    def pause_environment(self) -> None:
        """
        Pause the environment's training loop.
        """
        if getattr(self, 'running', False):
            self.env.logger.info("Pausing environment...")
            self.running = False

    def reset_environment(self) -> None:
        """
        Reset the environment.
        """
        self.env.reset()
        self.env.logger.info("Environment has been reset.")
        self._update_gui()

    def _run_steps(self) -> None:
        """
        Run environment steps in a loop.
        """
        if self.running and self.env.steps < self.env.max_steps:
            action = self._select_action()
            observation, reward, done, _, info = self.env.step(action)
            self._update_gui(observation)
            if done:
                self.running = False
                self.env.logger.info("Episode finished.")
                return
            # Schedule next step
            self.root.after(100, self._run_steps)

    def _select_action(self) -> int:
        """
        Select an action for the agent. Placeholder for agent logic.

        Returns:
            int: Selected action index.
        """
        # Placeholder: Select random action
        return np.random.randint(0, self.env.action_space.n)

    def _update_gui(self, observation: Optional[Dict[str, np.ndarray]] = None) -> None:
        """
        Update GUI elements with the latest observation.

        Args:
            observation (Optional[Dict[str, np.ndarray]], optional): Latest observation. Defaults to None.
        """
        if observation:
            self.step_var.set(f"Step: {observation['step'][0]}")
            self.accuracy_var.set(f"Accuracy: {observation['current_accuracy'][0]:.4f}")
            self.attack_impact_var.set(f"Attack Impact: {observation['attack_impact'][0]:.4f}")
        else:
            self.step_var.set(f"Step: {self.env.steps}")
            if self.env.performance_history:
                self.accuracy_var.set(f"Accuracy: {self.env.performance_history[-1]:.4f}")
                attack_impact = self.env.original_accuracy - self.env.performance_history[-1]
                self.attack_impact_var.set(f"Attack Impact: {attack_impact:.4f}")

    def run(self) -> None:
        """
        Start the GUI event loop.
        """
        self.running = False
        self.root.mainloop()


class SimpleModel(nn.Module):
    """
    A simple neural network model for demonstration purposes.
    """

    def __init__(self, input_size: int, num_classes: int) -> None:
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


def main() -> None:
    """
    Main function to initialize the environment and launch the GUI.
    """
    # Define input size and number of classes
    input_size = 10
    num_classes = 2  # Binary classification

    # Initialize the model
    model = SimpleModel(input_size=input_size, num_classes=num_classes)

    # Create a dummy dataset: Replace with your actual dataset
    dummy_inputs = torch.randn(100, input_size)  # Example input features
    dummy_labels = torch.randint(0, num_classes, (100,))  # Example labels for binary classification
    dataset = TensorDataset(dummy_inputs, dummy_labels)

    # Initialize environment
    env = RobustModelEnvironment(
        model=model,
        dataset=dataset,
        batch_size=32,
        device='cpu',  # Use 'cuda' if GPU is available
        max_steps=100,
        num_classes=num_classes  # Optional: Specify number of classes
    )

    # Initialize and run GUI
    gui = EnvironmentGUI(env)
    gui.run()


if __name__ == "__main__":
    main()