Fix styles and deleted unused code of wordle_env

wordle_env/__init__.py

@@ -7,9 +7,6 @@ from gym.envs.registration import (
 )
 import os
 from . import wordle
-# Classic
-# ----------------------------------------
-current_dir =  os.path.realpath(os.path.dirname(__file__))
 
 
 register(

wordle_env/state.py

@@ -2,13 +2,11 @@
 Keep the state in a 1D int array
 
 index[0] = remaining steps
-Rest of data is laid out as binary array
-
-[1..27] = whether char has been guessed or not
 
 [[status, status, status, status, status]
  for _ in "ABCD..."]
 where status has codes
+ [0, 0, 0] - no information about the char
  [1, 0, 0] - char is definitely not in this spot
  [0, 1, 0] - char is maybe in this spot
  [0, 0, 1] - char is definitely in this spot
@@ -126,6 +124,7 @@ def get_mask(word: str, goal_word: str) -> List[int]:
 
     return mask
 
+
 def update_mask(state: WordleState, word: str, goal_word: str) -> WordleState:
     """
     return a copy of state that has been updated to new state
@@ -171,4 +170,3 @@ def update(state: WordleState, word: str, goal_word: str) -> WordleState:
                 state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
             processed_letters.append(c)
     return state, reward
-

wordle_env/wordle.py

@@ -12,9 +12,11 @@ from .words import complete_vocabulary, target_vocabulary
 
 import random
 
-def _load_words(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 _load_words(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]
+
 
 class WordleEnvBase(gym.Env):
     """
@@ -23,21 +25,21 @@ class WordleEnvBase(gym.Env):
         * 13k for full vocab
     State space is defined as:
         * 6 possibilities for turns (WORDLE_TURNS)
-        * Each VALID_CHAR has a state of 0/1 for whether it's been guessed before
         * For each in VALID_CHARS [A-Z] can be in one of 3^WORDLE_N states: (No, Maybe, Yes)
         for full game, this is (3^5)^26
         Each state has 1 + 5*26 possibilities
     Reward:
         Reward is 10 for guessing the right word, -10 for not guessing the right word after 6 guesses.
+        1 from every letter correctly guessed on each try
     Starting State:
         Random goal word
-        Initial state with turn 0, all chars Unvisited + Maybe
+        Initial state with turn 0, all chars Unvisited
     """
+
     def __init__(self, words: List[str],
-                 max_turns: int=6,
-                 allowable_words: Optional[int]=None,
-                 frequencies: Optional[List[float]]=None,
-                 mask_based_state_updates: bool=False):
+                 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}'
         self.words = words
         self.max_turns = max_turns
@@ -46,16 +48,11 @@ class WordleEnvBase(gym.Env):
         if not self.allowable_words:
             self.allowable_words = len(self.words)
 
-        self.frequencies = None
-        if frequencies:
-            assert len(words) == len(frequencies), f'{len(words), len(frequencies)}'
-            self.frequencies = np.array(frequencies, dtype=np.float32) / sum(frequencies)
-
         self.action_space = spaces.Discrete(self.words_as_action_space())
         self.observation_space = spaces.MultiDiscrete(state.get_nvec(self.max_turns))
 
         self.done = True
-        self.goal_word: Tuple = tuple(tuple([tuple([-1]) * WORDLE_N]) *len(WORDLE_CHARS))
+        self.goal_word: int = -1
 
         self.state: state.WordleState = None
         self.state_updater = state.update
@@ -74,17 +71,17 @@ class WordleEnvBase(gym.Env):
         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)
+                                           word=word,
+                                           goal_word=goal_word)
 
         reward = r
         if action == self.goal_word:
             self.done = True
-            #reward = REWARD
+            # reward = REWARD
             if state.remaining_steps(self.state) == self.max_turns-1:
-                reward = 0#-10*REWARD  # No reward for guessing off the bat
+                reward = 0  # -10*REWARD  # No reward for guessing off the bat
             else:
-                #reward = REWARD*(self.state.remaining_steps() + 1) / self.max_turns
+                # reward = REWARD*(self.state.remaining_steps() + 1) / self.max_turns
                 reward = REWARD
         elif state.remaining_steps(self.state) == 0:
             self.done = True
@@ -108,25 +105,6 @@ class WordleEnvBase(gym.Env):
     def words_as_action_space(self):
         return len(self.words)
 
-    def encode_word(self, word):
-        encoded_word = np.array(
-            [[0] * WORDLE_N] * len(WORDLE_CHARS),
-            dtype=np.int32
-        )
-        for index, letter in enumerate(word):
-            cint = WORDLE_CHARS.index(letter)
-            encoded_word[cint][index] = 1
-        return encoded_word
-
-    def decode_word(self, action):
-        word = [''] * WORDLE_N
-        for index, letter_vec in enumerate(action):
-            if 1 in letter_vec:
-                for i, j in enumerate(letter_vec):
-                    if j == 1:
-                        word[i] = WORDLE_CHARS[index]
-        return ''.join(word)
-
 
 class WordleEnv100OneAction(WordleEnvBase):
     def __init__(self):