---
type: reference
title: "Pokémon TCG AI Battle — MLX Migration (historical record)"
description: "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."
tags: [pokemon-tcg, mlx, apple-silicon, fp16, gradient-accumulation, training-pipeline, historical]
timestamp: "2026-08-15T16:47:44-03:00"
---

# 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:

- [[pokemon_tcg_training_pipeline]] — end-to-end current training pipeline
- [[pokemon_tcg_train_config_reference]] — live flag reference
- [[pokemon_tcg_kv_cache_hierarchical]] — the row-group cache that landed after this migration
- [[pokemon_tcg_parquet_dataset]] — current dataset format (the migration wrote to Parquet, not `.npy`)
- [[pokemon_tcg_bc_curriculum_ablation]] — validation-at-scale evidence

## 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:

```text
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:

```text
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.

| Component | Required representation |
|---|---|
| IDs, positions, labels | `int32` |
| Masks | `bool` or `uint8` |
| Numeric inputs, embeddings, linears, Q/K/V, residuals | `float16` |
| Logits used by the loss | `float32` |
| Loss reductions | `float32` |
| Accumulated gradients | preferably `float32` |
| Adam state | `float32` |
| Metrics | `float32` 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

| Release | Content | Exit condition |
|---|---|---|
| RC1 | Corrected MLX, FP16 trainer, valid loss/checkpoint | Semantic smoke tests and self-contained checkpoint pass |
| RC2 | Exact compaction, complete logs, autoregressive multi-select | Engine/inference action semantics pass |
| RC3 | Persistent registers, sequential data and TBPTT | Reset/isolation/order tests pass |
| RC4 | Episode deduplication and refined corpus mixture | Episode 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 phase | Landed | Live-state page |
|---|---|---|
| A — Canonical MLX contract | `517ee0f`, `704b127`, `596643a`, `3e4da66`, `9422dad` (2026-07-25) | [[pokemon_tcg_agent_architecture]] |
| B — Semantic P0 fixes | `2daf8be` (B.1–B.3), `0d1e5f4` (B.4–B.6) (2026-07-25) | [[pokemon_tcg_agent_architecture]] + [[pokemon_tcg_training_pipeline]] |
| C — FP16 trainer | `9e92d1d` (2026-07-25) | [[pokemon_tcg_training_pipeline]] |
| D — Compaction + episode metadata + val split | `277a5b5` (2026-07-25), superseded by Parquet on `a942373` (2026-08-03) | [[pokemon_tcg_parquet_dataset]] |
| E — Inference semantics + autoregressive multi-select | `428be76` (2026-07-26), `45a6f43` (2026-07-29) | [[pokemon_tcg_torch_inference]] |
| F — Minimal recurrence + TBPTT | `75ae036` + 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:

- **Parquet corpus and integrated aux heads** — `a942373` (2026-08-03) pivoted the on-disk format and absorbed the aux signal into the primary model. See [[pokemon_tcg_parquet_dataset]] and [[pokemon_tcg_agent_architecture]].
- **Hierarchical KV cache** — `e772fe0` + `fa38caa` (2026-08-06) landed the row-group retention story that made 30k–150k rows/day tractable on 24 GiB unified memory. See [[pokemon_tcg_kv_cache_hierarchical]].
- **Top-Elo curriculum filter** — landed alongside the cache in `e772fe0`; see [[pokemon_tcg_top_elo_curriculum_filter]] for its exact semantics and its evidence-based findings from the ablation.
- **BC curriculum ablation suite + tournament JSON reports + intra-suite round-robin** — `bb72620` (2026-08-06); see [[pokemon_tcg_bc_curriculum_ablation]] and [[pokemon_tcg_tournament_system]]. The round-robin structure is a concrete step toward self-play.
- **Prospective sidecar (rejected)** — the 2026-07-29 sidecar branch was rolled back at `a942373`; see [[pokemon_tcg_prospective_v2]] for the design record and what survived from it.

## Related pages

- [[pokemon_tcg_ai_battle]] — project boundary and evidence model.
- [[pokemon_tcg_agent_architecture]] — current model and its ceiling.
- [[pokemon_tcg_training_pipeline]] — live current pipeline (canonical after this migration).
- [[pokemon_tcg_torch_inference]] — live current arena runtime.
- [[pokemon_tcg_parquet_dataset]] — dataset format that supersedes the D-phase NPY layout.
- [[pokemon_tcg_kv_cache_hierarchical]] — retention infra added after this migration.
- [[pokemon_tcg_bc_curriculum_ablation]] — the historical ablation that provided an architecture-sanity signal within its own cohort; later Stage 3/4 validation defects prevent a universal data-limited claim.
- [[pokemon_tcg_temporal_learning]] — sequence metadata, recurrence and data hygiene.
- [[pokemon_tcg_ladder_and_research]] — Elo-oriented evaluation and deferred research.
- [[pokemon_tcg_repository_timeline]] — commit-level story of this migration and everything that came after.
