Fix bug in state update

Fix tests to include that bug test and refactor step_test

wordle_env/state.py

@@ -161,8 +161,11 @@ def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState,
         offset = 1 + cint * WORDLE_N * 3
         if goal_word[i] != c:
             if c in goal_word and goal_word.count(c) > processed_letters.count(c):
-                # Char at position i = no, and in other positions maybe, other chars stay as they are
-                state[offset:offset + 3 * WORDLE_N] = [0, 1, 0] * WORDLE_N
+                # Char at position i = no, and in other positions maybe except it had a value before, other chars stay as they are
+                for char_idx in range(0, WORDLE_N * 3, 3):
+                    char_offset = offset + char_idx
+                    if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
+                        state[char_offset: char_offset + 3] = [0, 1, 0]
                 state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
                 reward += CHAR_REWARD * 0.1
             else:

wordle_env/test_wordle.py

@@ -7,7 +7,7 @@ TESTWORDS = [
     "APPAA",
     "APPAB",
     "APPAC",
-    "APPAD",
+    "APAPD",
 
     "BPPAB",
     "BPPAC",
@@ -111,6 +111,9 @@ def test_lose_reward(wordleEnv):
     except ValueError:
         pass
 
+def letter_test(char, state, letter_state):
+    offset = 1+3*5*(ord(char)-ord('A'))
+    assert tuple(state[offset:offset+15]) == letter_state
 
 def test_step(wordleEnv):
     wordleEnv.reset(seed=13)
@@ -121,98 +124,98 @@ def test_step(wordleEnv):
     assert state.remaining_steps(cur_state) == wordleEnv.max_turns
     assert state.remaining_steps(new_state) == wordleEnv.max_turns-1
     # Expect B to be all 1,0,0
-    offset = 1+3*5*(ord('B')-ord('A'))
-    assert tuple(new_state[offset:offset+15]) == tuple([1, 0, 0]*5)
-
-    # Expect A to be right in position 0 4 and maybe otherwise
-    offset = 1
-    assert tuple(new_state[offset:offset+15]) == (0, 0, 1,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 0)
-
-    # Expect P to be right in position 2 3 and maybe otherwise
-    offset = 1 + 3*5*(ord('P') - ord('A'))
-    assert tuple(new_state[offset:offset+15]) == (1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 1,
-                                                  1, 0, 0,
-                                                  0, 0, 0)
-
-    # Expect C to be maybes
-    offset = 1 + 3*5*(ord('C') - ord('A'))
-    assert tuple(new_state[offset:offset+15]) == (1, 0, 0,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  0, 0, 0)
+    letter_test('B', new_state, tuple([1, 0, 0]*5))
+
+    # Expect A to be right in position 0 3, no in 1 2 and 0 otherwise
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 0)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2, no in 0 2 and 0 otherwise
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    0, 0, 0)
+    letter_test('P', new_state, letter_state)
+
+    # Expect C to be no everywhere except in 4
+    letter_state = (1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 0)
+    letter_test('C', new_state, letter_state)
     cur_state = wordleEnv.state
-    new_state, reward, done, _ = wordleEnv.step(1)
+    new_state, reward, done, _ = wordleEnv.step(3)
     assert state.remaining_steps(cur_state) == wordleEnv.max_turns-1
     assert state.remaining_steps(new_state) == wordleEnv.max_turns-2
-    # Expect B to be all 1,0,0
-    offset = 1+3*5*(ord('B')-ord('A'))
-    assert tuple(new_state[offset:offset+15]) == tuple([1, 0, 0]*5)
-
-    # Expect A to be right in position 0 4 and maybe otherwise
-    offset = 1
-    assert tuple(new_state[offset:offset+15]) == (0, 0, 1,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 0)
-
-    # Expect P to be right in position 2 3 and maybe otherwise
-    offset = 1 + 3*5*(ord('P') - ord('A'))
-    assert tuple(new_state[offset:offset+15]) == (1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 1,
-                                                  1, 0, 0,
-                                                  0, 0, 0)
+    # Expect D to be all 1,0,0
+    letter_state = tuple([1, 0, 0]*5)
+    letter_test('D', new_state, letter_state)
+
+    # Expect A to be right in position 0 3 and Maybe in 4
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 1, 0)
+    print(new_state, letter_state)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2 no in 3 and maybe in 4
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    0, 1, 0)
+    letter_test('P', new_state, letter_state)
 
     new_state, reward, done, _ = wordleEnv.step(2)
     assert state.remaining_steps(new_state) == wordleEnv.max_turns-3
     # Expect B to be all 1,0,0
-    offset = 1 + 3*5*(ord('B')-ord('A'))
-    assert tuple(new_state[offset:offset+15]) == tuple([1, 0, 0]*5)
+    letter_state = tuple([1, 0, 0]*5)
+    letter_test('C', new_state, letter_state)
 
     # Expect C to be all 1,0,0
-    offset = 1+3*5*(ord('C')-ord('A'))
-    assert tuple(new_state[offset:offset+15]) == tuple([1, 0, 0]*5)
-
-    # Expect A to be right in position 0 4 and maybe otherwise
-    offset = 1
-    assert tuple(new_state[offset:offset+15]) == (0, 0, 1,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 0)
-
-    # Expect P to be right in position 2 3 and maybe otherwise
-    offset = 1 + 3*5*(ord('P') - ord('A'))
-    assert tuple(new_state[offset:offset+15]) == (1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 1,
-                                                  1, 0, 0,
-                                                  0, 0, 0)
+    letter_state = tuple([1, 0, 0]*5)
+    letter_test('C', new_state, letter_state)
+
+    # Expect A to be right in position 0 3 and, no 3 and maybe otherwise
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 1, 0)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2, no in 0 3 and maybe otherwise
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    0, 1, 0)
+    letter_test('P', new_state, letter_state)
 
     new_state, reward, done, _ = wordleEnv.step(0)
-    # Expect A to be right in position 0 4 and 5
-    offset = 1
-    assert tuple(new_state[offset:offset+15]) == (0, 0, 1,
-                                                  1, 0, 0,
-                                                  1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 1)
-
-    # Expect P to be right in position 2 3 and not otherwise
-    offset = 1 + 3*5*(ord('P') - ord('A'))
-    assert tuple(new_state[offset:offset+15]) == (1, 0, 0,
-                                                  0, 0, 1,
-                                                  0, 0, 1,
-                                                  1, 0, 0,
-                                                  1, 0, 0)
+    # Expect A to be right in position 0 3 and 4
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2 and not otherwise
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0)
+    letter_test('P', new_state, letter_state)
     assert state.remaining_steps(new_state) == wordleEnv.max_turns-4
     assert done
     assert wordleEnv.done