WikifitaGitHub live67e8de5
pesquisa · kaggle/pokemon_tcg_torch_inference

Pokémon TCG AI Battle — PyTorch Inference: FP16 History and FP32 Current Contract

Historical FP16 arena packaging record reconciled with the current strict-FP32 converter, autoregressive multi-select semantics, would-KO integration and per-side tracker state.

Baixar raw

Pokémon TCG AI Battle — PyTorch Inference: FP16 History and FP32 Current Contract

Boundary

This page preserves the FP16 packaging/runtime record from the earlier delivery phase while identifying the current source correction. The current agent/main.py remains the arena entry point, but rl/policy_infer_torch.py now converts and validates floating tensors as strict torch.float32. Complements pokemon_tcg_training_pipeline and pokemon_tcg_agent_architecture.

Current source correction at 20d7d0d

The FP16 statements in the historical sections describe the artifact contract that produced the August 7 submission lineage. They are not the current dtype authority. The live converter uses dtype=torch.float32, rejects non-FP32 floating state, and the current trainer emits FP32 model parameters. Current status and direct source evidence are consolidated in pokemon_tcg_current_state_reconciliation.

Runtime split

Historical delivery -> MLX Metal, FP16 params + FP32 loss/optimizer
Current training    -> MLX Metal, strict FP32 model path
Current packaging   -> MLX checkpoint -> PyTorch FP32 conversion (`rl/policy_infer_torch.py`)
Current arena       -> PyTorch FP32 on CPU (or GPU when available in the sandbox)

The commitment remains MLX for training only. Every current checkpoint must round-trip through the converter to a strict-FP32 PyTorch artifact before tournament evaluation or shipment. The converter preserves static-feature identity, auxiliary-head shapes and split-head topology. The FP16 conversion path remains a historical comparison record.

The submission bundle

uv run tcg-build --checkpoint <mlx.pkl> (or --out <path> to steer the tarball location) produces a self-contained submission archive:

submission.tar.gz
  ├── main.py                                        (from agent/)
  ├── deck.csv                                       (from agent/deck.csv)
  ├── rl/                                            (encoder + inference modules)
  ├── EN_Card_Data.csv                               (static feature source)
  └── model/bc_model/bc_best_torch_fp16.pt          (converted checkpoint)

No external network access. No Google Drive dependency. No teacher service. No mlx wheels — the tarball ships torch-only.

agent/main.py — checkpoint discovery

The module resolves a single checkpoint at import time by walking a priority-ordered list:

_CHECKPOINT_CANDIDATES = [
    os.path.join(_PROJECT_ROOT, "model", "bc_model", "bc_best_torch_fp16.pt"),
    os.path.join(_PROJECT_ROOT, "model", "checkpoint", "bc_best_torch_fp16.pt"),
    os.path.join(_PROJECT_ROOT, "model", "bc_model", "bc_best_mlx_final.pkl"),
    os.path.join(_PROJECT_ROOT, "model", "checkpoint", "bc_best_mlx.pkl"),
    os.path.join(_PROJECT_ROOT, "model", "bc_model", "bc_best_final.pkl"),
    os.path.join(_PROJECT_ROOT, "model", "checkpoint", "bc_best.pkl"),
]

Priority 1 wins in every real deployment: tcg-build writes the converted PyTorch checkpoint to model/bc_model/bc_best_torch_fp16.pt. The MLX-pickle fallbacks 3–6 exist so agent/main.py can still boot on a bare checkpoint if the converter has not run yet — they trigger a converter call at load time.

If no checkpoint is found, the module prints a warning and falls back to a random policy. This is a genuine safety net; in normal operation the checkpoint search always resolves.

Load flow

_LOADED_MODEL, _MODEL_METADATA, _RUNTIME_DATA = _load_model()

def _load_model():
    net, metadata = load_inference_checkpoint(_MODEL_PATH, _CARD_TABLE)
    return net, metadata, ...

# rl/policy_infer_torch.py
def load_inference_checkpoint(path, card_table):
    if path.endswith(".pt"):
        return load_torch_inference_checkpoint(path, card_table)
    elif path.endswith(".pkl"):
        return load_mlx_checkpoint(path, card_table, dtype=torch.float32)

Loaded model is an instance of TokenTransformerTorchInference (subclass of the reference PyTorch TokenTransformer), initialized from arch_config in the checkpoint. The load validates:

  • architecture config matches what the checkpoint claims,
  • static-feature buffer SHA256 matches the CSV in the bundle,
  • aux head shapes match,
  • padding_idx=0 semantics preserved.

Metadata carried out of the load: bc_would_ko, bc_wk_nvar, seed, provenance, mode (baseline / b1 / b2).

Per-side tracker state

Two independent GameTracker instances (one per side) hold state that persists across decisions inside a single episode:

def _get_tracker(side: int):
    st = _TRACKERS.setdefault(side, {
        "tracker": GameTracker(),
        "ability_tracker": AbilityTracker(),
        "logs_seen": 0,
        "would_ko_rng": None,
    })
    ...

The trackers eat every log the engine emits and update:

  • revealed cards / serials,
  • zone movements (deck → hand, hand → field, field → discard, ...),
  • attack and ability activations,
  • effect timers.

Full logs are consumed, not truncated. The old logs=[] shortcut is gone; passing partial logs to choose() would create a train/inference belief-state mismatch (the dataset is built from the same full-log stream).

would_ko_rng is seeded from seed + side at first use for reproducible would-KO determinizations at inference time.

Autoregressive multi-select

Some engine decisions are multi-select: pick 2 attached energies to discard, pick 3 cards to shuffle back. The dataset encodes these as substeps under one decision_id (see pokemon_tcg_parquet_dataset). The trainer learns the substep sequence with substep=0,1,2,... labels.

Inference must reproduce that sequence, not just call topk(count). The algorithm:

def _autoregressive_select_mode(logits_fn, ob, min_count, max_count, memory_in):
    picked_set = []
    for _ in range(max_count):
        logits, memory_out = logits_fn(ob, memory_in=memory_in)
        logits_np = _logits_to_numpy(logits)
        action_mask = ob["action_mask"]  # rebuilt to exclude already-picked
        n = ob["n_options"]
        action = _select_action_from_logits(logits_np, picked_set, action_mask,
                                            n, min_count, results=picked_set)
        if action == SUBMIT_ACTION and len(picked_set) >= min_count:
            break
        picked_set.append(action)
        # rebuild observation with picked_set in the encoded option state
        ob = _rebuild_with_picked(ob, picked_set)
        memory_in = memory_out
    return picked_set

Contract satisfied:

  • One action per substep.
  • action_mask is rebuilt between substeps so already-picked options are excluded.
  • SUBMIT_ACTION only accepted when min_count is met.
  • memory_out from each substep is carried into the next as memory_in, so scratch registers evolve within a multi-select as they did in training.
  • No topk, no batched sampling. One forward per substep.

SUBMIT_ACTION is the reserved option index (MAX_OPTIONS = 192) that means "stop this multi-select and commit the buffered set." The opt_picked column of the observation (already-picked flag) marks which options are in the set for the current substep so the encoder sees the buffered state, not just its cardinality.

Would-KO at inference

When bc_would_ko=True in the loaded checkpoint, the same offline search agent used by the dataset builder computes the per-option would-KO trio (would_ko rate, expected prizes taken, P(ends game)) at inference time and packs it into the option features. This is a pure input feature; no policy head reads or writes it. Enabling it at inference on a model trained with --zero-wouldko (or vice versa) is contractually inconsistent — the loader will refuse.

Would-KO variance count at inference is bc_wk_nvar from the checkpoint metadata.

Forward modes

Three modes, selected by the metadata's mode:

  • baseline — the autoregressive path above. Default.
  • b1 / b2 — experimental variants (_forward_b1, _forward_b2). Not part of the current supported runtime; kept as A/B hooks.

Deck at inference

agent/deck.csv is the deck the submission plays. It is read once at import (load_deck) and cached in the module-global DECK. reload_deck(path) is available specifically for the tournament sweep — see pokemon_tcg_tournament_system — and re-reads the file, verifying that what the agent holds matches what was written.

The submission deck is deliberately immutable per submission (see pokemon_tcg_deck_strategy). Local exploration uses reload_deck, but the tarball ships one canonical deck.csv.

Current conversion boundary

The current converter calls load_mlx_checkpoint(..., dtype=torch.float32) and rejects any floating tensor that is not torch.float32. The legacy output filename bc_best_torch_fp16.pt remains in the checkpoint search list for artifact compatibility, but the current file contents are governed by the strict-FP32 loader, not by that historical name.

Historical FP16 conversion record

Historical path: rl/policy_infer_torch.load_mlx_checkpoint(path, card_table, dtype=torch.float16).

  • Reads the MLX pickle.
  • Reconstructs a TokenTransformerTorchInference from arch_config.
  • Loads the parameter tree by name (flat mapping via _flatten_checkpoint).
  • Casts to torch FP16 (matching MLX FP16 params).
  • Copies the static-card-features buffer directly.
  • Persists to bc_best_torch_fp16.pt via save_torch_inference_checkpoint.

The historical saved .pt file was what agent/main.py preferred at load time. The legacy-named path remains compatible with current discovery, but current conversion uses FP32. Structurally the two backends match at every leaf; numeric parity is validated by scripts/validate/compare_backends.py.

compare_backends.py — parity gate

Independent test: run the same fixed set of real decisions through both backends and count where they disagree:

  • Tied — same top action.
  • NOT tied — both backends had a clear winner (> TIE_TOL) and disagreed.

The NOT tied rate is expected to be very small; a non-zero systemic value points at a conversion drift (dtype casting order, bias placement, mask sign) that the load-time validation missed.

What can go wrong at inference (evidence)

From the historical pokemon_tcg_bc_curriculum_ablation cohort (10 models × ~540 games per per-run tournament × 45 pairs × 30 games in the round-robin, approximately 6,750 games total, none reported as opponent ERR):

  • Zero illegal actions across the whole ablation.
  • Zero missing checkpoint events.
  • Zero would-KO contract mismatch at load.
  • Zero MLX→PyTorch load failures.

No broad inference failure was observed in that historical cohort. This is an inference-path observation, not a claim that the later training objective and validation contract were sound; see pokemon_tcg_stage3_training_failure_postmortem.

Cross-references