"""Extrapolate along the update an RL step made: theta = theta_new + beta * (theta_new - theta_old).
The delta is computed in fp32 per tensor; non-weight files are copied from <new>.
usage: extrap.py <new_dir> <old_dir> <beta> <out_dir>"""
import glob, os, shutil, sys, torch
from safetensors.torch import load_file, save_file
new_dir, old_dir, beta, out = sys.argv[1], sys.argv[2], float(sys.argv[3]), sys.argv[4]


def load_all(d):
    sd = {}
    for f in sorted(glob.glob(os.path.join(d, "*.safetensors"))):
        sd.update(load_file(f))
    return sd


new, old = load_all(new_dir), load_all(old_dir)
res, rel = {}, []
for k, v in new.items():
    o = old.get(k)
    if o is None or o.shape != v.shape:
        res[k] = v
        continue
    d = v.float() - o.float()
    res[k] = (v.float() + beta * d).to(v.dtype)
    if d.abs().sum() > 0:
        rel.append(d.norm().item() / o.float().norm().item())
os.makedirs(out, exist_ok=True)
save_file(res, os.path.join(out, "model.safetensors"), metadata={"format": "pt"})
for f in os.listdir(new_dir):
    p = os.path.join(new_dir, f)
    if os.path.isfile(p) and not f.endswith(".safetensors") and f != "model.safetensors.index.json":
        shutil.copy(p, out)
print(f"{len(rel)} tensors changed, mean relative update {sum(rel) / max(1, len(rel)):.2e}, beta {beta} -> {out}")