Policy・Value学習モジュール

対象: pca.training.train, pca.training.policy_value

役割

self-play JSONL から Policy/Value model を学習する。legacy ActionConditionedPolicyValueNet と v13 UnifiedTokenPolicyValueNet の両方を扱い、aux prize heads と integrated belief heads の loss もここで合成する。

モジュール一覧

モジュール	役割	実装の要点
pca.training.train	CLI / facade	train_policy_value() と CLI。policy_value helper を旧 import path 互換で re-export する。
pca.training.champion_adoption	Champion管理	外部学習checkpointの昇格artifactを検証し、alias、ledger、replay registryを同期する。
pca.training.policy_value.config	config	TrainConfig と SearchDataSummary。
pca.training.policy_value.data	dataloader	search record prescan、streaming dataloader、data summary。
pca.training.policy_value.losses	loss	policy/value/turn/aux prize/integrated belief loss と metrics。
pca.training.policy_value.gradient_diagnostics	gradient診断	shared Transformerに対するloss別gradient norm、比率、cosineを集計する。
pca.training.policy_value.metadata	metadata tables	JSONL embedded card/attack metadata から unified feature tables と cache を作る。
pca.training.policy_value.model_config	model config	checkpoint/input/CLI から model config を解決し、互換 state dict を load する。
pca.training.policy_value.checkpointing	checkpoint payload	model state、optimizer-independent metadata、training config を保存形式にまとめる。
pca.training.policy_value.runtime	runtime helpers	device selection、AMP、grad scaler、memory log、LR schedule、progress line。
pca.training.policy_value.training_state	batch resume	model/optimizer/RNG/batch cursorを原子的に保存・復元する。
pca.training.value_drift_analysis	value drift診断	baseline/candidateのbackboneとvalue headを交差適用し、value悪化の発生箇所を分解する。
pca.training.telemetry	live metrics	append-only JSONLと既存MLflow runへの失敗許容metric出力。
pca.training.tracing	validation trace	epoch validationの入力record数、metric、所要時間、成功/失敗をMLflow Traceへ記録する。

公開API

API	用途
TrainConfig	training loop 設定。
train_policy_value(input_path, output_path, config, best_output_path)	学習実行。
policy_value_loss(output, batch, ...)	model output と batch から loss/metrics を作る。
resolve_model_config(...)	model construction config を決める。
load_compatible_model_state(...)	checkpoint 互換 load。
load_unified_feature_tables_from_jsonl(...)	JSONL metadata から static feature tables を作る。
policy_value_checkpoint_payload(...)	checkpoint 保存 payload。
adopt_promoted_checkpoint(...)	評価済みcheckpointを運用Championへ反映する。

CLIの使い方

PYTHONPATH=src uv run python -m pca.training.train \
  --input data/selfplay/example.jsonl \
  --output checkpoints/policy_value.pt \
  --best-output checkpoints/policy_value_best.pt \
  --metrics-output logs/train-metrics.json \
  --validation-input data/selfplay/example.validation.jsonl \
  --epochs 3 \
  --batch-size 16 \
  --gradient-accumulation-steps 4 \
  --gradient-diagnostics-every-steps 250 \
  --checkpoint-every-batches 3000 \
  --restart-every-batches 3000 \
  --mps-restart-driver-memory-mb 22000 \
  --batch-padding-mode mps_stable \
  --resume \
  --model-class unified

value lossが学習後に悪化した場合は、固定validationに対してbackboneとvalue headを交差適用する。

pca value-drift \
  --baseline-checkpoint checkpoints/champion.pt \
  --candidate-checkpoint checkpoints/candidate.pt \
  --validation-input data/selfplay/run.validation.jsonl \
  --output logs/value-drift/run.json

train-compareなどself-play pipeline外で作ったcheckpointは、promotion-runに合格しただけではlatest aliasへ反映されない。評価artifactとcheckpoint hash、現在Championを再検証してから次のコマンドで正式採用する。

pca promote-checkpoint \
  --run-name v14-gumbel-iterative-10k \
  --run-id 20260714-005806-cycle003 \
  --candidate-checkpoint checkpoints/experiments/example/attempt-best.pt \
  --final-checkpoint checkpoints/experiments/example/attempt-final.pt \
  --promotion-result data/eval/promotion-gate/example/promotion-benchmark.json \
  --data data/selfplay/example-cycle003.train.jsonl \
  --training-input \
  data/selfplay/example-cycle003.train.jsonl \
  data/replay/example/cycle003-history.jsonl \
  --training-metrics logs/optimizer-comparison/example/attempt-metrics.json \
  --attempt-label attempt-m

更新対象はofficial / trainingのglobal・run-name alias、latest manifest、cycle manifest、run ledger、replay registryである。benchmarkがpromotedでない、Candidate hashが異なる、または現在latestがbenchmark時のChampion / 同じCandidate以外へ進んでいる場合は何も更新せず失敗する。実更新前の検証には--dry-runを使う。

baseline_backbone_candidate_headはvalue headだけを更新した場合、 candidate_backbone_baseline_headは共有backboneだけを更新した場合に相当する。通常のbaseline/candidateと同一batchでMSE、MAE、平均bias、最終勝敗ラベル別MSEを比較する。targetはtrain configと同じ (1 - turn_value_weight) * final_result + turn_value_weight * search_turn_valueを使う。

注意点

• TrainConfig.streaming=True の場合、大きい JSONL を memory に載せずに学習する。
• V14の標準train profileはCycle 3のmulti-head balance比較とpromotion benchmarkで採用したAttempt Mを使う。epochs=1、batch-size=16、gradient-accumulation-steps=4、 learning-rate=1e-5、gradient-clip-norm=10.0、value-loss-weight=0.5、 aux-prize-loss-scale=1.25、belief-backbone-gradient-scale=1.0である。これはV14の現行baselineであり、各cycleのpromotion gateを通じて継続的に非回帰を確認する。
• --gradient-accumulation-stepsは複数micro-batchのlossを正規化してから1回のoptimizer更新へまとめる。V14の256次元profileはbatch-size=16, gradient-accumulation-steps=4でeffective batch 64を維持する。
• --gradient-diagnostics-every-steps Nは、step 1と以後N optimizer stepごとに、実際に重み付けされたPolicy、Value、Auxiliary、Belief lossがshared Transformerへ与えるgradientを測る。gradient accumulationを使う場合は、同じoptimizer更新に含まれる全micro-batchのgradientを合算してからnormとcosineを計算する。診断用のautograd.gradは.gradを変更せず、通常の合成loss、backward、clip、optimizer更新は従来どおり実行する。
• --checkpoint-every-batchesを有効にすると<output>.train-state.ptへmodel、optimizer、RNG、epoch、処理済みmicro-batch数、途中metricを保存する。streaming stateには入力byte位置、shuffle RNG、buffer内recordのfile referenceも含めるため、新形式stateからのresumeは先頭recordを再decodeしない。正常完了時はstateを小さい完了markerへ置き換える。
• --restart-every-batchesはMPSだけで有効になり、state保存後に同じcommandを--resume付きでexecする。MPS/MPSGraph contextをprocess単位で解放しながらoptimizerとデータ順を継続する。CPU/CUDAでは無視する。
• --mps-restart-driver-memory-mbはoptimizer更新境界でdriver_allocated_memory()を確認し、固定batch間隔より先に危険域へ達した場合もstateを保存してprocessを置き換える。V14は22,000 MiBを使う。
• --batch-padding-mode mps_stableはstate、history、action、structured objectをV14で通常観測される範囲まで固定paddingする。値は上限ではないため、大きいrecordを切り捨てない。これにより可変長batchごとに生成されるMPSGraph input shapeの種類を減らす。
• MPSではdriver cacheの増加を抑えるため--empty-cache-every-batchesを使う。watermark無効化はsystem memoryを使い切る可能性があるため標準運用では行わない。
• --validation-input を指定するとoptimizer更新前のcheckpointを同じholdoutで一度評価し、 pretrain_val_*として保存する。各epoch末にも評価し、--best-output は val_loss 最小のepochを保存する。
• --metrics-output は学習前baseline、final/bestのloss内訳、train/validation対象record数、validation - trainのloss gap、学習前からのvalidation_*_improvementをJSONに保存する。gradient診断を有効にした場合は、stepごとの結果をgradient_diagnostics.samplesにも保存する。
• pca selfplay-train経由では表示間隔ごとにtrain/windowとtrain/runningを分けて記録する。policy lossに加えてtarget/predicted entropy、KL、top-1一致率、LR、gradient norm、memoryをMLflowで確認できる。gradientは互換名grad_normに加え、grad_norm_pre_clip、 grad_norm_post_clip、grad_clip_scale、grad_clip_appliedを記録する。window平均の grad_clip_appliedはclip適用率になる。
• loss別gradient診断はMLflowのtrain/gradient_diagnosticsに記録する。*_grad_normは各lossの更新強度、*_to_policy_norm_ratioはPolicyに対する相対強度、*_cosineは更新方向の一致を示す。cosineが負なら競合、0付近なら独立、正なら同じ方向への更新を意味する。診断は追加のbackward相当処理を伴うため、通常は250step程度の間隔で使う。
• V14 train profileはgradient-diagnostics-every-steps=250を既定値にする。1 epoch約2,265 optimizer stepのCycle 3相当データでは、step 1を含めて約10sampleを記録する。
• epoch終了時はtrain/epochにvalidationとvalidation_*_gapを含むmetricを記録し、validation区間を train.validation Traceにも保存する。学習前評価はtrain/baseline metricと train.pretrain_validation Traceに保存する。telemetry/tracing障害はtrainingを停止させない。
• policy_target_source は search / oracle などを切り替える。
• aux prize / integrated belief は config の profile と weight が 0 でないと loss に入らない。
• aux-prize-loss-scaleは3つのAuxiliary prize lossの内訳を保ったまま、total lossへの寄与だけを一括調整する。
• belief-backbone-gradient-scaleはBelief headのforward値とhead parameter勾配を保ち、shared backboneへ戻るBelief勾配だけを0..1で縮小する。