# Process-supervised PPO stage -- DRAFT config (freezes with the pre-registration).
#
#   consumed by scripts/ppo.py via align.ppo.stage_config (struct mode: an unknown
#   key anywhere raises naming its dotted path). Values are the LEGACY defaults of
#   mistral-prm-ppo; every `[TBD pre-reg freeze]` is a decision NOT made by the port.
#
#   NOT LAUNCHABLE AS SHIPPED, on purpose:
#     * prm.checkpoint does not exist. The legacy PRM is a peft LoRA adapter; this
#       stack reads align.ppo `save_prm_state` files. Re-train the PRM with align
#       or write a converter -- do not point this at the peft directory.
#     * scripts/ppo.py needs the spool rollout worker (own venv) and the 7B base
#       snapshot; neither is on this box as of 2026-10-12.
#     * shapes must satisfy rows % (world * mini_batch_rows) == 0
#       (PPOStageConfig.validate_world): 16 x 4 = 64 rows does NOT split over 3 ranks;
#       e.g. 3 ranks -> 24 prompts x 4 = 96 rows = 3 x 32, mini_batch_rows 8.
sft_checkpoint = "/apps/development/rlhf/nano-mistral-o1/checkpoints/sft/checkpoints/micro_00002250"
run_dir = "/apps/development/rlhf/align/runs/ppo_prm"
boundary_spec_path = "/apps/development/rlhf/align/artifacts/boundary_spec.json"
arm = "prm"  # "orm" = the MANDATORY control (reward table must then carry w_prm = 0)

[lora]
r = 64
alpha = 128.0

[rollout]
prompts_per_iteration = 16
n_samples_per_prompt = 4  # [TBD pre-reg freeze]
max_response_len = 768
temperature = 1.0

[ppo]
ppo_epochs = 2
mini_batch_rows = 8  # per rank
micro_batch_rows = 2
value_warmup_iterations = 3  # < train.iterations
temperature = 1.0  # MUST equal rollout.temperature (validated)
eos_token_id = 2  # validated against config.json AND the tokenizer at launch

[ppo.reward]
transform = "clip_delta"
eta = 0.80  # ABSOLUTE P(step correct) threshold, NOT a quantile  # [TBD pre-reg freeze]
lam_actor = 0.99  # must stay < 1 under delta-form shaping (validated)

[data]
band_file = "/apps/development/rlhf/nano-mistral-o1/data/stageb/bigmath_band0.jsonl"

[train]
iterations = 20  # [TBD pre-reg freeze]
lr = 1e-6  # [TBD pre-reg freeze]; critic lr = lr * value_lr_mult (9)

[canary]
every = 5
n_problems = 100

[prm]
checkpoint = "/apps/development/rlhf/align/artifacts/prm/prm_state.pt"  # DOES NOT EXIST yet