Batch JAX differential verification

2026-07-04 19:42:05 +09:00
parent f872204b13
commit 30ccc3cf41
2 changed files with 112 additions and 42 deletions
@@ -129,10 +129,12 @@ def legal_action_mask(state: RefState) -> list[bool]:
    return mask


-def step(state: RefState, action: int) -> tuple[RefState, list[float], bool]:
+def step(
+    state: RefState, action: int, *, validate: bool = True
+) -> tuple[RefState, list[float], bool]:
    if state.done or action < 0 or action >= N_ACTIONS:
        return clone_state(state), [0.0, 0.0], state.done
-    if not legal_action_mask(state)[action]:
+    if validate and not legal_action_mask(state)[action]:
        return clone_state(state), [0.0, 0.0], state.done

    next_state = clone_state(state)