-
Notifications
You must be signed in to change notification settings - Fork 0
/
DQN_attack_XORPUF.py
178 lines (138 loc) · 6.34 KB
/
DQN_attack_XORPUF.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 25 12:57:16 2023
@author: Mieszko Ferens
Script to run an experiment for modelling an Arbiter or XOR PUF using Double
DQN. Requires ChainerRL v0.8.0.
"""
import argparse
import pandas as pd
from pathlib import Path
import numpy as np
from chainer import optimizers
from chainerrl import agents, q_functions, explorers, replay_buffer
from PUF_env import PUF_env
def main():
# Setup logging
import logging
logging.basicConfig(level=logging.INFO)
# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("--outdir", type=str, default="./Results/")
parser.add_argument("--seed", type=int, default=0, help="Random seed.")
parser.add_argument("--n-bits", type=int, default=64,
help="Challenge length in bits.")
parser.add_argument("--k", type=int, default=1,
help="Number of parallel APUFs in the XOR PUF.")
parser.add_argument("--train-data", type=int, default=10000,
help="Number of training data samples for the model.")
parser.add_argument("--test-data", type=int, default=1000,
help="Number of testing data samples for the model.")
parser.add_argument("--hidden-layer-size", type=int, default=16,
help="Size of all hidden layers of the nueral network.")
parser.add_argument("--n-hidden-layers", type=int, default=2,
help="Number of hidden layers of the neural network.")
parser.add_argument("--epsilon", type=float, default=0.2,
help="Probability of exploratory (random) action.")
parser.add_argument("--replay-buffer-size", type=int, default=25000,
help="Experience replay buffer size.")
parser.add_argument("--replay-start-size", type=int, default=5000,
help="Min replay buffer size to use experience replay.")
parser.add_argument("--batch-size", type=int, default=32,
help="Experience replay batch size.")
parser.add_argument("--gamma", type=float, default=0.1,
help="Discount factor.")
parser.add_argument("--update-interval", type=int, default=1,
help="Number of steps for updating main model.")
parser.add_argument("--target-update-interval", type=int, default=500,
help="Number of steps for updating target model.")
parser.add_argument("--max-episode-len", type=int, default=2,
help="Max number of steps per training episode.")
args = parser.parse_args()
# --- Create PUF environment
env = PUF_env(args.n_bits, args.k, seed=args.seed)
env.reset(args.train_data, args.test_data)
# --- Create agent
# Q function
q_func = q_functions.FCStateQFunctionWithDiscreteAction(
env.obs_size, env.n_actions, args.hidden_layer_size,
args.n_hidden_layers)
# Explorer
random_selector = lambda : np.random.choice(env.n_actions)
explorer = explorers.ConstantEpsilonGreedy(
args.epsilon, random_action_func=random_selector)
# Optimizer
optimizer = optimizers.Adam(eps=1e-3)
optimizer.setup(q_func)
# Replay buffer
r_buffer = replay_buffer.ReplayBuffer(capacity=args.replay_buffer_size)
# Input converter
phi = lambda x : x.astype(np.float32, copy=False)
# Agent
agent = agents.DoubleDQN(
q_func, optimizer, r_buffer, args.gamma, explorer,
replay_start_size=args.replay_start_size,
minibatch_size=args.batch_size, update_interval=args.update_interval,
target_update_interval=args.target_update_interval, phi=phi)
# --- Train the agent
print("Training...")
steps = 0 # Total number of training steps
for i in range(args.train_data):
# Set CRP and get challenge
obs = env.set_CRP(i)
reward = 0
done = False
t = 0 # Time step
while(t < args.max_episode_len and not done):
# Increment total and episodic time steps
steps += 1
t += 1
# Act on the challenge and train
action = agent.act_and_train(obs, reward)
# Check result
_ , reward, done = env.step(action)
# Train on episode end
agent.stop_episode_and_train(obs, reward, done)
if(not (i+1) % 1000):
print("Episode: " + str(i+1))
print("Steps: " + str(steps))
print("Training complete")
# --- Test the agent ---
print("Testing...")
total_reward = 0
for i in range(args.train_data, args.train_data + args.test_data):
# Set CRP and get challenge
obs = env.set_CRP(i)
# Predict response
action = agent.act(obs)
# Check prediction
_ , reward, _ = env.step(action)
total_reward += reward
print("Testing complete")
# Calculate accuracy
accuracy = (total_reward/args.test_data)
print("Final accuracy: " + str(accuracy))
# Log data into csv format
data = pd.DataFrame({"seed": [args.seed],
"n_bits": [args.n_bits],
"k": [args.k],
"train_data": [args.train_data],
"test_data": [args.test_data],
"gamma": [args.gamma],
"epsilon": [args.epsilon],
"n_hidden_layers": [args.n_hidden_layers],
"hidden_layer_size": [args.hidden_layer_size],
"replay_buffer_size": [args.replay_buffer_size],
"batch_size": [args.batch_size],
"update_interval": [args.update_interval],
"target_update_interval": [args.target_update_interval],
"max_episode_len": [args.max_episode_len],
"accuracy": [accuracy]})
filepath = Path(args.outdir + "out_DQN_" + str(args.k) + "XOR.csv")
if(filepath.is_file()):
data.to_csv(filepath, header=False, index=False, mode='a')
else:
filepath.parent.mkdir(parents=True, exist_ok=True)
data.to_csv(filepath, header=True, index=False, mode='a')
if(__name__ == "__main__"):
main()