Now the one hot encoding of the words is embeded inside the net, is not necesary to one hot encod on the env now

a3c/discrete_A3C.py

@@ -11,29 +11,52 @@ import torch.nn.functional as F
 import gym
 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 = 3000
+MAX_EP = 100000
 
 class Net(nn.Module):
-    def __init__(self, s_dim, a_dim):
+    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
-        self.pi1 = nn.Linear(s_dim, 128)
-        self.pi2 = nn.Linear(128, 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, 128)
-        self.v2 = nn.Linear(128, 1)
-        set_init([self.pi1, self.pi2, self.v1, self.v2])
+        self.v2 = nn.Linear(128, n_emb)
+        self.v3 = nn.Linear(n_emb, 1)
+        set_init([ self.v1, self.v2]) # n_emb
         self.distribution = torch.distributions.Categorical
+        assert a_dim == len(word_list), "putos"
+        word_width = 26 * words_width
+        word_array = np.zeros((len(word_list), word_width))
+        self.actor_head = nn.Linear(n_emb, n_emb)
+        for i, word in enumerate(word_list):
+            for j, c in enumerate(word):
+                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.Tanh(),
+            nn.Linear(128, n_emb),
+        )
 
     def forward(self, x):
-        pi1 = torch.tanh(self.pi1(x))
-        logits = self.pi2(pi1)
+        # 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)
+        logits = torch.log_softmax(
+            torch.tensordot(self.actor_head(values), fw,
+                            dims=((1,), (0,))),
+            dim=-1)
+        values = self.v3(values)
         return logits, values
 
     def choose_action(self, s):
@@ -58,13 +81,14 @@ class Net(nn.Module):
 
 
 class Worker(mp.Process):
-    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name, N_S, N_A):
+    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name, env, N_S, N_A, words_list, word_width):
         super(Worker, self).__init__()
         self.name = 'w%02i' % name
         self.g_ep, self.g_ep_r, self.res_queue = global_ep, global_ep_r, res_queue
         self.gnet, self.opt = gnet, opt
-        self.lnet = Net(N_S, N_A)           # local network
-        self.env = gym.make('WordleEnv100OneAction-v0').unwrapped
+        self.word_list = words_list
+        self.lnet = Net(N_S, N_A, words_list, word_width)           # local network
+        self.env = env.unwrapped
 
     def run(self):
         total_step = 1
@@ -76,8 +100,7 @@ class Worker(mp.Process):
                 if self.name == 'w00':
                     self.env.render()
                 a = self.lnet.choose_action(v_wrap(s[None, :]))
-                s_, r, done, _ = self.env.step(a)
-                if done: r = -1
+                s_, r, done, _ = self.env.step(self.env.encode_word(self.word_list[a]))
                 ep_r += r
                 buffer_a.append(a)
                 buffer_s.append(s)
@@ -89,8 +112,10 @@ class Worker(mp.Process):
                     buffer_s, buffer_a, buffer_r = [], [], []
 
                     if done:  # done and print information
-                        record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name)
+                        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])
                         break
+
                 s = s_
                 total_step += 1
         self.res_queue.put(None)

a3c/utils.py

@@ -47,7 +47,7 @@ def push_and_pull(opt, lnet, gnet, done, s_, bs, ba, br, gamma):
     lnet.load_state_dict(gnet.state_dict())
 
 
-def record(global_ep, global_ep_r, ep_r, res_queue, name):
+def record(global_ep, global_ep_r, ep_r, res_queue, name, goal_word, action):
     with global_ep.get_lock():
         global_ep.value += 1
     with global_ep_r.get_lock():
@@ -56,8 +56,11 @@ def record(global_ep, global_ep_r, ep_r, res_queue, name):
         else:
             global_ep_r.value = global_ep_r.value * 0.99 + ep_r * 0.01
     res_queue.put(global_ep_r.value)
-    print(
-        name,
-        "Ep:", global_ep.value,
-        "| Ep_r: %.0f" % global_ep_r.value,
-    )
+    if goal_word == action:
+        print(
+            name,
+            "Ep:", global_ep.value,
+            "| Ep_r: %.0f" % global_ep_r.value,
+            "| Goal :", goal_word,
+            "| Action: ", action
+        )

main.py

@@ -10,18 +10,20 @@ from wordle_env.wordle import WordleEnvBase
 
 os.environ["OMP_NUM_THREADS"] = "1"
 
-env = gym.make('WordleEnv100OneAction-v0')
+env = gym.make('WordleEnv100FullAction-v0')
 N_S = env.observation_space.shape[0]
 N_A = env.action_space.shape[0]
 
 if __name__ == "__main__":
-    gnet = Net(N_S, N_A)        # global network
+    words_list = env.words
+    word_width = len(env.words[0])
+    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 = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()
 
     # parallel training
-    workers = [Worker(gnet, opt, global_ep, global_ep_r, res_queue, i, N_S = N_S, N_A=N_A) for i in range(mp.cpu_count())]
+    workers = [Worker(gnet, opt, global_ep, global_ep_r, res_queue, i, env, N_S = N_S, N_A=N_A, words_list=words_list, word_width=word_width) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []                    # record episode reward to plot
     while True: