WikifitaGitHub live67e8de5
pesquisa · kaggle/pokemon_tcg_mlx_migration

Pokémon TCG AI Battle — MLX Migration (historical record)

Historical migration contract and phased correction plan for the Mikaelzinho MLX port. Superseded for live-state description by pokemon_tcg_training_pipeline; kept for the record of what the migration promised and delivered.

Baixar raw

Pokémon TCG AI Battle — MLX Migration (historical record)

Status

The migration is complete. The live-state description of the current training pipeline is in pokemon_tcg_training_pipeline. This page is preserved as the record of what the migration contract promised and how it was delivered, phase by phase. Use the current-pipeline pages for anything operational; use this page only for context on how we got here. The migration's FP16 target was a historical contract; the current source snapshot has a strict-FP32 trainer and converter, as recorded in pokemon_tcg_current_state_reconciliation.

Related current-state pages:

Scope

A preserved baseline checkpoint contains 14 epochs, and the user reports an earlier seven-epoch submission at roughly 900–930 ladder rating. Current temporal corrections and the MLX-training/PyTorch-arena split are recorded in pokemon_tcg_training_overhaul_2026_07_29.

The migration target is exactly:

Mikaelzinho PyTorch behavioral-cloning snapshot
        -> Mikaelzinho MLX port

The original author's full PPO/self-play repository is not the parity target. At the time of this migration, MLX was the training runtime and PyTorch FP16 was the arena-inference runtime. That dtype statement is historical; the current strict-FP32 correction is recorded in the current-state pages. The goal is semantic correctness, a reliable training pipeline and better Elo-oriented behavior, not a backend benchmark.

The development target is the M3 Pro with 24 GB unified memory. The M1 Air and transparent multi-Mac memory are outside this phase. The shipped artifact must be self-contained and must not require Drive, network access, APIs or external teachers during inference.

Phase A — canonical contract

Create one versioned architecture/token schema consumed by the encoder, policy, trainer, loader, checkpoint and agent. It must distinguish at least:

CLS, SELECT_TYPE, SELECT_CONTEXT,
SELF_DECK, OPP_DECK,
SELF_PRIZE, OPP_PRIZE,
SELF_HAND, OPP_HAND,
SELF_DISCARD, OPP_DISCARD,
STADIUM,
SELF_ACTIVE, SELF_BENCH,
OPP_ACTIVE, OPP_BENCH,
EFFECT, OPTION, SCRATCH

The MLX port currently contains a semantic collision in the unit token mapping: own bench, opponent active and opponent bench are not all distinct. This is a correctness defect.

The current configuration is frozen at 128/4/3/512/4 with static features and split policy/value heads. A checkpoint must carry and validate architecture version, token-schema version, dimensions, scratch count, option capacity, flags and dtype. The misleading --ff setting must be removed or made consistent with the actual 4 * d_model width.

Phase B — semantic P0 corrections

Attention mask

The MLX port constructs a Boolean padding mask, while the intended MLX attention contract is additive. Available keys receive zero; padded keys receive a dtype-compatible large negative value:

[ M_{i,j}=0\text{ for available keys}, \qquad M_{i,j}\approx-\infty\text{ for padding}. ]

Attention bias

The PyTorch reference uses bias in attention projections. MLX attention must request the matching behavior explicitly rather than inheriting a different default.

Padding ID zero

Card and attack ID zero is padding, including empty entries inside pre-evolution, tool, energy and attack bags. It must contribute a zero vector. The port should mask embedding outputs for ids == 0, reproducing the reference padding_idx=0 semantics.

Static domain tables

card_feat and atom_support are immutable domain data. They must not accidentally become trainable parameters or receive optimizer state. The learned projection may train; the source table must remain fixed.

Categorical value head

If categorical value is enabled, return the scalar expectation:

[ V=\sum_i\operatorname{softmax}(z)_i,s_i, ]

not the atom logits themselves.

Complete inference observation

The live agent currently constructs an observation with logs=[] before updating trackers. The tracker documentation and dataset builder treat logs as incremental deltas carrying reveals, serials, movements, attacks and effects. The live path must pass the complete observation to tracker, ability tracker, encoder and memory.

Phase C — FP16-native trainer (historical contract)

The migration specification treated FP16 as the correct representation contract for this workload, not an optional speed experiment. It described a trainer that read FP16 data and converted numeric arrays back to NumPy FP32 before creating MLX arrays, then required removal of that round trip. This section records the historical target and its rationale; it is not the current dtype authority.

ComponentRequired representation
IDs, positions, labelsint32
Masksbool or uint8
Numeric inputs, embeddings, linears, Q/K/V, residualsfloat16
Logits used by the lossfloat32
Loss reductionsfloat32
Accumulated gradientspreferably float32
Adam statefloat32
Metricsfloat32 or host values

Gradient accumulation is part of the first functional trainer. For K microbatches, accumulate in FP32, normalize by the real example count, clip once after accumulation, perform one optimizer update and advance the scheduler by one optimizer step. The scheduler must count updates, not forwards.

The compiled MLX graph should cover loss, backward, clipping and optimizer update wherever shapes are stable. Gradient norms must not be converted to Python float on every step. mx.eval belongs at explicit state/update and buffer-reuse boundaries.

The current port computes learning-rate steps for one epoch while global_step crosses all epochs. Total steps must include all epochs and accumulation updates. Validation must use real cross-entropy:

[ \mathrm{CE}(z,y)=-(z_y-\operatorname{logsumexp}(z)). ]

Checkpoints must restore model, optimizer, scheduler/global step, architecture configuration, seed and dataset manifest.

Phase D — exact data and shapes

The dataset is columnar and streamed in slabs. The current MLX default of 262,144 rows is too aggressive for the M3 Pro once opt_attr, prefetched slabs, optimizer state, activations and the operating system are included. Start with a manifest-derived conservative policy around 32k–64k rows and prefetch depth one.

Port the exact compaction behavior already available in the PyTorch path:

  • option buckets 32/64/128/192;
  • remove state columns only when padded for every row in the batch;
  • preserve mandatory context, scratch, value and submit positions;
  • remap option source/target references after state compaction;
  • keep a finite set of compiled shapes.

This is exact removal of inaccessible capacity, not an approximation.

Phase E — inference action semantics

Replace the one-pass topk(count) path with autoregressive multi-select. After each selected option:

  1. update picked;
  2. rebuild the relevant option representation and mask;
  3. prevent duplicate selection;
  4. run the next forward pass;
  5. stop only at legal SUBMIT or the legal maximum count.

This aligns inference with the factorization represented by the dataset.

Phase F — minimal recurrence

The corrected baseline reuses the existing scratch-token interface with 16 registers:

[ J_0^{\mathrm{in}}=J_{\mathrm{init}}, \qquad J_{t+1}^{\mathrm{in}}=J_t^{\mathrm{out}}. ]

Expose memory_in and memory_out in the model API. Store memory per match and side, reset it at match start and never share it across sides, matches or processes.

Release candidates

ReleaseContentExit condition
RC1Corrected MLX, FP16 trainer, valid loss/checkpointSemantic smoke tests and self-contained checkpoint pass
RC2Exact compaction, complete logs, autoregressive multi-selectEngine/inference action semantics pass
RC3Persistent registers, sequential data and TBPTTReset/isolation/order tests pass
RC4Episode deduplication and refined corpus mixtureEpisode holdout, rare matchup/action tracking and data validation pass

The release record must include release ID, dataset manifest, architecture config, training state, deck, artifact, observed rating/Elo, matchup results and regressions. Elo and matchup robustness are the objectives; throughput is only a means to train the correct unit.

Functional acceptance

  • No engine files change.
  • Token types are canonical and collision-free.
  • Padding, masks, static tables and categorical value semantics are correct.
  • Losses and gradients are finite.
  • Resume restores training state.
  • Labels are legal under masks.
  • Episodes and temporal order are preserved.
  • Multi-select follows the dataset's sequential semantics.
  • Logs reach inference trackers.
  • Memory resets and isolates correctly.
  • The bundle runs without external files or network.

Delivery vs plan (2026-08-07 audit)

All six migration phases the plan announced landed on develop. Anchor commits and live-state docs:

Plan phaseLandedLive-state page
A — Canonical MLX contract517ee0f, 704b127, 596643a, 3e4da66, 9422dad (2026-07-25)pokemon_tcg_agent_architecture
B — Semantic P0 fixes2daf8be (B.1–B.3), 0d1e5f4 (B.4–B.6) (2026-07-25)pokemon_tcg_agent_architecture + pokemon_tcg_training_pipeline
C — FP16 trainer9e92d1d (2026-07-25)pokemon_tcg_training_pipeline
D — Compaction + episode metadata + val split277a5b5 (2026-07-25), superseded by Parquet on a942373 (2026-08-03)pokemon_tcg_parquet_dataset
E — Inference semantics + autoregressive multi-select428be76 (2026-07-26), 45a6f43 (2026-07-29)pokemon_tcg_torch_inference
F — Minimal recurrence + TBPTT75ae036 + fixes (2026-07-26)pokemon_tcg_tbptt_training_contract

Post-plan work that went beyond the original scope but stayed inside the "correctness before novelty" boundary the plan set:

Related pages