From d9bc67000a31c66a8e29b335001060b7d5541387 Mon Sep 17 00:00:00 2001 From: Jeroen Oudshoorn Date: Sat, 7 Oct 2023 00:46:05 +0200 Subject: [PATCH] v2.4.7 Signed-off-by: Jeroen Oudshoorn --- pwnagotchi/ai/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pwnagotchi/ai/__init__.py b/pwnagotchi/ai/__init__.py index 073f699e..ad41fc10 100644 --- a/pwnagotchi/ai/__init__.py +++ b/pwnagotchi/ai/__init__.py @@ -21,15 +21,18 @@ def load(config, agent, epoch, from_disk=True): SB_BACKEND = "stable_baselines3" try: - from stable_baselines3.common.vec_env import DummyVecEnv - from stable_baselines3.a2c import a2c + from stable_baselines3 import A2C logging.debug("[ai] A2C imported in %.2fs" % (time.time() - start)) start = time.time() - from stable_baselines3.a2c.policies import MlpPolicy + 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)) + except Exception as e: logging.debug("[ai] stable_baselines3 not accessible. Trying stable_baselines")