Fix code style with black and isort

a3c/eval.py

@@ -1,4 +1,5 @@
 import os
+
 import torch
 
 from .net import GreedyNet
@@ -14,9 +15,7 @@ def evaluate_checkpoints(dir, env):
             wins, guesses = evaluate(env, pretrained_model_path)
             results[checkpoint] = wins, guesses
     return dict(
-        sorted(results.items(), key=lambda x: (
-            x[1][0], -x[1][1]), reverse=True
-        )
+        sorted(results.items(), key=lambda x: (x[1][0], -x[1][1]), reverse=True)
     )
 
 
@@ -39,4 +38,4 @@ def evaluate(env, pretrained_model_path):
         took {n_win_guesses/n_wins} guesses per win, "
         f"{n_guesses / N} including losses."
     )
-    return n_wins/N*100, n_win_guesses/n_wins
+    return n_wins / N * 100, n_win_guesses / n_wins

a3c/net.py

@@ -1,7 +1,7 @@
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import numpy as np
 
 
 class Net(nn.Module):
@@ -23,15 +23,15 @@ class Net(nn.Module):
         word_array = np.zeros((word_width, len(word_list)))
         for i, word in enumerate(word_list):
             for j, c in enumerate(word):
-                word_array[j*26 + (ord(c) - ord('A')), i] = 1
+                word_array[j * 26 + (ord(c) - ord("A")), i] = 1
         self.words = torch.Tensor(word_array)
 
     def forward(self, x):
         values = self.v1(x.float())
         logits = torch.log_softmax(
-            torch.tensordot(self.actor_head(values), self.words,
-                            dims=((1,), (0,))),
-            dim=-1)
+            torch.tensordot(self.actor_head(values), self.words, dims=((1,), (0,))),
+            dim=-1,
+        )
         values = self.v4(values)
         return logits, values
 

a3c/play.py

@@ -1,15 +1,18 @@
 import os
+
 import torch
 from dotenv import load_dotenv
+
 from wordle_env.state import update_from_mask
+
 from .net import GreedyNet
 from .utils import v_wrap
 
 
 def get_play_model_path():
     load_dotenv()
-    model_name = os.getenv('RS_WORDLE_MODEL_NAME')
-    model_checkpoint_dir = os.path.join('checkpoints', 'best_models')
+    model_name = os.getenv("RS_WORDLE_MODEL_NAME")
+    model_checkpoint_dir = os.path.join("checkpoints", "best_models")
     return os.path.join(model_checkpoint_dir, model_name)
 
 
@@ -28,12 +31,7 @@ def get_initial_state(env):
     return state
 
 
-def suggest(
-        env,
-        words,
-        states,
-        pretrained_model_path
-) -> str:
+def suggest(env, words, states, pretrained_model_path) -> str:
     """
     Given a list of words and masks, return the next suggested word
 

a3c/shared_adam.py

@@ -6,20 +6,18 @@ import torch
 
 
 class SharedAdam(torch.optim.Adam):
-    def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), eps=1e-8,
-                 weight_decay=0):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), eps=1e-8, weight_decay=0):
         super(SharedAdam, self).__init__(
-            params, lr=lr,
-            betas=betas, eps=eps, weight_decay=weight_decay
+            params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay
         )
         # State initialization
         for group in self.param_groups:
-            for p in group['params']:
+            for p in group["params"]:
                 state = self.state[p]
-                state['step'] = 0
-                state['exp_avg'] = torch.zeros_like(p.data)
-                state['exp_avg_sq'] = torch.zeros_like(p.data)
+                state["step"] = 0
+                state["exp_avg"] = torch.zeros_like(p.data)
+                state["exp_avg_sq"] = torch.zeros_like(p.data)
 
                 # share in memory
-                state['exp_avg'].share_memory_()
-                state['exp_avg_sq'].share_memory_()
+                state["exp_avg"].share_memory_()
+                state["exp_avg_sq"].share_memory_()

a3c/train.py

@@ -1,10 +1,12 @@
 import os
-import numpy as np
 import random
+
+import numpy as np
 import torch
 import torch.multiprocessing as mp
-from .shared_adam import SharedAdam
+
 from .net import Net
+from .shared_adam import SharedAdam
 from .worker import Worker
 
 
@@ -25,12 +27,12 @@ def train(
     env,
     max_ep,
     model_checkpoint_dir,
-    gamma=0.,
+    gamma=0.0,
     seed=100,
     pretrained_model_path=None,
     save=False,
     min_reward=9.9,
-    every_n_save=100
+    every_n_save=100,
 ):
     os.environ["OMP_NUM_THREADS"] = "1"
     if not os.path.exists(model_checkpoint_dir):
@@ -45,18 +47,40 @@ def train(
     if pretrained_model_path:
         gnet.load_state_dict(torch.load(pretrained_model_path))
     gnet.share_memory()  # share the global parameters in multiprocessing
-    opt = SharedAdam(gnet.parameters(), lr=1e-4,
-                     betas=(0.92, 0.999))  # global optimizer
-    global_ep, global_ep_r, res_queue, win_ep = mp.Value(
-        'i', 0), mp.Value('d', 0.), mp.Queue(), mp.Value('i', 0)
+    opt = SharedAdam(
+        gnet.parameters(), lr=1e-4, betas=(0.92, 0.999)
+    )  # global optimizer
+    global_ep, global_ep_r, res_queue, win_ep = (
+        mp.Value("i", 0),
+        mp.Value("d", 0.0),
+        mp.Queue(),
+        mp.Value("i", 0),
+    )
 
     # parallel training
     workers = [
         Worker(
-            max_ep, gnet, opt, global_ep, global_ep_r, res_queue, i, env,
-            n_s, n_a, words_list, word_width, win_ep, model_checkpoint_dir,
-            gamma, pretrained_model_path, save, min_reward, every_n_save
-        ) for i in range(mp.cpu_count())
+            max_ep,
+            gnet,
+            opt,
+            global_ep,
+            global_ep_r,
+            res_queue,
+            i,
+            env,
+            n_s,
+            n_a,
+            words_list,
+            word_width,
+            win_ep,
+            model_checkpoint_dir,
+            gamma,
+            pretrained_model_path,
+            save,
+            min_reward,
+            every_n_save,
+        )
+        for i in range(mp.cpu_count())
     ]
     [w.start() for w in workers]
     res = []  # record episode reward to plot
@@ -68,6 +92,8 @@ def train(
             break
     [w.join() for w in workers]
     if save:
-        torch.save(gnet.state_dict(), os.path.join(
-            model_checkpoint_dir, f'model_{env.unwrapped.spec.id}.pth'))
+        torch.save(
+            gnet.state_dict(),
+            os.path.join(model_checkpoint_dir, f"model_{env.unwrapped.spec.id}.pth"),
+        )
     return global_ep, win_ep, gnet, res

a3c/utils.py

@@ -1,5 +1,5 @@
-import torch
 import numpy as np
+import torch
 
 
 def v_wrap(np_array, dtype=np.float32):

a3c/worker.py

@@ -2,40 +2,42 @@
 Worker class implementation of the a3c discrete algorithm
 """
 import os
-import torch
+
 import numpy as np
+import torch
 import torch.multiprocessing as mp
 from torch import nn
+
 from .net import Net
 from .utils import v_wrap
 
 
 class Worker(mp.Process):
     def __init__(
-            self,
-            max_ep,
-            gnet,
-            opt,
-            global_ep,
-            global_ep_r,
-            res_queue,
-            name,
-            env,
-            N_S,
-            N_A,
-            words_list,
-            word_width,
-            winning_ep,
-            model_checkpoint_dir,
-            gamma=0.,
-            pretrained_model_path=None,
-            save=False,
-            min_reward=9.9,
-            every_n_save=100
+        self,
+        max_ep,
+        gnet,
+        opt,
+        global_ep,
+        global_ep_r,
+        res_queue,
+        name,
+        env,
+        N_S,
+        N_A,
+        words_list,
+        word_width,
+        winning_ep,
+        model_checkpoint_dir,
+        gamma=0.0,
+        pretrained_model_path=None,
+        save=False,
+        min_reward=9.9,
+        every_n_save=100,
     ):
         super(Worker, self).__init__()
         self.max_ep = max_ep
-        self.name = 'w%02i' % name
+        self.name = "w%02i" % name
         self.g_ep = global_ep
         self.g_ep_r = global_ep_r
         self.res_queue = res_queue
@@ -57,7 +59,7 @@ class Worker(mp.Process):
         while self.g_ep.value < self.max_ep:
             s = self.env.reset()
             buffer_s, buffer_a, buffer_r = [], [], []
-            ep_r = 0.
+            ep_r = 0.0
             while True:
                 a = self.lnet.choose_action(v_wrap(s[None, :]))
                 s_, r, done, _ = self.env.step(a)
@@ -68,11 +70,9 @@ class Worker(mp.Process):
 
                 if done:  # update global and assign to local net
                     # sync
-                    self.push_and_pull(done, s_, buffer_s,
-                                       buffer_a, buffer_r)
+                    self.push_and_pull(done, s_, buffer_s, buffer_a, buffer_r)
                     goal_word = self.word_list[self.env.goal_word]
-                    self.record(ep_r, goal_word,
-                                self.word_list[a], len(buffer_a))
+                    self.record(ep_r, goal_word, self.word_list[a], len(buffer_a))
                     self.save_model()
                     buffer_s, buffer_a, buffer_r = [], [], []
                     break
@@ -81,22 +81,22 @@ class Worker(mp.Process):
 
     def push_and_pull(self, done, s_, bs, ba, br):
         if done:
-            v_s_ = 0.               # terminal
+            v_s_ = 0.0  # terminal
         else:
-            v_s_ = self.lnet.forward(v_wrap(
-                s_[None, :]))[-1].data.numpy()[0, 0]
+            v_s_ = self.lnet.forward(v_wrap(s_[None, :]))[-1].data.numpy()[0, 0]
 
         buffer_v_target = []
-        for r in br[::-1]:    # reverse buffer r
+        for r in br[::-1]:  # reverse buffer r
             v_s_ = r + self.gamma * v_s_
             buffer_v_target.append(v_s_)
         buffer_v_target.reverse()
 
         loss = self.lnet.loss_func(
             v_wrap(np.vstack(bs)),
-            v_wrap(np.array(ba), dtype=np.int64) if
-            ba[0].dtype == np.int64 else v_wrap(np.vstack(ba)),
-            v_wrap(np.array(buffer_v_target)[:, None])
+            v_wrap(np.array(ba), dtype=np.int64)
+            if ba[0].dtype == np.int64
+            else v_wrap(np.vstack(ba)),
+            v_wrap(np.array(buffer_v_target)[:, None]),
         )
 
         # calculate local gradients and push local parameters to global
@@ -110,16 +110,21 @@ class Worker(mp.Process):
         self.lnet.load_state_dict(self.gnet.state_dict())
 
     def save_model(self):
-        if (self.save and self.g_ep_r.value >= self.min_reward and
-                self.g_ep.value % self.every_n_save == 0):
-            torch.save(self.gnet.state_dict(), os.path.join(
-                self.model_checkpoint_dir, f'model_{self.g_ep.value}.pth'))
+        if (
+            self.save
+            and self.g_ep_r.value >= self.min_reward
+            and self.g_ep.value % self.every_n_save == 0
+        ):
+            torch.save(
+                self.gnet.state_dict(),
+                os.path.join(self.model_checkpoint_dir, f"model_{self.g_ep.value}.pth"),
+            )
 
     def record(self, ep_r, goal_word, action, action_number):
         with self.g_ep.get_lock():
             self.g_ep.value += 1
         with self.g_ep_r.get_lock():
-            if self.g_ep_r.value == 0.:
+            if self.g_ep_r.value == 0.0:
                 self.g_ep_r.value = ep_r
             else:
                 self.g_ep_r.value = self.g_ep_r.value * 0.99 + ep_r * 0.01
@@ -129,9 +134,13 @@ class Worker(mp.Process):
             if self.g_ep.value % 100 == 0:
                 print(
                     self.name,
-                    "Ep:", self.g_ep.value,
+                    "Ep:",
+                    self.g_ep.value,
                     "| Ep_r: %.0f" % self.g_ep_r.value,
-                    "| Goal :", goal_word,
-                    "| Action: ", action,
-                    "| Actions: ", action_number
+                    "| Goal :",
+                    goal_word,
+                    "| Action: ",
+                    action,
+                    "| Actions: ",
+                    action_number,
                 )

api_rest/api.py

@@ -1,30 +1,32 @@
 import random
-from a3c.play import get_play_model_path, play
-from flask import Flask, request, jsonify
+
+from flask import Flask, jsonify, request
 from flask_cors import cross_origin
-from wordle_env.words import target_vocabulary
+
+from a3c.play import get_play_model_path, play
 from wordle_env.wordle import get_env
+from wordle_env.words import target_vocabulary
 
 app = Flask(__name__)
 
 
 def validate_goal_word(word):
     if not word:
-        return True, 'Goal word not provided'
+        return True, "Goal word not provided"
     if word.upper() not in target_vocabulary:
-        return True, 'Goal word not in vocabulary'
-    return False, ''
+        return True, "Goal word not in vocabulary"
+    return False, ""
 
 
-@app.route('/play_word', methods=['GET'])
-@cross_origin(origin='*', headers=['Content-Type', 'Authorization'])
+@app.route("/play_word", methods=["GET"])
+@cross_origin(origin="*", headers=["Content-Type", "Authorization"])
 def get_play():
     # Get the goal word from the request
-    word = request.args.get('goal_word')
+    word = request.args.get("goal_word")
 
     error, msge = validate_goal_word(word)
     if error:
-        return jsonify({'error': msge}), 400
+        return jsonify({"error": msge}), 400
 
     word = word.upper()
     env = get_env()
@@ -32,16 +34,16 @@ def get_play():
     # Call the play function with the goal word
     # and return the attempts and the result
     won, attempts = play(env, model_path, word)
-    return jsonify({'attempts': attempts, 'won': won})
+    return jsonify({"attempts": attempts, "won": won})
 
 
-@app.route('/word', methods=['GET'])
-@cross_origin(origin='*', headers=['Content-Type', 'Authorization'])
+@app.route("/word", methods=["GET"])
+@cross_origin(origin="*", headers=["Content-Type", "Authorization"])
 def get_word():
     # Get a random word from the target vocabulary used to train the model
     word = random.choice(target_vocabulary)
     word = word.upper()
-    return jsonify({'word': word})
+    return jsonify({"word": word})
 
 
 def create_app(settings_override=None):
@@ -58,5 +60,5 @@ def create_app(settings_override=None):
     return app
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     app.run(debug=True)

main.py

@@ -3,23 +3,33 @@
 import argparse
 import os
 import time
+
 import matplotlib.pyplot as plt
-from a3c.train import train
+
 from a3c.eval import evaluate, evaluate_checkpoints
 from a3c.play import suggest
+from a3c.train import train
 from wordle_env.wordle import get_env
 
 
 def training_mode(args, env, model_checkpoint_dir):
     max_ep = args.games
     start_time = time.time()
-    pretrained_model_path = os.path.join(
-        model_checkpoint_dir, args.model_name
-    ) if args.model_name else args.model_name
+    pretrained_model_path = (
+        os.path.join(model_checkpoint_dir, args.model_name)
+        if args.model_name
+        else args.model_name
+    )
     global_ep, win_ep, gnet, res = train(
-        env, max_ep, model_checkpoint_dir, args.gamma,
-        args.seed, pretrained_model_path, args.save,
-        args.min_reward, args.every_n_save
+        env,
+        max_ep,
+        model_checkpoint_dir,
+        args.gamma,
+        args.seed,
+        pretrained_model_path,
+        args.save,
+        args.min_reward,
+        args.every_n_save,
     )
     print("--- %.0f seconds ---" % (time.time() - start_time))
     print_results(global_ep, win_ep, res)
@@ -34,8 +44,8 @@ def evaluation_mode(args, env, model_checkpoint_dir):
 
 def play_mode(args, env, model_checkpoint_dir):
     print("Play mode")
-    words = [word.strip() for word in args.words.split(',')]
-    states = [state.strip() for state in args.states.split(',')]
+    words = [word.strip() for word in args.words.split(",")]
+    states = [state.strip() for state in args.states.split(",")]
     pretrained_model_path = os.path.join(model_checkpoint_dir, args.model_name)
     word = suggest(env, words, states, pretrained_model_path)
     print(word)
@@ -45,8 +55,8 @@ def print_results(global_ep, win_ep, res):
     print("Jugadas:", global_ep.value)
     print("Ganadas:", win_ep.value)
     plt.plot(res)
-    plt.ylabel('Moving average ep reward')
-    plt.xlabel('Step')
+    plt.ylabel("Moving average ep reward")
+    plt.xlabel("Step")
     plt.show()
 
 
@@ -55,90 +65,80 @@ if __name__ == "__main__":
     parser.add_argument(
         "enviroment",
         help="Enviroment (type of wordle game) used for training, \
-            example: WordleEnvFull-v0"
+            example: WordleEnvFull-v0",
     )
     parser.add_argument(
         "--models_dir",
         help="Directory where models are saved (default=checkpoints)",
-        default='checkpoints'
+        default="checkpoints",
     )
-    subparsers = parser.add_subparsers(help='sub-command help')
+    subparsers = parser.add_subparsers(help="sub-command help")
 
     parser_train = subparsers.add_parser(
-        'train',
-        help='Train a model from scratch or train from pretrained model'
+        "train", help="Train a model from scratch or train from pretrained model"
     )
     parser_train.add_argument(
-        "--games",
-        "-g",
-        help="Number of games to train",
-        type=int,
-        required=True
+        "--games", "-g", help="Number of games to train", type=int, required=True
     )
     parser_train.add_argument(
         "--model_name",
         "-m",
         help="If want to train from a pretrained model, \
-            the name of the pretrained model file"
+            the name of the pretrained model file",
     )
     parser_train.add_argument(
         "--gamma",
         help="Gamma hyperparameter (discount factor) value",
         type=float,
-        default=0.
+        default=0.0,
     )
     parser_train.add_argument(
-        "--seed",
-        help="Seed used for random numbers generation",
-        type=int,
-        default=100
+        "--seed", help="Seed used for random numbers generation", type=int, default=100
     )
     parser_train.add_argument(
         "--save",
-        '-s',
+        "-s",
         help="Save instances of the model while training",
-        action='store_true'
+        action="store_true",
     )
     parser_train.add_argument(
         "--min_reward",
         help="The minimun global reward value achieved for saving the model",
         type=float,
-        default=9.9
+        default=9.9,
     )
     parser_train.add_argument(
         "--every_n_save",
         help="Check every n training steps to save the model",
         type=int,
-        default=100
+        default=100,
     )
     parser_train.set_defaults(func=training_mode)
 
     parser_eval = subparsers.add_parser(
-        'eval', help='Evaluate saved models for the enviroment')
+        "eval", help="Evaluate saved models for the enviroment"
+    )
     parser_eval.set_defaults(func=evaluation_mode)
 
     parser_play = subparsers.add_parser(
-        'play',
-        help='Give the model a word and the state result \
-            and the model will try to predict the goal word'
+        "play",
+        help="Give the model a word and the state result \
+            and the model will try to predict the goal word",
     )
     parser_play.add_argument(
-        "--words",
-        "-w",
-        help="List of words played in the wordle game",
-        required=True
+        "--words", "-w", help="List of words played in the wordle game", required=True
     )
     parser_play.add_argument(
         "--states",
         "-st",
         help="List of states returned by playing each of the words",
-        required=True
+        required=True,
     )
     parser_play.add_argument(
         "--model_name",
         "-m",
         help="Name of the pretrained model file thich will play the game",
-        required=True
+        required=True,
     )
     parser_play.set_defaults(func=play_mode)
 

rs_wordle_player/firebase_connector.py

@@ -1,13 +1,12 @@
 import os
-import firebase_admin
-from firebase_admin import credentials
-from firebase_admin import firestore
 from datetime import datetime
-from dotenv import load_dotenv
 
+import firebase_admin
+from dotenv import load_dotenv
+from firebase_admin import credentials, firestore
 
-class FirebaseConnector():
 
+class FirebaseConnector:
     def __init__(self):
         load_dotenv()
         cert_path = self.get_credentials_path()
@@ -20,32 +19,34 @@ class FirebaseConnector():
         return db
 
     def get_credentials_path(self):
-        credentials_path = os.getenv('RS_FIREBASE_CREDENTIALS_PATH')
+        credentials_path = os.getenv("RS_FIREBASE_CREDENTIALS_PATH")
         return credentials_path
 
     def get_user(self):
-        user = os.getenv('RS_WORDLE_USER')
+        user = os.getenv("RS_WORDLE_USER")
         return user
 
     def get_state_from_fb_result(self, firebase_result):
-        result_number_map = {'incorrect': '0',
-                             'misplaced': '1',
-                             'correct': '2'}
+        result_number_map = {"incorrect": "0", "misplaced": "1", "correct": "2"}
         char_result_map = map(
             lambda char_res: result_number_map[char_res], firebase_result
         )
-        return ''.join(char_result_map)
+        return "".join(char_result_map)
 
     def today(self):
-        return datetime.today().strftime('%Y%m%d')
+        return datetime.today().strftime("%Y%m%d")
 
     def today_user_results(self):
-        daily_results_col = 'dailyResults'
+        daily_results_col = "dailyResults"
         currentUser = self.get_user()
         # Execute the query and get the first result
-        docs = self.db.collection(daily_results_col).where(
-            'user.email', '==', currentUser).where(
-            'date', '==', self.today()).limit(1).get()
+        docs = (
+            self.db.collection(daily_results_col)
+            .where("user.email", "==", currentUser)
+            .where("date", "==", self.today())
+            .limit(1)
+            .get()
+        )
         return docs
 
     def today_user_attempts(self):
@@ -53,10 +54,10 @@ class FirebaseConnector():
         attempted_words = []
         if len(docs) > 0:
             doc = docs[0]
-            attempted_words = doc.to_dict().get('attemptedWords')
+            attempted_words = doc.to_dict().get("attemptedWords")
         return attempted_words
 
     def today_word(self):
-        words_col = 'words'
+        words_col = "words"
         doc = self.db.collection(words_col).document(self.today())
-        return doc.get().get('word')
+        return doc.get().get("word")

rs_wordle_player/rs_wordle_player.py

@@ -1,5 +1,6 @@
 from a3c.play import get_play_model_path, suggest
 from wordle_env.wordle import get_env
+
 from .firebase_connector import FirebaseConnector
 from .selenium_player import SeleniumPlayer
 
@@ -17,7 +18,7 @@ def get_attempts(fb_connector):
 
 def is_game_finished(states):
     if states:
-        return states[-1] == '22222' or len(states) == 6
+        return states[-1] == "22222" or len(states) == 6
     return False
 
 
@@ -49,5 +50,5 @@ def play():
     return words, won
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     print(play())

rs_wordle_player/selenium_player.py

@@ -1,5 +1,6 @@
 import os
 import time
+
 from dotenv import load_dotenv
 from selenium import webdriver
 from selenium.common.exceptions import UnexpectedAlertPresentException
@@ -8,8 +9,7 @@ from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 
 
-class SeleniumPlayer():
-
+class SeleniumPlayer:
     def __init__(self):
         self.wordle_url = self.get_wordle_url()
         self.driver = self.get_driver()
@@ -24,22 +24,22 @@ class SeleniumPlayer():
 
     def get_wordle_url(self):
         load_dotenv()
-        return os.getenv('RS_WORDLE_URL')
+        return os.getenv("RS_WORDLE_URL")
 
     def get_credentials(self):
         load_dotenv()
-        username = os.getenv('RS_WORDLE_USER')
-        password = os.getenv('RS_WORDLE_PASSWORD')
+        username = os.getenv("RS_WORDLE_USER")
+        password = os.getenv("RS_WORDLE_PASSWORD")
         return username, password
 
     def logged_in(self):
-        return self.driver.current_url != self.wordle_url + '/login'
+        return self.driver.current_url != self.wordle_url + "/login"
 
     def log_in(self):
         if not self.logged_in():
             time.sleep(2)
-            login_div = self.driver.find_element(By.CLASS_NAME, 'login-button')
-            login_btns = login_div.find_elements(By.TAG_NAME, 'button')
+            login_div = self.driver.find_element(By.CLASS_NAME, "login-button")
+            login_btns = login_div.find_elements(By.TAG_NAME, "button")
             login_btn = login_btns[0]
             login_btn.click()
             time.sleep(10)
@@ -47,32 +47,31 @@ class SeleniumPlayer():
             login_window = self.driver.window_handles[1]
             self.driver.switch_to.window(login_window)
             username, password = self.get_credentials()
-            element = self.driver.find_element(By.ID, 'identifierId')
+            element = self.driver.find_element(By.ID, "identifierId")
             element.send_keys(username)
             element.send_keys(Keys.ENTER)
             time.sleep(10)
-            element = self.driver.find_element(By.NAME, 'password')
+            element = self.driver.find_element(By.NAME, "password")
             element.send_keys(password)
             element.send_keys(Keys.ENTER)
             self.driver.switch_to.window(wordle_window)
             time.sleep(5)
             onboard_div = self.driver.find_element(
-                By.CLASS_NAME,
-                'onboarding-modal-container'
+                By.CLASS_NAME, "onboarding-modal-container"
             )
-            onboard_btn = onboard_div.find_elements(By.TAG_NAME, 'button')
+            onboard_btn = onboard_div.find_elements(By.TAG_NAME, "button")
             onboard_btn[-1].click()
 
     def play_word(self, word):
         try:
-            element = self.driver.find_element(By.TAG_NAME, 'html')
+            element = self.driver.find_element(By.TAG_NAME, "html")
             # simulate typing the letters in the word into the input field
             element.send_keys(word)
             # simulate pressing the Enter key
             element.send_keys(Keys.ENTER)
             time.sleep(5)
         except UnexpectedAlertPresentException:
-            print('Won game alert on screen')
+            print("Won game alert on screen")
 
     def finish(self):
         self.driver.quit()

wordle_env/__init__.py

@@ -1,13 +1,9 @@
-from gym.envs.registration import (
-    registry,
-    register,
-    make,
-    spec,
-    load_env_plugins as _load_env_plugins,
-)
 import os
-from . import wordle
 
+from gym.envs.registration import load_env_plugins as _load_env_plugins
+from gym.envs.registration import make, register, registry, spec
+
+from . import wordle
 
 register(
     id="WordleEnv100OneAction-v0",

wordle_env/const.py

@@ -1,4 +1,4 @@
-WORDLE_CHARS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+WORDLE_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 WORDLE_N = 5
 REWARD = 10
 CHAR_REWARD = 0.1

wordle_env/state.py

@@ -13,11 +13,11 @@ where status has codes
 """
 import collections
 from typing import List, Tuple
+
 import numpy as np
 
 from .const import CHAR_REWARD, WORDLE_CHARS, WORDLE_N
 
-
 WordleState = np.ndarray
 
 
@@ -27,8 +27,8 @@ def get_nvec(max_turns: int):
 
 def new(max_turns: int) -> WordleState:
     return np.array(
-        [max_turns] + [0, 0, 0] * WORDLE_N * len(WORDLE_CHARS),
-        dtype=np.int32)
+        [max_turns] + [0, 0, 0] * WORDLE_N * len(WORDLE_CHARS), dtype=np.int32
+    )
 
 
 def remaining_steps(state: WordleState) -> int:
@@ -40,11 +40,7 @@ SOMEWHERE = 1
 YES = 2
 
 
-def update_from_mask(
-    state: WordleState,
-    word: str,
-    mask: List[int]
-) -> WordleState:
+def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleState:
     """
     return a copy of state that has been updated to new state
 
@@ -84,14 +80,14 @@ def update_from_mask(
             # Need to check this first in case there's prior maybe + yes
             if c in prior_maybe:
                 # Then the maybe could be anywhere except here
-                state[offset+3*i:offset+3*i+3] = [1, 0, 0]
+                state[offset + 3 * i : offset + 3 * i + 3] = [1, 0, 0]
             elif c in prior_yes:
                 # No maybe, definitely a yes,
                 # so it's zero everywhere except the yesses
                 for j in range(WORDLE_N):
                     # Only flip no if previously was maybe
-                    if state[offset + 3 * j:offset + 3 * j + 3][1] == 1:
-                        state[offset + 3 * j:offset + 3 * j + 3] = [1, 0, 0]
+                    if state[offset + 3 * j : offset + 3 * j + 3][1] == 1:
+                        state[offset + 3 * j : offset + 3 * j + 3] = [1, 0, 0]
             else:
                 # Just straight up no
                 _set_all_no(state, offset)
@@ -115,7 +111,7 @@ def get_mask(word: str, goal_word: str) -> List[int]:
                 mask[i] = 1
                 counts[c] -= 1
             else:
-                for j in range(i+1, len(mask)):
+                for j in range(i + 1, len(mask)):
                     if mask[j] == 2:
                         continue
                     mask[j] = 0
@@ -136,11 +132,7 @@ def update_mask(state: WordleState, word: str, goal_word: str) -> WordleState:
     return update_from_mask(state, word, mask)
 
 
-def update(
-    state: WordleState,
-    word: str,
-    goal_word: str
-) -> Tuple[WordleState, float]:
+def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState, float]:
     state = state.copy()
     reward = 0
     state[0] -= 1
@@ -158,8 +150,7 @@ def update(
         cint = ord(c) - ord(WORDLE_CHARS[0])
         offset = 1 + cint * WORDLE_N * 3
         if goal_word[i] != c:
-            if (c in goal_word and
-                    goal_word.count(c) > processed_letters.count(c)):
+            if c in goal_word and goal_word.count(c) > processed_letters.count(c):
                 # Char at position i = no,
                 # and in other positions maybe except it had a value before,
                 # other chars stay as they are
@@ -184,27 +175,27 @@ def _set_if_cero(state, offset, value):
     # but only if it didnt have a value before
     for char_idx in range(0, WORDLE_N * 3, 3):
         char_offset = offset + char_idx
-        if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
-            state[char_offset: char_offset + 3] = value
+        if tuple(state[char_offset : char_offset + 3]) == (0, 0, 0):
+            state[char_offset : char_offset + 3] = value
 
 
 def _set_yes(state, offset, char_int, char_pos):
     # char at position char_pos = yes,
     # all other chars at position char_pos == no
     pos_offset = 3 * char_pos
-    state[offset + pos_offset:offset + pos_offset + 3] = [0, 0, 1]
+    state[offset + pos_offset : offset + pos_offset + 3] = [0, 0, 1]
     for ocint in range(len(WORDLE_CHARS)):
         if ocint != char_int:
             oc_offset = 1 + ocint * WORDLE_N * 3
             yes_index = oc_offset + pos_offset
-            state[yes_index:yes_index + 3] = [1, 0, 0]
+            state[yes_index : yes_index + 3] = [1, 0, 0]
 
 
 def _set_no(state, offset, char_pos):
     # Set offset character = no at char_pos position
-    state[offset + 3 * char_pos:offset + 3 * char_pos + 3] = [1, 0, 0]
+    state[offset + 3 * char_pos : offset + 3 * char_pos + 3] = [1, 0, 0]
 
 
 def _set_all_no(state, offset):
     # Set offset character = no at all positions
-    state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
+    state[offset : offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N

wordle_env/test_wordle.py

@@ -1,7 +1,6 @@
 import pytest
 
-from . import wordle
-from . import state
+from . import state, wordle
 
 TESTWORDS = [
     "APPAA",

wordle_env/wordle.py

@@ -1,24 +1,22 @@
+import random
+from typing import List, Optional
+
 import gym
 from gym import spaces
-from typing import Optional, List
 
 from . import state
-from .const import WORDLE_N, REWARD, WORDLE_CHARS
-
+from .const import REWARD, WORDLE_CHARS, WORDLE_N
 from .words import complete_vocabulary, target_vocabulary
 
-import random
-
 
 def _load_words(
-    limit: Optional[int] = None,
-    complete: Optional[bool] = False
+    limit: Optional[int] = None, complete: Optional[bool] = False
 ) -> List[str]:
     words = complete_vocabulary if complete else target_vocabulary
     return words if not limit else words[:limit]
 
 
-def get_env(env_id='WordleEnvFull-v0'):
+def get_env(env_id="WordleEnvFull-v0"):
     return gym.make(env_id)
 
 
@@ -42,13 +40,16 @@ class WordleEnvBase(gym.Env):
         Initial state with turn 0, all chars Unvisited
     """
 
-    def __init__(self, words: List[str],
-                 max_turns: int = 6,
-                 allowable_words: Optional[int] = None,
-                 mask_based_state_updates: bool = False):
+    def __init__(
+        self,
+        words: List[str],
+        max_turns: int = 6,
+        allowable_words: Optional[int] = None,
+        mask_based_state_updates: bool = False,
+    ):
         assert all(
             len(w) == WORDLE_N for w in words
-        ), f'Not all words of length {WORDLE_N}, {words}'
+        ), f"Not all words of length {WORDLE_N}, {words}"
         self.words = words
         self.max_turns = max_turns
         self.allowable_words = allowable_words
@@ -57,8 +58,7 @@ class WordleEnvBase(gym.Env):
             self.allowable_words = len(self.words)
 
         self.action_space = spaces.Discrete(self.words_as_action_space())
-        self.observation_space = spaces.MultiDiscrete(
-            state.get_nvec(self.max_turns))
+        self.observation_space = spaces.MultiDiscrete(state.get_nvec(self.max_turns))
 
         self.done = True
         self.goal_word: int = -1
@@ -79,15 +79,15 @@ class WordleEnvBase(gym.Env):
         word = self.words[action]
         goal_word = self.words[self.goal_word]
         # assert word in self.words, f'{word} not in words list'
-        self.state, r = self.state_updater(state=self.state,
-                                           word=word,
-                                           goal_word=goal_word)
+        self.state, r = self.state_updater(
+            state=self.state, word=word, goal_word=goal_word
+        )
 
         reward = r
         if action == self.goal_word:
             self.done = True
             # reward = REWARD
-            if state.remaining_steps(self.state) == self.max_turns-1:
+            if state.remaining_steps(self.state) == self.max_turns - 1:
                 reward = 0  # -10*REWARD  # No reward for guessing off the bat
             else:
                 reward = REWARD
@@ -100,7 +100,7 @@ class WordleEnvBase(gym.Env):
     def reset(self):
         self.state = state.new(self.max_turns)
         self.done = False
-        random_word = random.choice(self.words[:self.allowable_words])
+        random_word = random.choice(self.words[: self.allowable_words])
         self.goal_word = self.words.index(random_word)
         return self.state.copy()
 
@@ -121,8 +121,7 @@ class WordleEnv100OneAction(WordleEnvBase):
 
 class WordleEnv100WithMask(WordleEnvBase):
     def __init__(self):
-        super().__init__(words=_load_words(100),
-                         mask_based_state_updates=True)
+        super().__init__(words=_load_words(100), mask_based_state_updates=True)
 
 
 class WordleEnv100TwoAction(WordleEnvBase):
@@ -142,8 +141,7 @@ class WordleEnv100FullAction(WordleEnvBase):
 
 class WordleEnv1000WithMask(WordleEnvBase):
     def __init__(self):
-        super().__init__(words=_load_words(1000),
-                         mask_based_state_updates=True)
+        super().__init__(words=_load_words(1000), mask_based_state_updates=True)
 
 
 class WordleEnv1000FullAction(WordleEnvBase):
@@ -158,5 +156,6 @@ class WordleEnvFull(WordleEnvBase):
 
 class WordleEnvRealWithMask(WordleEnvBase):
     def __init__(self):
-        super().__init__(words=_load_words(), allowable_words=2315,
-                         mask_based_state_updates=True)
+        super().__init__(
+            words=_load_words(), allowable_words=2315, mask_based_state_updates=True
+        )

wordle_env/words.py

@@ -7,7 +7,7 @@ _COMPLETE_VOCABULARY_URL = "https://gist.githubusercontent.com/scholtes/\
 _TARGET_VOCABULARY_URL = "https://gist.githubusercontent.com/scholtes/\
     94f3c0303ba6a7768b47583aff36654d/raw/\
     d9cddf5e16140df9e14f19c2de76a0ef36fd2748/wordle-La.txt"
-_DOWNLOADS_DIR = '.'
+_DOWNLOADS_DIR = "."
 _COMPLETE_VOCABULARY_FILENAME = "complete_vocabulary.txt"
 _TARGET_VOCABULARY_FILENAME = "target_vocabulary.txt"
 
@@ -24,7 +24,11 @@ def _retrieve_vocabulary(url, filename, dir):
 
 
 target_vocabulary = _retrieve_vocabulary(
-    _TARGET_VOCABULARY_URL, _TARGET_VOCABULARY_FILENAME, _DOWNLOADS_DIR)
-complete_vocabulary = _retrieve_vocabulary(
-    _COMPLETE_VOCABULARY_URL, _COMPLETE_VOCABULARY_FILENAME, _DOWNLOADS_DIR
-) + target_vocabulary
+    _TARGET_VOCABULARY_URL, _TARGET_VOCABULARY_FILENAME, _DOWNLOADS_DIR
+)
+complete_vocabulary = (
+    _retrieve_vocabulary(
+        _COMPLETE_VOCABULARY_URL, _COMPLETE_VOCABULARY_FILENAME, _DOWNLOADS_DIR
+    )
+    + target_vocabulary
+)

wordle_game.py

@@ -1,30 +1,28 @@
-from rich.prompt import Prompt
-from rich.console import Console
 from random import choice
-from wordle_env.words import target_vocabulary, complete_vocabulary
 
-SQUARES = {
-    'correct_place': '🟩',
-    'correct_letter': '🟨',
-    'incorrect_letter': '⬛'
-}
+from rich.console import Console
+from rich.prompt import Prompt
+
+from wordle_env.words import complete_vocabulary, target_vocabulary
+
+SQUARES = {"correct_place": "🟩", "correct_letter": "🟨", "incorrect_letter": "⬛"}
 
-WELCOME_MESSAGE = f'\n[white on blue] WELCOME TO WORDLE [/]\n'
+WELCOME_MESSAGE = f"\n[white on blue] WELCOME TO WORDLE [/]\n"
 PLAYER_INSTRUCTIONS = "You may start guessing\n"
 GUESS_STATEMENT = "\nEnter your guess"
 ALLOWED_GUESSES = 6
 
 
 def correct_place(letter):
-    return f'[black on green]{letter}[/]'
+    return f"[black on green]{letter}[/]"
 
 
 def correct_letter(letter):
-    return f'[black on yellow]{letter}[/]'
+    return f"[black on yellow]{letter}[/]"
 
 
 def incorrect_letter(letter):
-    return f'[black on white]{letter}[/]'
+    return f"[black on white]{letter}[/]"
 
 
 def check_guess(guess, answer):
@@ -34,19 +32,20 @@ def check_guess(guess, answer):
     for i, letter in enumerate(guess):
         if answer[i] == guess[i]:
             guessed[i] = correct_place(letter)
-            wordle_pattern.append(SQUARES['correct_place'])
+            wordle_pattern.append(SQUARES["correct_place"])
             processed_letters.append(letter)
     for i, letter in enumerate(guess):
         if answer[i] != guess[i]:
-            if (letter in answer and
-                    answer.count(letter) > processed_letters.count(letter)):
+            if letter in answer and answer.count(letter) > processed_letters.count(
+                letter
+            ):
                 guessed[i] = correct_letter(letter)
-                wordle_pattern.append(SQUARES['correct_letter'])
+                wordle_pattern.append(SQUARES["correct_letter"])
             else:
                 guessed[i] = incorrect_letter(letter)
-                wordle_pattern.append(SQUARES['incorrect_letter'])
+                wordle_pattern.append(SQUARES["incorrect_letter"])
             processed_letters.append(letter)
-    return ''.join(guessed), ''.join(wordle_pattern)
+    return "".join(guessed), "".join(wordle_pattern)
 
 
 def game(console, chosen_word):
@@ -57,12 +56,15 @@ def game(console, chosen_word):
 
     while not end_of_game:
         guess = Prompt.ask(GUESS_STATEMENT).upper()
-        while (len(guess) != 5 or guess in already_guessed or
-               guess not in complete_vocabulary):
+        while (
+            len(guess) != 5
+            or guess in already_guessed
+            or guess not in complete_vocabulary
+        ):
             if guess in already_guessed:
                 console.print("[red]You've already guessed this word!!\n[/]")
             else:
-                console.print('[red]Please enter a valid 5-letter word!!\n[/]')
+                console.print("[red]Please enter a valid 5-letter word!!\n[/]")
             guess = Prompt.ask(GUESS_STATEMENT).upper()
         already_guessed.append(guess)
         guessed, pattern = check_guess(guess, chosen_word)
@@ -74,14 +76,13 @@ def game(console, chosen_word):
             end_of_game = True
     if len(already_guessed) == ALLOWED_GUESSES and guess != chosen_word:
         console.print(f"\n[red]WORDLE X/{ALLOWED_GUESSES}[/]")
-        console.print(f'\n[green]Correct Word: {chosen_word}[/]')
+        console.print(f"\n[green]Correct Word: {chosen_word}[/]")
     else:
-        console.print(
-            f"\n[green]WORDLE {len(already_guessed)}/{ALLOWED_GUESSES}[/]\n")
+        console.print(f"\n[green]WORDLE {len(already_guessed)}/{ALLOWED_GUESSES}[/]\n")
     console.print(*full_wordle_pattern, sep="\n")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     console = Console()
     chosen_word = choice(target_vocabulary)
     console.print(WELCOME_MESSAGE)