Upload folder using huggingface_hub

code/eval_value_search.py

@@ -25,7 +25,12 @@ FIXTURE = Path(__file__).parent / "fixtures" / "canonical_states.json"
 def main():
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--value", required=True,
-                   help="Path to value.pt, or a Hub repo id to download it from.")
+                   help="Path to a value checkpoint, or a Hub repo id to fetch it "
+                        "from. Either flavour works -- supervised (train_value.py) "
+                        "or value iteration (train_value_iteration.py); search.py "
+                        "detects which from the checkpoint.")
+    p.add_argument("--value-file", default="value.pt",
+                   help="Filename within the Hub repo, when --value is a repo id.")
     p.add_argument("--beam", default="64",
                    help="Beam width, or a comma-separated list to sweep. A sweep "
                         "reuses one box for every width, which matters because "
@@ -47,10 +52,12 @@ def main():
     path = args.value
     if not os.path.exists(path):
         from huggingface_hub import hf_hub_download
-        path = hf_hub_download(args.value, "value.pt", token=os.environ.get("HF_TOKEN"))
+        path = hf_hub_download(args.value, args.value_file,
+                               token=os.environ.get("HF_TOKEN"))
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = load_value_model(path, device)
-    print(f"value model on {device} | beam {args.beam} | lam {args.lam}", flush=True)
+    print(f"value model ({model.value_kind}) on {device} | beam {args.beam} | "
+          f"lam {args.lam}", flush=True)
 
     states = json.loads(FIXTURE.read_text())
     if args.depths:

code/search.py

@@ -13,11 +13,17 @@ as the transition function and a network supplying the heuristic. The weighting
 greedy on the heuristic (fast, longer solutions), higher lambda behaves more like
 uniform-cost search.
 
-The heuristic is Kociemba length, an upper bound rather than true optimal
-distance, so it is inadmissible and the search is not guaranteed to return a
-shortest solution. That is fine: the engine verifies that whatever comes back
-actually solves, and finding *a* solution at depth 20 is the open problem, not
-finding the shortest.
+Two kinds of value model can drive this, and `load_value_model` tells them apart
+by what the checkpoint carries:
+
+- **supervised** (`train_value.py`): a distribution over distances, trained on
+  Kociemba lengths. Those are upper bounds rather than true optimal distances, so
+  the heuristic is inadmissible and no shortest-solution guarantee holds.
+- **value iteration** (`train_value_iteration.py`): a scalar, bootstrapped from
+  the solved state with no solver involved.
+
+Either way the engine verifies that whatever comes back actually solves, and
+finding *a* solution at depth 20 is the open problem, not finding the shortest.
 """
 import sys
 from pathlib import Path
@@ -30,29 +36,55 @@ from gen_data import SOLVED, MOVES, apply_move
 
 
 def load_value_model(path, device):
-    from train_value import ValueNet
+    """Loads either value-model flavour, detected from the checkpoint's own keys.
+
+    The two trainers save under different keys ("state_dict" vs "model"), which is
+    the only reliable discriminator: both carry identical hidden/layers/heads, so
+    shape alone cannot distinguish a 27-way distribution head from a scalar one
+    without guessing. The flavour is recorded on the model so `heuristic` reads
+    each correctly -- feeding a scalar head through softmax would silently return
+    a constant and turn the search into an untargeted walk.
+    """
     ckpt = torch.load(path, map_location=device)
-    model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
-    model.load_state_dict(ckpt["state_dict"])
+    if "state_dict" in ckpt:
+        from train_value import ValueNet
+        model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
+        model.load_state_dict(ckpt["state_dict"])
+        kind = "distribution"
+    elif "model" in ckpt:
+        from train_value_iteration import ValueNet as ValueNetScalar
+        model = ValueNetScalar(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
+        model.load_state_dict(ckpt["model"])
+        kind = "scalar"
+    else:
+        raise ValueError(f"{path}: no 'state_dict' or 'model' key; not a value "
+                         f"checkpoint (keys: {sorted(ckpt)[:8]})")
     model.eval()
+    model.value_kind = kind
     return model
 
 
 @torch.no_grad()
 def heuristic(model, states, device, batch_size=1024):
-    """Expected distance-to-solved for each state.
+    """Estimated distance-to-solved for each state.
 
-    Uses the expectation over the predicted distribution rather than the argmax:
-    the extra resolution matters when ranking siblings that all round to the same
-    integer, which is most of the frontier.
+    For a distribution head, the expectation over the predicted distribution
+    rather than the argmax: the extra resolution matters when ranking siblings
+    that all round to the same integer, which is most of the frontier. A scalar
+    head already has that resolution and is read directly.
     """
+    kind = getattr(model, "value_kind", "distribution")
     out = []
     for i in range(0, len(states), batch_size):
         chunk = states[i:i + batch_size]
         ids = torch.tensor([T.encode_state(s) for s in chunk], device=device)
-        probs = torch.softmax(model(ids).float(), dim=-1)
-        values = torch.arange(probs.shape[-1], device=device, dtype=torch.float32)
-        out.extend((probs * values).sum(-1).tolist())
+        raw = model(ids).float()
+        if kind == "scalar":
+            out.extend(raw.clamp(min=0.0).tolist())
+        else:
+            probs = torch.softmax(raw, dim=-1)
+            values = torch.arange(probs.shape[-1], device=device, dtype=torch.float32)
+            out.extend((probs * values).sum(-1).tolist())
     return out
 
 

code/vast_bootstrap.sh

@@ -79,6 +79,7 @@ if [ -n "${EVAL_VALUE_ONLY:-}" ]; then
   cd "${CODE_DIR:-/workspace/code}"
   echo "=== evaluating value search from ${EVAL_VALUE_ONLY} ==="
   /opt/venv/bin/python eval_value_search.py --value "$EVAL_VALUE_ONLY" \
+    --value-file "${VALUE_FILE:-value.pt}" \
     --beam "${SEARCH_BEAM:-64}" --limit-per-depth "${SEARCH_N:-25}" \
     ${SEARCH_DEPTHS:+--depths "$SEARCH_DEPTHS"} ${SEARCH_LAM:+--lam "$SEARCH_LAM"}
   echo "=== done ==="