Allow to pass hiperparameters as command line arguments

a3c/discrete_A3C.py

@@ -13,9 +13,7 @@ import torch.multiprocessing as mp
 from .utils import v_wrap, set_init, push_and_pull, record
 import numpy as np
 
-UPDATE_GLOBAL_ITER = 5
 GAMMA = 0.9
-MAX_EP = 500000
 
 class Net(nn.Module):
     def __init__(self, s_dim, a_dim, word_list, words_width):
@@ -25,8 +23,8 @@ class Net(nn.Module):
         n_emb = 32
         # self.pi1 = nn.Linear(s_dim, 128)
         # self.pi2 = nn.Linear(128, a_dim)
-        self.v1 = nn.Linear(s_dim, 128)
-        self.v2 = nn.Linear(128, n_emb)
+        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
@@ -38,9 +36,9 @@ class Net(nn.Module):
                 word_array[i, j*26 + (ord(c) - ord('A'))] = 1
         self.words = torch.Tensor(word_array)
         self.f_word = nn.Sequential(
-            nn.Linear(word_width, 128),
+            nn.Linear(word_width, 64),
             nn.Tanh(),
-            nn.Linear(128, n_emb),
+            nn.Linear(64, n_emb),
         )
 
     def forward(self, x):
@@ -80,8 +78,9 @@ class Net(nn.Module):
 
 
 class Worker(mp.Process):
-    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name, env, N_S, N_A, words_list, word_width, winning_ep):
+    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):
         super(Worker, self).__init__()
+        self.max_ep = max_ep
         self.name = 'w%02i' % name
         self.g_ep, self.g_ep_r, self.res_queue, self.winning_ep = global_ep, global_ep_r, res_queue, winning_ep
         self.gnet, self.opt = gnet, opt
@@ -90,33 +89,26 @@ class Worker(mp.Process):
         self.env = env.unwrapped
 
     def run(self):
-        total_step = 1
-        while self.g_ep.value < MAX_EP:
+        while self.g_ep.value < self.max_ep:
             s = self.env.reset()
             buffer_s, buffer_a, buffer_r = [], [], []
             ep_r = 0.
             while True:
-                if self.name == 'w00':
-                    self.env.render()
                 a = self.lnet.choose_action(v_wrap(s[None, :]))
-                s_, r, done, _ = self.env.step(self.env.encode_word(self.word_list[a]))
+                s_, r, done, _ = self.env.step(a)
                 ep_r += r
                 buffer_a.append(a)
                 buffer_s.append(s)
                 buffer_r.append(r)
 
-                if total_step % UPDATE_GLOBAL_ITER == 0 or done:  # update global and assign to local net
+                if done:  # update global and assign to local net
                     # sync
                     push_and_pull(self.opt, self.lnet, self.gnet, done, s_, buffer_s, buffer_a, buffer_r, GAMMA)
-
-                    if done:  # done and print information
-                        goal_word = self.env.decode_word(self.env.goal_word)
-                        record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name, goal_word, self.word_list[a], len(buffer_a), self.winning_ep)
-                        break
+                    goal_word = self.word_list[self.env.goal_word]
+                    record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name, goal_word, self.word_list[a], len(buffer_a), self.winning_ep)
                     buffer_s, buffer_a, buffer_r = [], [], []
-
+                    break
                 s = s_
-                total_step += 1
         self.res_queue.put(None)
 
 

a3c/utils.py

@@ -58,11 +58,12 @@ def record(global_ep, global_ep_r, ep_r, res_queue, name, goal_word, action, act
     res_queue.put(global_ep_r.value)
     if goal_word == action:
         winning_ep.value += 1
-        print(
-            name,
-            "Ep:", global_ep.value,
-            "| Ep_r: %.0f" % global_ep_r.value,
-            "| Goal :", goal_word,
-            "| Action: ", action,
-            "| Actions: ", action_number
-        )
+        if global_ep.value % 100 == 0:
+            print(
+                name,
+                "Ep:", global_ep.value,
+                "| Ep_r: %.0f" % global_ep_r.value,
+                "| Goal :", goal_word,
+                "| Action: ", action,
+                "| Actions: ", action_number
+            )

main.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import gym
 import matplotlib.pyplot as plt
 import torch.multiprocessing as mp
@@ -7,23 +8,23 @@ from a3c.discrete_A3C import Net, Worker
 from a3c.shared_adam import SharedAdam
 from wordle_env.wordle import WordleEnvBase
 
-
 os.environ["OMP_NUM_THREADS"] = "1"
 
-env = gym.make('WordleEnv100FullAction-v0')
-N_S = env.observation_space.shape[0]
-N_A = env.action_space.shape[0]
-
 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'
+    env = gym.make(env_id)
+    n_s = env.observation_space.shape[0]
+    n_a = env.action_space.n
     words_list = env.words
     word_width = len(env.words[0])
-    gnet = Net(N_S, N_A, words_list, word_width)        # global network
+    gnet = Net(n_s, n_a, words_list, word_width)        # global network
     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)
 
     # parallel training
-    workers = [Worker(gnet, opt, global_ep, global_ep_r, res_queue, i, env, N_S, N_A, words_list, word_width, win_ep) for i in range(mp.cpu_count())]
+    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) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []                    # record episode reward to plot
     while True:

wordle_env/__init__.py

@@ -35,6 +35,12 @@ register(
     max_episode_steps=500,
 )
 
+register(
+    id="WordleEnv100fiftyAction-v0",
+    entry_point=wordle.WordleEnv100fiftyAction,
+    max_episode_steps=500,
+)
+
 register(
     id="WordleEnv100FullAction-v0",
     entry_point=wordle.WordleEnv100FullAction,

wordle_env/wordle.py

@@ -51,7 +51,7 @@ class WordleEnvBase(gym.Env):
             assert len(words) == len(frequencies), f'{len(words), len(frequencies)}'
             self.frequencies = np.array(frequencies, dtype=np.float32) / sum(frequencies)
 
-        self.action_space = spaces.MultiDiscrete(self.words_as_action_space())
+        self.action_space = spaces.Discrete(self.words_as_action_space())
         self.observation_space = spaces.MultiDiscrete(state.get_nvec(self.max_turns))
 
         self.done = True
@@ -70,15 +70,14 @@ class WordleEnvBase(gym.Env):
                 "should always call 'reset()' once you receive 'done = "
                 "True' -- any further steps are undefined behavior."
             )
-        word = self.decode_word(action)
-        goal_word = self.decode_word(self.goal_word)
+        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,
                                         word=word,
                                         goal_word=goal_word)
 
         reward = 0
-        action = tuple(map(tuple, action))
         if action == self.goal_word:
             self.done = True
             #reward = REWARD
@@ -97,20 +96,17 @@ class WordleEnvBase(gym.Env):
         self.state = state.new(self.max_turns)
         self.done = False
         random_word = random.choice(self.words[:self.allowable_words])
-        encoded_random_word = self.encode_word(random_word)
-        self.goal_word = tuple(map(tuple, encoded_random_word))
+        self.goal_word = self.words.index(random_word)
         return self.state.copy()
 
     def set_goal_word(self, goal_word: str):
-        encoded_word = self.encode_word(goal_word.upper())
-        self.goal_word = tuple(map(tuple, encoded_word))
+        self.goal_word = self.words.index(goal_word)
 
     def set_goal_encoded(self, goal_encoded: int):
-        goal_encoded = tuple(map(tuple, goal_encoded))
         self.goal_word = goal_encoded
 
     def words_as_action_space(self):
-        return [[[2] * WORDLE_N] * len(WORDLE_CHARS)] * len(self.words)
+        return len(self.words)
 
     def encode_word(self, word):
         encoded_word = np.array(
@@ -157,6 +153,9 @@ class WordleEnv100TwoAction(WordleEnvBase):
     def __init__(self):
         super().__init__(words=_load_words(100), allowable_words=2)
 
+class WordleEnv100fiftyAction(WordleEnvBase):
+    def __init__(self):
+        super().__init__(words=_load_words(100), allowable_words=50)
 
 class WordleEnv100FullAction(WordleEnvBase):
     def __init__(self):