class CustomEnv(ParallelEnv, EzPickle):
def __init__(self, question, template, response, judgement_model, embedding_model, mutate, prompt_compose, logger, question_seed, template_seed, question_pool, template_pool):
super(CustomEnv, self).__init__()
action_space = spaces.Discrete(5) # 5 mutators for each agent
self.embedding_model = embedding_model
self.question = question
self.template = template
self.response = response
self.mutate = mutate
self.prompt_compose = prompt_compose
self.judgement_model = judgement_model
self.question_seed = question_seed
self.template_seed = template_seed
self.question_pool = question_pool
self.template_pool = template_pool
self.logger = logger
self.state = self._get_embedding(question, template)
self.state1 = self.embedding_model.encode([response])[0]
self.combined_state = np.concatenate((self.state, self.state1))
observation_space = spaces.Box(
low=-np.inf, high=np.inf, shape=(self.combined_state.shape[0],), dtype=np.float32
)
self.iq = 0
self.previous_iq = 0
self.previous_score = 0
self.score = 0
self.metadata = {
"render_modes": ["human", "rgb_array"],
"name": "Fuzzing",
"render_fps": 60,
}
agent_names = ["question_mutator", "template_mutator"]
self.agents = [agent_names[n] for n in range(2)]
self.possible_agents = self.agents[:]
self.action_spaces = {agent: action_space for agent in self.agents}
self.observation_spaces = {agent: observation_space for agent in self.agents}
def _get_embedding(self, question, template):
prompt = self.prompt_compose(question, template)
prompt_embedding = self.embedding_model.encode([prompt])[0]
return prompt_embedding
def step(self, action_dict):
retries = 5
backoff_factor = 1
for attempt in range(retries):
try:
question_policy = question_policies[action_dict['question_mutator']]
template_policy = template_policies[action_dict['template_mutator']]
self.logger.debug(f'question: {self.question}\n')
self.logger.debug(f'template: {self.template}\n')
self.logger.debug("mutating....\n")
self.question_mutated, self.template_mutated, self.response, iq = self.mutate(self.question, self.template, question_policy, template_policy, self.logger)
self.logger.debug('embedding....')
self.iq = iq
self.logger.debug('judging....')
judge, judge_score = self.judgement_model.evaluate(self.response, self.question_seed, self.template_seed, self.question_pool, self.template_pool)
self.score = judge_score
reward = self._compute_reward()
self.previous_iq = iq
self.previous_score = judge_score
if judge == SeedOperation.insert:
self.template_seed.success_attack_num += 1
self.template_pool.add_seed(self.template_mutated, score=judge_score)
self.template_pool.save_current()
# save_template_pool(self.template_pool)
info = {'response': self.response, 'judge_score': judge_score}
self.logger.debug(f'score: {judge_score}, iq:{iq}')
done = False
truncated = False
self.question_seed = self.question_pool.select()
self.question = self.question_seed.visit()
self.template_seed = self.template_pool.select()
self.template = self.template_seed.visit()
self.state = self._get_embedding(self.question, self.template)
self.state1 = self.embedding_model.encode([self.response])[0]
self.combined_state = np.concatenate((self.state, self.state1))
return {agent: self.combined_state for agent in self.agents}, {agent: reward for agent in self.agents}, {agent: done for agent in self.agents}, {agent: truncated for agent in self.agents}, {agent: info for agent in self.agents}
except Exception as e:
self.logger.error(f"Attempt {attempt + 1} failed with error: {e}")
if attempt < retries - 1:
sleep_time = backoff_factor * (2 ** attempt)
self.logger.debug(f"Retrying in {sleep_time} seconds...")
time.sleep(sleep_time)
else:
self.logger.error("All retry attempts failed.")
return self.state, 0, False, False, {}
def reset(self, seed=None, options=None):
# super().reset(seed=seed)
self.previous_iq = 0
self.agents = self.possible_agents[:]
self.terminations = {agent: False for agent in self.agents}
obs = self._observe()
infos = {agent: {} for agent in self.agents}
return {agent: obs for agent in self.agents}, infos
def observation_space(self, agent):
return self.observation_spaces[agent]
def action_space(self, agent):
return self.action_spaces[agent]
def _observe(self):
return self.combined_state
def render(self, mode='human'):
pass
def close(self):
pass
def _compute_reward(self):
iq_diff = self.iq - self.previous_iq
score_diff = self.score - self.previous_score
score_reward = np.log(1 + abs(score_diff))
iq_reward = np.log(1 + abs(iq_diff))
if iq_diff < 0:
iq_reward = -iq_reward
if score_diff < 0:
score_reward = -score_reward
reward = (iq_reward + score_reward) / 2
return reward
def seed(self, seed=None):
self.np_random, seed = seeding.np_random(seed)
return [seed]
```