File size: 3,760 Bytes
9eb2e79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Vectorized cube moves over numpy arrays of facelet indices.
`gen_data.apply_move` rebuilds a 54-character Python string per move, which is
fine when the bottleneck is a Kociemba call. Value iteration has no solver in
the loop: every training step expands all 18 children of every state in the
batch, so at batch 1024 that is ~18k move applications per step and the string
path becomes the whole cost.
A move is a fixed permutation of the 54 facelets -- `new[i] = old[perm[i]]` --
so a batch of states is one fancy-index away: `states[:, perm]`. Quarter turns
compose by indexing a permutation with itself (`p2 = p[p]`), which is how the
half- and counter-turns are built rather than by applying the quarter turn
repeatedly.
States are (N, 54) uint8 of colour indices into COLORS, not strings. Use
`encode`/`decode` at the boundaries only.
"""
import numpy as np
from gen_data import QT, MOVES, SOLVED
COLORS = "URFDLB"
_C2I = {c: i for i, c in enumerate(COLORS)}
# MOVES order is (face, suffix) with suffix in ("", "'", "2") -- 1, 3 and 2
# quarter turns respectively. PERMS is indexed by position in MOVES so a move
# index from the model maps straight to a permutation.
_TURNS = {"": 1, "'": 3, "2": 2}
def _compose(perm: np.ndarray, times: int) -> np.ndarray:
out = np.arange(54, dtype=np.int64)
for _ in range(times):
out = out[perm]
return out
PERMS = np.stack([_compose(np.array(QT[m[0]], dtype=np.int64), _TURNS[m[1:]])
for m in MOVES])
NUM_MOVES = len(MOVES)
assert PERMS.shape == (NUM_MOVES, 54)
def encode(states) -> np.ndarray:
"""Facelet string(s) -> (N, 54) uint8."""
if isinstance(states, str):
states = [states]
return np.array([[_C2I[c] for c in s] for s in states], dtype=np.uint8)
def decode(arr: np.ndarray):
"""(N, 54) uint8 -> list of facelet strings."""
return ["".join(COLORS[i] for i in row) for row in np.asarray(arr)]
SOLVED_ARR = encode(SOLVED)[0]
def apply_moves(states: np.ndarray, move_ids) -> np.ndarray:
"""Applies one move per state. `move_ids` is an int array of length N."""
rows = np.arange(len(states))[:, None]
return states[rows, PERMS[np.asarray(move_ids)]]
def all_children(states: np.ndarray) -> np.ndarray:
"""(N, 54) -> (N, 18, 54): every state's 18 successors, in MOVES order."""
return states[:, PERMS]
def is_solved(states: np.ndarray) -> np.ndarray:
"""(N, 54) -> (N,) bool."""
return (states == SOLVED_ARR).all(axis=1)
def random_walk(n: int, max_depth: int, rng: np.random.Generator):
"""n states from independent walks off solved, depths uniform in 1..max_depth.
Returns (states, depths). The walk never turns the same face twice in a row,
matching gen_data.random_state -- consecutive turns of one face are a single
different move, so allowing them biases the effective depth downward.
Depths are returned for reporting only. They are an upper bound on the true
distance, never the label: value iteration derives its own targets, and a
walk of length 12 routinely lands somewhere far nearer to solved.
"""
depths = rng.integers(1, max_depth + 1, size=n)
states = np.repeat(SOLVED_ARR[None, :], n, axis=0)
prev_face = np.full(n, -1)
for step in range(max_depth):
live = np.where(depths > step)[0]
if not len(live):
break
mv = rng.integers(0, NUM_MOVES, size=len(live))
clash = (mv // 3) == prev_face[live]
while clash.any():
mv[clash] = rng.integers(0, NUM_MOVES, size=int(clash.sum()))
clash = (mv // 3) == prev_face[live]
states[live] = states[live][np.arange(len(live))[:, None], PERMS[mv]]
prev_face[live] = mv // 3
return states, depths
|