Update build

Signed-off-by: jayofelony <oudshoorn.jeroen@gmail.com>

pwnagotchi/agent.py

@@ -15,12 +15,11 @@ from pwnagotchi.automata import Automata
 from pwnagotchi.log import LastSession
 from pwnagotchi.bettercap import Client
 from pwnagotchi.mesh.utils import AsyncAdvertiser
-from pwnagotchi.ai.train import AsyncTrainer
 
 RECOVERY_DATA_FILE = '/root/.pwnagotchi-recovery'
 
 
-class Agent(Client, Automata, AsyncAdvertiser, AsyncTrainer):
+class Agent(Client, Automata, AsyncAdvertiser):
     def __init__(self, view, config, keypair):
         Client.__init__(self,
                         "127.0.0.1" if "hostname" not in config['bettercap'] else config['bettercap']['hostname'],
@@ -30,7 +29,6 @@ class Agent(Client, Automata, AsyncAdvertiser, AsyncTrainer):
                         "pwnagotchi" if "password" not in config['bettercap'] else config['bettercap']['password'])
         Automata.__init__(self, config, view)
         AsyncAdvertiser.__init__(self, config, view, keypair)
-        AsyncTrainer.__init__(self, config)
 
         self._started_at = time.time()
         self._current_channel = 0
@@ -130,7 +128,6 @@ class Agent(Client, Automata, AsyncAdvertiser, AsyncTrainer):
                 time.sleep(1)
 
     def start(self):
-        self.start_ai()
         self._wait_bettercap()
         self.setup_events()
         self.set_starting()

pwnagotchi/ai/__init__.py

@@ -1,74 +0,0 @@
-import os
-import time
-import logging
-
-# https://stackoverflow.com/questions/40426502/is-there-a-way-to-suppress-the-messages-tensorflow-prints/40426709
-# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # or any {'0', '1', '2'}
-
-
-def load(config, agent, epoch, from_disk=True):
-    config = config['ai']
-    if not config['enabled']:
-        logging.info("ai disabled")
-        return False
-
-    try:
-        begin = time.time()
-
-        logging.info("[AI] bootstrapping dependencies ...")
-
-        start = time.time()
-        SB_BACKEND = "stable_baselines3"
-
-        from stable_baselines3 import A2C
-        logging.debug("[AI] A2C imported in %.2fs" % (time.time() - start))
-
-        # remove invalid ai.parameters leftover from tensor_flow, if present
-        for key in [ 'alpha', 'epsilon', 'lr_schedule' ]:
-            if key in config['params']:
-                logging.info("Removing legacy ai parameter %s" % key);
-                del config['params'][key]
-        
-        start = time.time()
-        from stable_baselines3.a2c import MlpPolicy
-        logging.debug("[AI] MlpPolicy imported in %.2fs" % (time.time() - start))
-        SB_A2C_POLICY = MlpPolicy
-
-        start = time.time()
-        from stable_baselines3.common.vec_env import DummyVecEnv
-        logging.debug("[AI] DummyVecEnv imported in %.2fs" % (time.time() - start))
-
-        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)
-        env = DummyVecEnv([lambda: env])
-
-        logging.info("[AI] creating model ...")
-
-        start = time.time()
-        a2c = A2C(SB_A2C_POLICY, env, **config['params'])
-        logging.debug("[AI] A2C created in %.2fs" % (time.time() - start))
-
-        if from_disk and os.path.exists(config['path']):
-            logging.info("[AI] loading %s ..." % config['path'])
-            start = time.time()
-            a2c.load(config['path'], env)
-            logging.debug("[AI] A2C loaded in %.2fs" % (time.time() - start))
-        else:
-            logging.info("[AI] model created:")
-            for key, value in config['params'].items():
-                logging.info("      %s: %s" % (key, value))
-
-        logging.debug("[AI] total loading time is %.2fs" % (time.time() - begin))
-
-        return a2c
-    except Exception as e:
-        logging.info("[AI] Error while starting AI")
-        logging.debug("[AI] error while starting AI (%s)", e)
-        logging.info("[AI] Deleting brain and restarting.")
-        os.system("rm /root/brain.nn && service pwnagotchi restart")
-
-    logging.warning("[AI] AI not loaded!")
-    return False

pwnagotchi/ai/featurizer.py

@@ -1,61 +0,0 @@
-import numpy as np
-
-import pwnagotchi.mesh.wifi as wifi
-
-MAX_EPOCH_DURATION = 1024
-
-
-histogram_size = wifi.NumChannels
-
-shape = (1,
-         # aps per channel
-         histogram_size +
-         # clients per channel
-         histogram_size +
-         # peers per channel
-         histogram_size +
-         # duration
-         1 +
-         # inactive
-         1 +
-         # active
-         1 +
-         # missed
-         1 +
-         # hops
-         1 +
-         # deauths
-         1 +
-         # assocs
-         1 +
-         # handshakes
-         1)
-
-
-def featurize(state, step):
-    tot_epochs = step + 1e-10
-    tot_interactions = (state['num_deauths'] + state['num_associations']) + 1e-10
-    return np.concatenate((
-        # aps per channel
-        state['aps_histogram'],
-        # clients per channel
-        state['sta_histogram'],
-        # peers per channel
-        state['peers_histogram'],
-        # duration
-        [np.clip(state['duration_secs'] / MAX_EPOCH_DURATION, 0.0, 1.0)],
-        # inactive
-        [state['inactive_for_epochs'] / tot_epochs],
-        # active
-        [state['active_for_epochs'] / tot_epochs],
-        # missed
-        [state['missed_interactions'] / tot_interactions],
-        # hops
-        [state['num_hops'] / wifi.NumChannels],
-        # deauths
-        [state['num_deauths'] / tot_interactions],
-        # assocs
-        [state['num_associations'] / tot_interactions],
-        # handshakes
-        [state['num_handshakes'] / tot_interactions],
-    ))

pwnagotchi/ai/gym.py

@@ -1,148 +0,0 @@
-import logging
-import gymnasium as gym
-from gymnasium import spaces
-import numpy as np
-
-import pwnagotchi.ai.featurizer as featurizer
-import pwnagotchi.ai.reward as reward
-from pwnagotchi.ai.parameter import Parameter
-
-
-class Environment(gym.Env):
-    render_mode = "human"
-    metadata = {'render_modes': ['human']}
-    params = [
-        Parameter('min_rssi', min_value=-200, max_value=-50),
-        Parameter('ap_ttl', min_value=30, max_value=600),
-        Parameter('sta_ttl', min_value=60, max_value=300),
-
-        Parameter('recon_time', min_value=5, max_value=60),
-        Parameter('max_inactive_scale', min_value=3, max_value=10),
-        Parameter('recon_inactive_multiplier', min_value=1, max_value=3),
-        Parameter('hop_recon_time', min_value=5, max_value=60),
-        Parameter('min_recon_time', min_value=1, max_value=30),
-        Parameter('max_interactions', min_value=1, max_value=25),
-        Parameter('max_misses_for_recon', min_value=3, max_value=10),
-        Parameter('excited_num_epochs', min_value=5, max_value=30),
-        Parameter('bored_num_epochs', min_value=5, max_value=30),
-        Parameter('sad_num_epochs', min_value=5, max_value=30),
-    ]
-
-    def __init__(self, agent, epoch):
-        super(Environment, self).__init__()
-        self._agent = agent
-        self._epoch = epoch
-        self._epoch_num = 0
-        self._last_render = None
-
-        channels = agent.supported_channels()
-
-        Environment.params += [
-            Parameter('_channel_%d' % ch, min_value=0, max_value=1, meta=ch + 1) for ch in
-            range(featurizer.histogram_size) if ch + 1 in channels
-        ]
-
-        self.last = {
-            'reward': 0.0,
-            'observation': None,
-            'policy': None,
-            'params': {},
-            'state': None,
-            'state_v': None
-        }
-
-        self.action_space = spaces.MultiDiscrete([p.space_size() for p in Environment.params if p.trainable])
-        self.observation_space = spaces.Box(low=0, high=1, shape=featurizer.shape, dtype=np.float32)
-        self.reward_range = reward.range
-
-    @staticmethod
-    def policy_size():
-        return len(list(p for p in Environment.params if p.trainable))
-
-    @staticmethod
-    def policy_to_params(policy):
-        num = len(policy)
-        params = {}
-
-        assert len(Environment.params) == num
-
-        channels = []
-
-        for i in range(num):
-            param = Environment.params[i]
-
-            if '_channel' not in param.name:
-                params[param.name] = param.to_param_value(policy[i])
-            else:
-                has_chan = param.to_param_value(policy[i])
-                # print("%s policy:%s bool:%s" % (param.name, policy[i], has_chan))
-                chan = param.meta
-                if has_chan:
-                    channels.append(chan)
-
-        params['channels'] = channels
-
-        return params
-
-    def _next_epoch(self):
-        logging.debug("[ai] waiting for epoch to finish ...")
-        return self._epoch.wait_for_epoch_data()
-
-    def _apply_policy(self, policy):
-        new_params = Environment.policy_to_params(policy)
-        self.last['policy'] = policy
-        self.last['params'] = new_params
-        self._agent.on_ai_policy(new_params)
-
-    def step(self, policy):
-        # create the parameters from the policy and update
-        # them in the algorithm
-        self._apply_policy(policy)
-        self._epoch_num += 1
-
-        # wait for the algorithm to run with the new parameters
-        state = self._next_epoch()
-
-        self.last['reward'] = state['reward']
-        self.last['state'] = state
-        self.last['state_v'] = featurizer.featurize(state, self._epoch_num)
-
-        self._agent.on_ai_step()
-
-        return self.last['state_v'], self.last['reward'], not self._agent.is_training(), {}
-
-    def reset(self):
-        # logging.info("[ai] resetting environment ...")
-        self._epoch_num = 0
-        state = self._next_epoch()
-        self.last['state'] = state
-        self.last['state_v'] = featurizer.featurize(state, 1)
-        return self.last['state_v']
-
-    def _render_histogram(self, hist):
-        for ch in range(featurizer.histogram_size):
-            if hist[ch]:
-                logging.info("      CH %d: %s" % (ch + 1, hist[ch]))
-
-    def render(self, mode='human', close=False, force=False):
-        # when using a vectorialized environment, render gets called twice
-        # avoid rendering the same data
-        if self._last_render == self._epoch_num:
-            return
-
-        if not self._agent.is_training() and not force:
-            return
-
-        self._last_render = self._epoch_num
-
-        logging.info("[AI] --- training epoch %d/%d ---" % (self._epoch_num, self._agent.training_epochs()))
-        logging.info("[AI] REWARD: %f" % self.last['reward'])
-
-        logging.debug(
-            "[AI] policy: %s" % ', '.join("%s:%s" % (name, value) for name, value in self.last['params'].items()))
-
-        logging.info("[AI] observation:")
-        for name, value in self.last['state'].items():
-            if 'histogram' in name:
-                logging.info("    %s" % name.replace('_histogram', ''))
-                self._render_histogram(value)

pwnagotchi/ai/parameter.py

@@ -1,30 +0,0 @@
-from gymnasium import spaces
-
-
-class Parameter(object):
-    def __init__(self, name, value=0.0, min_value=0, max_value=2, meta=None, trainable=True):
-        self.name = name
-        self.trainable = trainable
-        self.meta = meta
-        self.value = value
-        self.min_value = min_value
-        self.max_value = max_value + 1
-
-        # gymnasium.space.Discrete is within [0, 1, 2, ..., n-1]
-        if self.min_value < 0:
-            self.scale_factor = abs(self.min_value)
-        elif self.min_value > 0:
-            self.scale_factor = -self.min_value
-        else:
-            self.scale_factor = 0
-
-    def space_size(self):
-        return self.max_value + self.scale_factor
-
-    def space(self):
-        return spaces.Discrete(self.max_value + self.scale_factor)
-
-    def to_param_value(self, policy_v):
-        self.value = policy_v - self.scale_factor
-        assert self.min_value <= self.value <= self.max_value
-        return int(self.value)

pwnagotchi/ai/train.py

@@ -1,198 +0,0 @@
-# import _thread
-import threading
-import time
-import random
-import os
-import json
-import logging
-
-import pwnagotchi.plugins as plugins
-import pwnagotchi.ai as ai
-
-
-class Stats(object):
-    def __init__(self, path, events_receiver):
-        self._lock = threading.Lock()
-        self._receiver = events_receiver
-
-        self.path = path
-        self.born_at = time.time()
-        # total epochs lived (trained + just eval)
-        self.epochs_lived = 0
-        # total training epochs
-        self.epochs_trained = 0
-
-        self.worst_reward = 0.0
-        self.best_reward = 0.0
-
-        self.load()
-
-    def on_epoch(self, data, training):
-        best_r = False
-        worst_r = False
-        with self._lock:
-            reward = data['reward']
-            if reward < self.worst_reward:
-                self.worst_reward = reward
-                worst_r = True
-
-            elif reward > self.best_reward:
-                best_r = True
-                self.best_reward = reward
-
-            self.epochs_lived += 1
-            if training:
-                self.epochs_trained += 1
-
-        self.save()
-
-        if best_r:
-            self._receiver.on_ai_best_reward(reward)
-        elif worst_r:
-            self._receiver.on_ai_worst_reward(reward)
-
-    def load(self):
-        with self._lock:
-            if os.path.exists(self.path) and os.path.getsize(self.path) > 0:
-                logging.info("[AI] loading %s" % self.path)
-                with open(self.path, 'rt') as fp:
-                    obj = json.load(fp)
-
-                self.born_at = obj['born_at']
-                self.epochs_lived, self.epochs_trained = obj['epochs_lived'], obj['epochs_trained']
-                self.best_reward, self.worst_reward = obj['rewards']['best'], obj['rewards']['worst']
-
-    def save(self):
-        with self._lock:
-            logging.info("[AI] saving %s" % self.path)
-
-            data = json.dumps({
-                'born_at': self.born_at,
-                'epochs_lived': self.epochs_lived,
-                'epochs_trained': self.epochs_trained,
-                'rewards': {
-                    'best': self.best_reward,
-                    'worst': self.worst_reward
-                }
-            })
-
-            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)
-
-
-class AsyncTrainer(object):
-    def __init__(self, config):
-        self._config = config
-        self._model = None
-        self._is_training = False
-        self._training_epochs = 0
-        self._nn_path = self._config['ai']['path']
-        self._stats = Stats("%s.json" % os.path.splitext(self._nn_path)[0], self)
-
-    def set_training(self, training, for_epochs=0):
-        self._is_training = training
-        self._training_epochs = for_epochs
-
-        if training:
-            plugins.on('ai_training_start', self, for_epochs)
-        else:
-            plugins.on('ai_training_end', self)
-
-    def is_training(self):
-        return self._is_training
-
-    def training_epochs(self):
-        return self._training_epochs
-
-    def start_ai(self):
-        #_thread.start_new_thread(self._ai_worker, ())
-        threading.Thread(target=self._ai_worker, args=(), name="AI Worker", daemon=True).start()
-
-    def _save_ai(self):
-        logging.info("[AI] saving model to %s ..." % self._nn_path)
-        temp = "%s.tmp" % self._nn_path
-        self._model.save(temp)
-        os.replace(temp, self._nn_path)
-
-    def on_ai_step(self):
-        self._model.env.render()
-
-        if self._is_training:
-            self._save_ai()
-
-        self._stats.on_epoch(self._epoch.data(), self._is_training)
-
-    def on_ai_training_step(self, _locals, _globals):
-        self._model.env.render()
-        plugins.on('ai_training_step', self, _locals, _globals)
-
-    def on_ai_policy(self, new_params):
-        plugins.on('ai_policy', self, new_params)
-        logging.info("[AI] setting new policy:")
-        for name, value in new_params.items():
-            if name in self._config['personality']:
-                curr_value = self._config['personality'][name]
-                if curr_value != value:
-                    logging.info("[AI] ! %s: %s -> %s" % (name, curr_value, value))
-                    self._config['personality'][name] = value
-            else:
-                logging.error("[AI] param %s not in personality configuration!" % name)
-
-        self.run('set wifi.ap.ttl %d' % self._config['personality']['ap_ttl'])
-        self.run('set wifi.sta.ttl %d' % self._config['personality']['sta_ttl'])
-        self.run('set wifi.rssi.min %d' % self._config['personality']['min_rssi'])
-
-    def on_ai_ready(self):
-        self._view.on_ai_ready()
-        plugins.on('ai_ready', self)
-
-    def on_ai_best_reward(self, r):
-        logging.info("[AI] best reward so far: %s" % r)
-        self._view.on_motivated(r)
-        plugins.on('ai_best_reward', self, r)
-
-    def on_ai_worst_reward(self, r):
-        logging.info("[AI] worst reward so far: %s" % r)
-        self._view.on_demotivated(r)
-        plugins.on('ai_worst_reward', self, r)
-
-    def _ai_worker(self):
-        self._model = ai.load(self._config, self, self._epoch)
-
-        if self._model:
-            self.on_ai_ready()
-
-            epochs_per_episode = self._config['ai']['epochs_per_episode']
-
-            obs = None
-            while True:
-                self._model.env.render()
-                # enter in training mode?
-                if random.random() > self._config['ai']['laziness']:
-                    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)
-                    except Exception as e:
-                        logging.exception("[AI] error while training (%s)", e)
-                    finally:
-                        self.set_training(False)
-                        obs = self._model.env.reset()
-                # init the first time
-                elif obs is None:
-                    obs = self._model.env.reset()
-
-                # run the inference
-                action, _ = self._model.predict(obs)
-                obs, _, _, _ = self._model.env.step(action)

pwnagotchi/ai/utils.py

@@ -1,16 +0,0 @@
-import numpy as np
-
-
-def normalize(v, min_v, max_v):
-    return (v - min_v) / (max_v - min_v)
-
-
-def as_batches(x, y, batch_size, shuffle=True):
-    x_size = len(x)
-    assert x_size == len(y)
-
-    indices = np.random.permutation(x_size) if shuffle else None
-
-    for offset in range(0, x_size - batch_size + 1, batch_size):
-        excerpt = indices[offset:offset + batch_size] if shuffle else slice(offset, offset + batch_size)
-        yield x[excerpt], y[excerpt]