Add files via upload

easyjailbreak/selector/EXP3SelectPolicy.py

@@ -0,0 +1,78 @@
+"""
+EXP3SelectPolicy class
+==========================
+"""
+import numpy as np
+from easyjailbreak.selector import SelectPolicy
+from easyjailbreak.datasets import Instance, JailbreakDataset
+
+__all__ = ["EXP3SelectPolicy"]
+
+class EXP3SelectPolicy(SelectPolicy):
+    """
+    A selection policy based on the Exponential-weight algorithm for Exploration and Exploitation (EXP3).
+    This policy is designed for environments with adversarial contexts, balancing between exploring new instances
+    and exploiting known rewards in a JailbreakDataset.
+    """
+
+    def __init__(self,
+                 Dataset: JailbreakDataset,
+                 energy: float = 1.0,
+                 gamma: float = 0.05,
+                 alpha: float = 25):
+        """
+        Initializes the EXP3SelectPolicy with a given JailbreakDataset and parameters for the EXP3 algorithm.
+
+        :param ~JailbreakDataset Dataset: The dataset from which instances will be selected.
+        :param float energy: Initial energy level (not used in current implementation).
+        :param float gamma: Parameter for controlling the exploration-exploitation trade-off.
+        :param float alpha: Learning rate for the weight updates.
+        """
+        super().__init__(Dataset)
+
+        self.energy = energy
+        self.gamma = gamma
+        self.alpha = alpha
+        self.last_choice_index = None
+
+        self.initial()
+
+    def initial(self):
+        """
+        Initializes or resets the weights and probabilities for each instance in the dataset.
+        """
+        self.weights = [1. for _ in range(len(self.Datasets))]
+        self.probs = [0. for _ in range(len(self.Datasets))]
+
+    def select(self) -> JailbreakDataset:
+        """
+        Selects an instance from the dataset based on the EXP3 algorithm.
+
+        :return ~JailbreakDataset: The selected instance from the dataset.
+        """
+        if len(self.Datasets) > len(self.weights):
+            self.weights.extend([1. for _ in range(len(self.Datasets) - len(self.weights))])
+        if len(self.Datasets) > len(self.probs):
+            self.probs.extend([0. for _ in range(len(self.Datasets) - len(self.probs))])
+
+        np_weights = np.array(self.weights)
+        probs = (1 - self.gamma) * np_weights / np_weights.sum() + self.gamma / len(self.Datasets)
+
+        self.last_choice_index = np.random.choice(len(self.Datasets), p=probs)
+
+        self.Datasets[self.last_choice_index].visited_num += 1
+        self.probs[self.last_choice_index] = probs[self.last_choice_index]
+
+        return JailbreakDataset([self.Datasets[self.last_choice_index]])
+
+    def update(self, prompt_nodes: JailbreakDataset):
+        """
+        Updates the weights of the last chosen instance based on the success of the prompts.
+
+        :param ~JailbreakDataset prompt_nodes: The dataset containing prompts used for updating weights.
+        """
+        succ_num = sum([prompt_node.num_jailbreak for prompt_node in prompt_nodes])
+
+        r = 1 - succ_num / len(prompt_nodes)
+        x = -1 * r / self.probs[self.last_choice_index]
+        self.weights[self.last_choice_index] *= np.exp(self.alpha * x / len(self.Datasets))