| """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)} |
|
|
| |
| |
| |
| _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 |
|
|
