diff --git a/pwnagotchi/ai/featurizer.py b/pwnagotchi/ai/featurizer.py index f8730f7d9..d8f27a743 100644 --- a/pwnagotchi/ai/featurizer.py +++ b/pwnagotchi/ai/featurizer.py @@ -4,31 +4,37 @@ 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 describe(extended=False): + if not extended: + histogram_size = wifi.NumChannels + else: + # see https://github.com/evilsocket/pwnagotchi/issues/583 + histogram_size = wifi.NumChannelsExt + + return histogram_size, (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): diff --git a/pwnagotchi/ai/gym.py b/pwnagotchi/ai/gym.py index 22e64417c..de5c94ce1 100644 --- a/pwnagotchi/ai/gym.py +++ b/pwnagotchi/ai/gym.py @@ -34,10 +34,14 @@ def __init__(self, agent, epoch): self._epoch_num = 0 self._last_render = None - channels = agent.supported_channels() + # see https://github.com/evilsocket/pwnagotchi/issues/583 + self._supported_channels = agent.supported_channels() + self._extended_spectrum = any(ch > 140 for ch in self._supported_channels) + self._histogram_size, self._observation_shape = featurizer.describe(self._extended_spectrum) + 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 + range(self._histogram_size) if ch + 1 in self._supported_channels ] self.last = { @@ -50,7 +54,7 @@ def __init__(self, agent, epoch): } 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.observation_space = spaces.Box(low=0, high=1, shape=self._observation_shape, dtype=np.float32) self.reward_range = reward.range @staticmethod @@ -118,7 +122,7 @@ def reset(self): return self.last['state_v'] def _render_histogram(self, hist): - for ch in range(featurizer.histogram_size): + for ch in range(self._histogram_size): if hist[ch]: logging.info(" CH %d: %s" % (ch + 1, hist[ch])) diff --git a/pwnagotchi/mesh/wifi.py b/pwnagotchi/mesh/wifi.py index 3bc714b04..ea81022c6 100644 --- a/pwnagotchi/mesh/wifi.py +++ b/pwnagotchi/mesh/wifi.py @@ -1,4 +1,5 @@ NumChannels = 140 +NumChannelsExt = 165 # see https://github.com/evilsocket/pwnagotchi/issues/583 def freq_to_channel(freq):