Fix bug on env initialization

a3c/discrete_A3C.py

@@ -13,47 +13,50 @@ import torch.multiprocessing as mp
 from .utils import v_wrap, set_init, push_and_pull, record
 import numpy as np
 
-GAMMA = 0.9
+GAMMA = 0.7
 
 class Net(nn.Module):
     def __init__(self, s_dim, a_dim, word_list, words_width):
         super(Net, self).__init__()
         self.s_dim = s_dim
         self.a_dim = a_dim
-        n_emb = 32
-        # self.pi1 = nn.Linear(s_dim, 128)
-        # self.pi2 = nn.Linear(128, a_dim)
-        self.v1 = nn.Linear(s_dim, 256)
-        self.v2 = nn.Linear(256, n_emb)
-        self.v3 = nn.Linear(n_emb, 1)
-        set_init([ self.v1, self.v2]) # n_emb
-        self.distribution = torch.distributions.Categorical
+        # n_emb = 32
+
         word_width = 26 * words_width
-        word_array = np.zeros((len(word_list), word_width))
-        self.actor_head = nn.Linear(n_emb, n_emb)
+        layers = [
+            nn.Linear(s_dim, word_width),
+            nn.Tanh(),
+            # nn.Linear(128, word_width),
+            # nn.Tanh(),
+            # nn.Linear(256, n_emb),
+            # nn.Tanh(),
+        ]
+        self.v1 = nn.Sequential(*layers)
+        self.v4 = nn.Linear(word_width, 1)
+        self.actor_head = nn.Linear(word_width, word_width)
+
+        self.distribution = torch.distributions.Categorical
+        word_array = np.zeros((word_width, len(word_list)))
         for i, word in enumerate(word_list):
             for j, c in enumerate(word):
-                word_array[i, j*26 + (ord(c) - ord('A'))] = 1
+                word_array[ j*26 + (ord(c) - ord('A')), i ] = 1
         self.words = torch.Tensor(word_array)
-        self.f_word = nn.Sequential(
-            nn.Linear(word_width, 64),
-            nn.Tanh(),
-            nn.Linear(64, n_emb),
-        )
+        # self.f_word = nn.Sequential(
+        #     nn.Linear(word_width, 64),
+        #     nn.ReLU(),
+        #     nn.Linear(64, n_emb),
+        # )
 
     def forward(self, x):
-        # pi1 = torch.tanh(self.pi1(x))
-        fw = self.f_word(
-            self.words.to(x.device.index),
-        ).transpose(0, 1)
-        # logits = self.pi2(pi1)
-        v1 = torch.tanh(self.v1(x))
-        values = self.v2(v1)
+        # fw = self.f_word(
+        #     self.words.to(x.device.index),
+        # ).transpose(0, 1)
+        values = self.v1(x.float())
         logits = torch.log_softmax(
-            torch.tensordot(self.actor_head(values), fw,
+            torch.tensordot(self.actor_head(values), self.words,
                             dims=((1,), (0,))),
             dim=-1)
-        values = self.v3(values)
+        values = self.v4(values)
         return logits, values
 
     def choose_action(self, s):

main.py

@@ -6,10 +6,44 @@ import torch.multiprocessing as mp
 
 from a3c.discrete_A3C import Net, Worker
 from a3c.shared_adam import SharedAdam
+from a3c.utils import v_wrap
 from wordle_env.wordle import WordleEnvBase
 
 os.environ["OMP_NUM_THREADS"] = "1"
 
+def evaluate(net, env):
+    print("Evaluation mode")
+    n_wins = 0
+    n_guesses = 0
+    n_win_guesses = 0
+    env = env.unwrapped
+    N = env.allowable_words
+    for goal_word in env.words[:N]:
+        win, outcomes = play(net, env)
+        if win:
+            n_wins += 1
+            n_win_guesses += len(outcomes)
+        else:
+            print("Lost!", goal_word, outcomes)
+        n_guesses += len(outcomes)
+
+    print(f"Evaluation complete, won {n_wins/N*100}% and took {n_win_guesses/n_wins} guesses per win, "
+          f"{n_guesses / N} including losses.")
+
+def play(net, env):
+    state = env.reset()
+    outcomes = []
+    win = False
+    for i in range(env.max_turns):
+        action = net.choose_action(v_wrap(state[None, :]))
+        state, reward, done, _ = env.step(action)
+        outcomes.append((env.words[action], reward))
+        if done:
+            if reward >= 0:
+                win = True
+            break
+    return win, outcomes
+
 if __name__ == "__main__":
     max_ep = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
     env_id = sys.argv[2] if len(sys.argv) > 2 else 'WordleEnv100FullAction-v0'
@@ -39,4 +73,5 @@ if __name__ == "__main__":
     plt.plot(res)
     plt.ylabel('Moving average ep reward')
     plt.xlabel('Step')
-    plt.show()
+    plt.show()
+    evaluate(gnet, env)

wordle_env/state.py

@@ -141,7 +141,7 @@ def update_mask(state: WordleState, word: str, goal_word: str) -> WordleState:
 
 def update(state: WordleState, word: str, goal_word: str) -> WordleState:
     state = state.copy()
-
+    reward = 0
     state[0] -= 1
     processed_letters = []
     for i, c in enumerate(word):
@@ -149,6 +149,8 @@ def update(state: WordleState, word: str, goal_word: str) -> WordleState:
         offset = 1 + cint * WORDLE_N * 3
         if goal_word[i] == c:
             # char at position i = yes, all other chars at position i == no
+            if state[offset + 3 * i:offset + 3 * i + 3][2] == 0:
+                reward += 0.1
             state[offset + 3 * i:offset + 3 * i + 3] = [0, 0, 1]
             for ocint in range(len(WORDLE_CHARS)):
                 if ocint != cint:
@@ -168,5 +170,5 @@ def update(state: WordleState, word: str, goal_word: str) -> WordleState:
                 # Char at all positions = no
                 state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
             processed_letters.append(c)
-    return state
+    return state, reward
 

wordle_env/wordle.py

@@ -73,11 +73,11 @@ 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 = self.state_updater(state=self.state,
+        self.state, r = self.state_updater(state=self.state,
                                         word=word,
                                         goal_word=goal_word)
 
-        reward = 0
+        reward = r
         if action == self.goal_word:
             self.done = True
             #reward = REWARD
@@ -159,7 +159,7 @@ class WordleEnv100fiftyAction(WordleEnvBase):
 
 class WordleEnv100FullAction(WordleEnvBase):
     def __init__(self):
-        super().__init__(words=_load_words(), allowable_words=100)
+        super().__init__(words=_load_words(100), allowable_words=100)
 
 
 class WordleEnv1000(WordleEnvBase):
@@ -175,7 +175,7 @@ class WordleEnv1000WithMask(WordleEnvBase):
 
 class WordleEnv1000FullAction(WordleEnvBase):
     def __init__(self):
-        super().__init__(words=_load_words(), allowable_words=1000)
+        super().__init__(words=_load_words(1000), allowable_words=1000)
 
 
 class WordleEnvFull(WordleEnvBase):