Fix bug in state update

Now only set the character as state "not" if it is not in the word at all

wordle_env/state.py

@@ -77,11 +77,11 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
         offset = 1 + cint * WORDLE_N * 3
         if mask[i] == SOMEWHERE:
             prior_maybe.append(c)
-            # Char at position i = no, and in other positions maybe, other chars stay as they are
-            for j in range(WORDLE_N):
-                # Only flip no if previously was yes or no
-                if sum(state[offset + 3 * j:offset + 3 * j + 3]) == 0:
-                    state[offset + 3 * j:offset + 3 * j + 3] = [0, 1, 0]
+            # 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]
         elif mask[i] == NO:
             # Need to check this first in case there's prior maybe + yes
@@ -168,7 +168,7 @@ def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState,
                         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:
+            elif c not in goal_word:
                 # Char at all positions = no
                 state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
             processed_letters.append(c)