chrissewell/pwnagotchi
1
0
Fork 0
mirror of https://github.com/jayofelony/pwnagotchi.git synced 2025-07-01 18:37:27 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'Snuf/master' into pwnagotchi-torch

Browse Source 
# Conflicts:
#	bin/pwnagotchi
#	builder/data/usr/bin/pwnlib
#	pwnagotchi/ai/__init__.py
This commit is contained in:
Jeroen Oudshoorn
2023-07-24 21:20:37 +02:00
parent 27da741005 8e10143339
commit 8329e282d3
19 changed files with 670 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
pwnagotchi/ai/__init__.py
View File

@ -14,10 +14,11 @@ def load(config, agent, epoch, from_disk=True):
try:
begin = time.time()
logging.info("[ai] bootstrapping dependencies ...")
start = time.time()
SB_BACKEND = "stable_baselines3"
SB_BACKEND = "stable_baselines3";
try:
from stable_baselines3 import A2C
@ -51,7 +52,6 @@ def load(config, agent, epoch, from_disk=True):
start = time.time()
import pwnagotchi.ai.gym as wrappers
logging.debug("[ai] gym wrapper imported in %.2fs" % (time.time() - start))
env = wrappers.Environment(agent, epoch)
@ -79,5 +79,5 @@ def load(config, agent, epoch, from_disk=True):
except Exception as e:
logging.exception("error while starting AI (%s)", e)
logging.warning("[ai] AI not loaded!")
return False
logging.warning("[ai] AI not loaded!")
return False

2
pwnagotchi/ai/epoch.py
View File

@ -177,7 +177,7 @@ class Epoch(object):
self.bored_for = 0
now = time.time()
cpu = pwnagotchi.cpu_load()
cpu = pwnagotchi.cpu_load("epoch")
mem = pwnagotchi.mem_usage()
temp = pwnagotchi.temperature()

9
pwnagotchi/ai/train.py
View File

@ -77,9 +77,12 @@ class Stats(object):
})
temp = "%s.tmp" % self.path
back = "%s.bak" % self.path
with open(temp, 'wt') as fp:
fp.write(data)
if os.path.isfile(self.path):
os.replace(self.path, back)
os.replace(temp, self.path)
@ -174,7 +177,13 @@ class AsyncTrainer(object):
logging.info("[ai] learning for %d epochs ..." % epochs_per_episode)
try:
self.set_training(True, epochs_per_episode)
# back up brain file before starting new training set
if os.path.isfile(self._nn_path):
back = "%s.bak" % self._nn_path
os.replace(self._nn_path, back)
self._view.set("mode", " ai")
self._model.learn(total_timesteps=epochs_per_episode, callback=self.on_ai_training_step)
self._view.set("mode", " AI")
except Exception as e:
logging.exception("[ai] error while training (%s)", e)
finally: