upload wordle env, fix indexing issue in wordle env, attempt to improve reward (no improvement)

dqn_wordle.py

@@ -1,38 +0,0 @@
-import gym
-import sys
-from stable_baselines3 import DQN
-from stable_baselines3.common.env_util import make_vec_env
-import wordle_gym
-import numpy as np
-from tqdm import tqdm
-
-def train (model, env, total_timesteps = 100000): 
-    model.learn(total_timesteps=total_timesteps, progress_bar=True)
-    model.save("dqn_wordle")
-
-def test(model, env, test_num=1000):
-
-    total_correct = 0
-
-    for i in tqdm(range(test_num)):
-
-        model = DQN.load("dqn_wordle")
-
-        env = gym.make("wordle-v0")
-        obs = env.reset()
-        done = False
-        while not done:
-            action, _states = model.predict(obs)
-            obs, rewards, done, info = env.step(action)
-            
-    return total_correct / test_num
-
-if __name__ == "__main__":
-    
-    env = gym.make("wordle-v0")
-    model = DQN("MlpPolicy", env, verbose=0)
-    print(env)
-    print(model)
-
-    train(model, env, total_timesteps=500000)
-    print(test(model, env))

gym_wordle/__init__.py

@@ -0,0 +1,7 @@
+from gym.envs.registration import register
+from .wordle import WordleEnv
+
+register(
+    id='Wordle-v0',
+    entry_point='gym_wordle.wordle:WordleEnv'
+)

gym_wordle/utils.py

@@ -0,0 +1,93 @@
+import numpy as np
+import numpy.typing as npt
+
+from pathlib import Path
+
+
+_chars = ' abcdefghijklmnopqrstuvwxyz'
+_char_d = {c: i for i, c in enumerate(_chars)}
+
+
+def to_english(array: npt.NDArray[np.int64]) -> str:
+    """Converts a numpy integer array into a corresponding English string.
+
+    Args:
+        array: Word in array (int) form. It is assumed that each integer in the
+          array is between 0,...,26 (inclusive).
+
+    Returns:
+        A (lowercase) string representation of the word.    
+    """
+    return ''.join(_chars[i] for i in array)
+
+
+def to_array(word: str) -> npt.NDArray[np.int64]:
+    """Converts a string of characters into a corresponding numpy array.
+
+    Args:
+        word: Word in string form. It is assumed that each character in the
+          string is either an empty space ' ' or lowercase alphabetical
+          character.
+
+    Returns:
+        An array representation of the word.
+    """
+    return np.array([_char_d[c] for c in word])
+
+
+def get_words(category: str, build: bool=False) -> npt.NDArray[np.int64]:
+    """Loads a list of words in array form. 
+
+    If specified, this will recompute the list from the human-readable list of
+    words, and save the results in array form.
+
+    Args:
+        category: Either 'guess' or 'solution', which corresponds to the list
+          of acceptable guess words and the list of acceptable solution words.
+        build: If True, recomputes and saves the array-version of the computed
+          list for future access.
+
+    Returns:
+        An array representation of the list of words specified by the category.
+        This array has two dimensions, and the number of columns is fixed at
+        five.
+    """
+    assert category in {'guess', 'solution'}
+    
+    arr_path = Path(__file__).parent / f'dictionary/{category}_list.npy'
+    if build:
+       list_path = Path(__file__).parent / f'dictionary/{category}_list.csv'
+
+       with open(list_path, 'r') as f:
+           words = np.array([to_array(line.strip()) for line in f])
+           np.save(arr_path, words)
+
+    return np.load(arr_path)
+
+
+def play():
+    """Play Wordle yourself!"""
+    import gym
+    import gym_wordle
+    
+    env = gym.make('Wordle-v0')  # load the environment
+    
+    env.reset()
+    solution = to_english(env.unwrapped.solution_space[env.solution]).upper()  # no peeking!
+
+    done = False
+    
+    while not done:
+        action = -1 
+
+        # in general, the environment won't be forgiving if you input an
+        # invalid word, but for this function I want to let you screw up user
+        # input without consequence, so just loops until valid input is taken
+        while not env.action_space.contains(action):
+            guess = input('Guess: ')
+            action = env.unwrapped.action_space.index_of(to_array(guess))
+    
+        state, reward, done, info = env.step(action)
+        env.render()
+    
+    print(f"The word was {solution}")

gym_wordle/wordle.py

@@ -0,0 +1,296 @@
+import gym
+import numpy as np
+import numpy.typing as npt
+from sty import fg, bg, ef, rs
+
+from collections import Counter 
+from gym_wordle.utils import to_english, to_array, get_words
+from typing import Optional 
+
+class WordList(gym.spaces.Discrete):
+    """Super class for defining a space of valid words according to a specified
+    list.
+
+    TODO: Fix these paragraphs
+    The space is a subclass of gym.spaces.Discrete, where each element
+    corresponds to an index of a valid word in the word list. The obfuscation
+    is necessary for more direct implementation of RL algorithms, which expect
+    spaces of less sophisticated form.
+
+    In addition to the default methods of the Discrete space, it implements
+    a __getitem__ method for easy index lookup, and an index_of method to
+    convert potential words into their corresponding index (if they exist).
+    """
+
+    def __init__(self, words: npt.NDArray[np.int64], **kwargs):
+        """
+        Args:
+            words: Collection of words in array form with shape (_, 5), where
+              each word is a row of the array. Each array element is an integer
+              between 0,...,26 (inclusive).
+            kwargs: See documentation for gym.spaces.MultiDiscrete
+        """
+        super().__init__(words.shape[0], **kwargs)
+        self.words = words
+
+    def __getitem__(self, index: int) -> npt.NDArray[np.int64]:
+        """Obtains the (int-encoded) word associated with the given index.
+
+        Args:
+            index: Index for the list of words.
+
+        Returns:
+            Associated word at the position specified by index.
+        """
+        return self.words[index]
+
+    def index_of(self, word: npt.NDArray[np.int64]) -> int:
+        """Given a word, determine its index in the list (if it exists),
+        otherwise returning -1 if no index exists.
+
+        Args:
+            word: Word to find in the word list.
+
+        Returns:
+            The index of the given word if it exists, otherwise -1.
+        """
+        try:
+            index, = np.nonzero((word == self.words).all(axis=1))
+            return index[0]
+        except:
+            return -1
+
+
+class SolutionList(WordList):
+    """Space for *solution* words to the Wordle environment.
+
+    In the game Wordle, there are two different collections of words:
+
+        * "guesses", which the game accepts as valid words to use to guess the
+          answer.
+        * "solutions", which the game uses to choose solutions from.
+
+    Of course, the set of solutions is a strict subset of the set of guesses.
+
+    Reference: https://fivethirtyeight.com/features/when-the-riddler-met-wordle/
+
+    This class represents the set of solution words.
+    """
+    def __init__(self, **kwargs):
+        """
+        Args:
+            kwargs: See documentation for gym.spaces.MultiDiscrete
+        """
+        words = get_words('solution')
+        super().__init__(words, **kwargs)
+
+
+class WordleObsSpace(gym.spaces.Box):
+    """Implementation of the state (observation) space in terms of gym
+    primatives, in this case, gym.spaces.Box.
+
+    The Wordle observation space can be thought of as a 6x5 array with two
+    channels:
+
+      - the character channel, indicating which characters are placed on the
+        board (unfilled rows are marked with the empty character, 0)
+      - the flag channel, indicating the in-game information associated with
+        each character's placement (green highlight, yellow highlight, etc.)
+
+    where there are 6 rows, one for each turn in the game, and 5 columns, since
+    the solution will always be a word of length 5.
+
+    For simplicity, and compatibility with the stable_baselines algorithms,
+    this multichannel is modeled as a 6x10 array, where the two channels are
+    horizontally appended (along columns). Thus each row in the observation
+    should be interpreted as 
+
+        c0 c1 c2 c3 c4 f0 f1 f2 f3 f4
+
+    when the word is c0...c4 and its associated flags are f0...f4.
+
+    While the superclass method `sample` is available to the WordleObsSpace, it
+    should be emphasized that the output of `sample` will (almost surely) not
+    correspond to a real game configuration, because the sampling is not out of
+    possible game configurations. Instead, the Box superclass just samples the
+    integer array space uniformly.
+    """
+
+    def __init__(self, **kwargs):
+        self.n_rows = 6
+        self.n_cols = 5
+        self.max_char = 26
+        self.max_flag = 4
+
+        low = np.zeros((self.n_rows, 2*self.n_cols))
+        high = np.c_[np.full((self.n_rows, self.n_cols), self.max_char),
+                     np.full((self.n_rows, self.n_cols), self.max_flag)]
+
+        super().__init__(low, high, dtype=np.int64, **kwargs)
+
+
+class GuessList(WordList):
+    """Space for *solution* words to the Wordle environment.
+
+    In the game Wordle, there are two different collections of words:
+
+        * "guesses", which the game accepts as valid words to use to guess the
+          answer.
+        * "solutions", which the game uses to choose solutions from.
+
+    Of course, the set of solutions is a strict subset of the set of guesses.
+
+    Reference: https://fivethirtyeight.com/features/when-the-riddler-met-wordle/
+
+    This class represents the set of guess words.
+    """
+    def __init__(self, **kwargs):
+        """
+        Args:
+            kwargs: See documentation for gym.spaces.MultiDiscrete
+        """
+        words = get_words('guess')
+        super().__init__(words, **kwargs)
+
+
+class WordleEnv(gym.Env):
+
+    metadata = {'render.modes': ['human']}
+
+    # character flag codes
+    no_char = 0
+    right_pos = 1
+    wrong_pos = 2
+    wrong_char = 3
+
+    def __init__(self):
+        super().__init__()
+
+        self.seed()
+        self.action_space = GuessList()
+        self.solution_space = SolutionList()
+
+        self.observation_space = WordleObsSpace()
+
+        self._highlights = {
+            self.right_pos: (bg.green, bg.rs),
+            self.wrong_pos: (bg.yellow, bg.rs),
+            self.wrong_char: ('', ''),
+            self.no_char: ('', ''),
+        }
+
+        self.n_rounds = 6
+        self.n_letters = 5
+
+    def _highlighter(self, char: str, flag: int) -> str:
+        """Terminal renderer functionality. Properly highlights a character
+        based on the flag associated with it.
+
+        Args:
+            char: Character in question.
+            flag: Associated flag, one of:
+                - 0: no character (render no background)
+                - 1: right position (render green background)
+                - 2: wrong position (render yellow background)
+                - 3: wrong character (render no background)
+
+        Returns:
+            Correct ASCII sequence producing the desired character in the
+            correct background.
+        """
+        front, back = self._highlights[flag]
+        return front + char + back
+
+    def reset(self):
+        self.round = 0
+        self.solution = self.solution_space.sample()
+
+        self.state = np.zeros((self.n_rounds, 2 * self.n_letters), dtype=np.int64)
+
+        return self.state
+
+    def render(self, mode: str ='human'):
+        """Renders the Wordle environment.
+
+        Currently supported render modes:
+
+        - human: renders the Wordle game to the terminal.
+
+        Args:
+            mode: the mode to render with
+        """
+        if mode == 'human':
+            for row in self.states:
+                text = ''.join(map(
+                    self._highlighter, 
+                    to_english(row[:self.n_letters]).upper(), 
+                    row[self.n_letters:]
+                ))
+
+                print(text)
+        else:
+            super(WordleEnv, self).render(mode=mode)
+                
+    def step(self, action):
+        """Run one step of the Wordle game. Every game must be previously
+        initialized by a call to the `reset` method.
+
+        Args:
+            action: Word guessed by the agent.
+        Returns:
+            state (object): Wordle game state after the guess.
+            reward (float): Reward associated with the guess (-1 for incorrect,
+              0 for correct)
+            done (bool): Whether the game has ended (by a correct guess or
+              after six guesses).
+            info (dict): Auxiliary diagnostic information (empty).
+        """
+        assert self.action_space.contains(action), 'Invalid word!'
+
+        # transform the action, solution indices to their words
+        action = self.action_space[action]  
+        solution = self.solution_space[self.solution]
+
+        # populate the word chars into the row (character channel)
+        self.state[self.round][:self.n_letters] = action
+
+        # populate the flag characters into the row (flag channel)
+        counter = Counter()
+        for i, char in enumerate(action):
+            flag_i = i + self.n_letters  # starts at 5
+            counter[char] += 1
+
+            if char == solution[i]:  # character is in correct position
+                self.state[self.round, flag_i] = self.right_pos
+            elif counter[char] <= (char == solution).sum():  
+                # current character has been seen within correct number of
+                # occurrences 
+                self.state[self.round, flag_i] = self.wrong_pos
+            else:
+                # wrong character, or "correct" character too many times
+                self.state[self.round, flag_i] = self.wrong_char
+
+        self.round += 1
+
+        correct = (action == solution).all()
+        game_over = (self.round == self.n_rounds)
+
+        done = correct or game_over
+
+        # Total reward equals -(number of incorrect guesses)
+        # reward = 0. if correct else -1.
+        
+        # correct +10
+        # guesses new letter +1
+        # guesses correct letter +1
+        # spent another guess -1
+
+        reward = 0
+        reward += np.sum(self.state[:, 5:] == 1) * 1
+        reward += np.sum(self.state[:, 5:] == 2) * 0.5
+        reward += np.sum(self.state[:, 5:] == 3) * -1
+        reward += 10 if correct else -10 if done else 0
+
+        info = {'correct': correct}
+
+        return self.state, reward, done, info

wordle_gym/__init__.py

@@ -1,9 +0,0 @@
-from gym.envs.registration import register
-
-register(
-    id="wordle-v0", entry_point="wordle_gym.envs.wordle_env:WordleEnv",
-)
-
-register(
-    id="wordle-alpha-v0", entry_point="wordle_gym.envs.wordle_alpha_env:WordleEnv",
-)

wordle_gym/envs/strategies/base.py

@@ -1,15 +0,0 @@
-from enum import Enum
-
-from typing import List
-
-class StrategyType(Enum):
-    RANDOM = 1
-    ELIMINATION = 2
-    PROBABILITY = 3
-
-class Strategy:
-    def __init__(self, type: StrategyType):
-        self.type = type
-    
-    def get_best_word(self, guesses: List[List[str]], state: List[List[int]]):
-        raise NotImplementedError("Strategy.get_best_word() not implemented")

wordle_gym/envs/strategies/elimination.py

@@ -1,2 +0,0 @@
-def get_best_word(state):
-    

wordle_gym/envs/strategies/probabilistic.py

@@ -1,20 +0,0 @@
-from random import sample
-from typing import List
-
-from base import Strategy
-from base import StrategyType
-
-from utils import freq
-
-class Random(Strategy):
-    def __init__(self):
-        self.words = freq.get_5_letter_word_freqs()
-        super().__init__(StrategyType.RANDOM)
-
-    def get_best_word(self, state: List[List[int]]):
-        
-
-
-if __name__ == "__main__":
-    r = Random()
-    print(r.get_best_word([]))

wordle_gym/envs/strategies/rand.py

@@ -1,29 +0,0 @@
-from random import sample
-from typing import List
-
-from base import Strategy
-from base import StrategyType
-
-from utils import freq
-
-class Random(Strategy):
-    def __init__(self):
-        self.words = freq.get_5_letter_word_freqs()
-        super().__init__(StrategyType.RANDOM)
-
-    def get_best_word(self, guesses: List[List[str]], state: List[List[int]]):
-        correct_letters = []
-        regex = ""
-        for g, s in zip(guesses, state):
-            for c, s in zip(g, s):
-                if s == 2:
-                    correct_letters.append(c)
-                    regex += c
-                
-
-
-
-
-if __name__ == "__main__":
-    r = Random()
-    print(r.get_best_word([]))

wordle_gym/envs/strategies/utils/freq.py

@@ -1,27 +0,0 @@
-from os import path
-
-def get_5_letter_word_freqs():
-    """
-    Returns a list of words with 5 letters.
-    """
-    FILEPATH = path.join(path.dirname(path.abspath(__file__)), "data/norvig.txt")
-    lines = read_file(FILEPATH)
-    return {k:v for k, v in get_freq(lines).items() if len(k) == 5}
-
-
-def read_file(filename):
-    """
-    Reads a file and returns a list of words and frequencies
-    """
-    with open(filename, 'r') as f:
-        return f.readlines()
-
-
-def get_freq(lines):
-    """
-    Returns a dictionary of words and their frequencies
-    """
-    freqs = {}
-    for word, freq in map(lambda x: x.split("\t"), lines):
-        freqs[word] = int(freq)
-    return freqs

wordle_gym/envs/wordle_env.py

@@ -1,131 +0,0 @@
-import os
-
-
-import gym
-from gym import error, spaces, utils
-from gym.utils import seeding
-
-from enum import Enum
-from collections import Counter
-import numpy as np
-
-WORD_LENGTH = 5
-TOTAL_GUESSES = 6
-SOLUTION_PATH = "../words/solution.csv"
-VALID_WORDS_PATH = "../words/guess.csv"
-
-class LetterState(Enum):
-    ABSENT = 0
-    PRESENT = 1
-    CORRECT_POSITION = 2
-
-
-class WordleEnv(gym.Env):
-    metadata = {"render.modes": ["human"]}
-
-    def _current_path(self):
-        return os.path.dirname(os.path.abspath(__file__))
-
-    def _read_solutions(self):
-        return open(os.path.join(self._current_path(), SOLUTION_PATH)).read().splitlines()
-    
-    def _get_valid_words(self):
-        words = []
-        for word in open(os.path.join(self._current_path(), VALID_WORDS_PATH)).read().splitlines():
-            words.append((word, Counter(word)))
-        return words
-
-    def get_valid(self):
-        return self._valid_words
-
-    def __init__(self):
-        self._solutions = self._read_solutions()
-        self._valid_words = self._get_valid_words()
-        self.action_space = spaces.Discrete(len(self._valid_words))
-        self.observation_space = spaces.MultiDiscrete([3] * TOTAL_GUESSES * WORD_LENGTH)
-        np.random.seed(0)
-        self.reset()
-    
-    def _check_guess(self, guess, guess_counter):
-        c = guess_counter & self.solution_ct
-        result = []
-        correct = True
-        reward = 0
-        for i, char in enumerate(guess):
-            if c.get(char, 0) > 0:
-                if self.solution[i] == char:
-                    result.append(2)
-                    reward += 2
-                else:
-                    result.append(1)
-                    correct = False
-                    reward += 1
-                c[char] -= 1
-            else:
-                result.append(0)
-                correct = False
-        return result, correct, reward
-
-    def step(self, action):
-        """
-        action: index of word in valid_words
-
-        returns:
-            observation: (TOTAL_GUESSES, WORD_LENGTH)
-            reward: 0 if incorrect, 1 if correct, -1 if game over w/o final answer being obtained
-            done: True if game over, w/ or w/o correct answer
-            additional_info: empty
-        """
-        guess, guess_counter = self._valid_words[action]
-        if guess in self.guesses:
-            return self.obs, -1, False, {}
-        self.guesses.append(guess)
-        result, correct, reward = self._check_guess(guess, guess_counter)
-        done = False
-
-        for i in range(self.guess_no*WORD_LENGTH, self.guess_no*WORD_LENGTH + WORD_LENGTH):
-            self.obs[i] = result[i - self.guess_no*WORD_LENGTH]
-        
-        self.guess_no += 1
-        if correct:
-            done = True
-            reward = 1200
-        if self.guess_no == TOTAL_GUESSES:
-            done = True
-            if not correct:
-                reward = -15
-        return self.obs, reward, done, {}
-
-    def reset(self):
-        self.solution = self._solutions[np.random.randint(len(self._solutions))]
-        self.solution_ct = Counter(self.solution)
-        self.guess_no = 0
-        self.guesses = []
-        self.obs = np.zeros((TOTAL_GUESSES * WORD_LENGTH, ))
-        return self.obs
-
-    def render(self, mode="human"):
-        m = {
-            0: "⬜",
-            1: "🟨",
-            2: "🟩"
-        }
-        print("Solution:", self.solution)
-        for g, o in zip(self.guesses, np.reshape(self.obs, (TOTAL_GUESSES, WORD_LENGTH))):
-            o_n = "".join(map(lambda x: m[x], o))
-            print(g, o_n)
-
-    def close(self):
-        pass
-
-if __name__ == "__main__":
-    env = WordleEnv()
-    print(env.action_space)
-    print(env.observation_space)
-    print(env.solution)
-    print(env.step(0))
-    print(env.step(0))
-    print(env.step(0))
-    print(env.step(0))
-    print(env.step(0))
-    print(env.step(0))