Check actual PufferLib 4 CPU training and saved-weight reload

.github/workflows/install.yml

@@ -61,3 +61,36 @@ jobs:
           tests/test_diplomacy.py
           tests/test_agent_worktree.py
           -q
+
+  pufferlib4-cpu-training:
+    name: pufferlib4-cpu-training
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - name: Set up Python
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: "3.12"
+
+      - name: Check out pinned PufferLib 4 source
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          repository: PufferAI/PufferLib
+          ref: 42f70d6932c30ac977736f861006809c50168ba9
+          path: build/pufferlib4-upstream
+          persist-credentials: false
+
+      - name: Install CPU training dependencies
+        run: |
+          sudo apt-get update && sudo apt-get install -y clang libomp-dev
+          python -m pip install torch==2.9.0 --index-url https://download.pytorch.org/whl/cpu
+          python -m pip install -r requirements/pufferlib4-train.txt
+
+      - name: Build PufferLib 4 CPU environment backend
+        run: python scripts/build_diplomacy_puffer4.py --upstream build/pufferlib4-upstream
+
+      - name: Check actual PufferLib 4 CPU training and weight reload
+        run: python scripts/smoke_pufferlib4_training.py --upstream build/pufferlib4-upstream

docs/development/pufferlib-4-migration.md

@@ -87,9 +87,9 @@ tests on 3.9 are not evidence that the full trainer supports 3.9.
    The tests cover positive deltas, no repeat reward, reset bookkeeping and the
    adapter retaining final rewards through autoreset. Review and full checks
    remain required; this fix alone does not establish a usable training game.
-3. **Real 4.0 training.** Isolate the pinned Python package and native module;
-   run a small CPU PyTorch training job on Python >=3.10. Verify nonzero reward
-   input, finite losses, policy updates, and checkpoint load behavior. Record
+3. **Real 4.0 training (issue #13).** The isolated smoke run below checks the
+   actual upstream CPU PyTorch trainer on Python >=3.10. It checks nonzero reward
+   input, finite losses, policy updates, and saved-weight reload. Record
    differences in clipping, recurrence, truncation and training configuration.
    Upstream CPU `create_vec` leaks temporary configuration dictionaries and
    leaves lifecycle management to the caller; inspect/fix these before long runs.
@@ -128,3 +128,56 @@ declarations. The build now requests POSIX 2008 declarations explicitly;
 implicit-declaration checking stays enabled. The fixed adapter passed Linux CI
 run 38182492879 (247 passed, 46 skipped). The Winter reward fix passed run
 38182603605 (252 passed, 46 skipped).
+
+## CPU training check
+
+Use a separate virtual environment with Python >=3.10. This script imports the
+pinned upstream Python files, verifies their bytes, and loads the native adapter
+under that package. It does not overwrite either checkout's Python package.
+Use a separate process so old and new packages cannot accidentally share imports.
+
+```bash
+python -m pip install -r requirements/pufferlib4-train.txt
+# macOS: use --torch-openmp to share Torch's bundled OpenMP runtime.
+python scripts/build_diplomacy_puffer4.py --upstream /path/to/PufferLib-4-source --torch-openmp
+python scripts/smoke_pufferlib4_training.py --upstream /path/to/PufferLib-4-source
+```
+
+On Linux omit `--torch-openmp`; CI installs a CPU Torch wheel and runs this as a
+separate job. The flag only changes the newly built adapter's library reference;
+it does not modify Torch, Homebrew, or upstream files. On macOS it uses Torch's
+headers/library, rewrites its retained Homebrew install name to the Torch rpath,
+and ad-hoc signs the generated library. This avoids loading two copies of libomp.
+The first training run aborted with OpenMP error 15; no unsafe duplicate-runtime
+override was used. The rebuilt run completed. Torch emits an upstream TF32 API
+deprecation warning; it has not been suppressed.
+
+The smoke policy uses upstream's MLP and trainer, with the same action mask in
+rollout and update: HOLD everywhere, or disband unit zero during Winter. It
+starts with a disband-logit bias so a short check actually sees positive rewards.
+The script verifies this restriction from recorded observations/actions. This
+avoids relying on the existing illegal non-Winter voluntary-disband behavior.
+It is a restricted integration test, not a learned or strategic baseline, and
+does not demonstrate full order legality or action-mask support in the GPU backend.
+
+On macOS arm64, Python 3.12.14 and Torch 2.9.0, a 336-step run received positive
+welfare rewards, had finite losses, changed eight parameter tensors, and restored
+identical logits/value predictions after saving, perturbing and reloading weights.
+The script prints exact configuration and measurements. Independent review
+reproduced the run and full suite (252 passed, 46 skipped); Linux CI run
+38183006374 passed both jobs at reviewed commit 3add2322. New heads still require
+CI/review re-checks after changes.
+
+Important limits: upstream clips rewards to [-1, 1], so this is not yet a faithful
+optimizer of cumulative welfare. The script reports how many reward entries that
+would change. Independent review also reproduced a dropped-batch-boundary bug:
+upstream starts each rollout with zero rewards/terminals. Before the second batch,
+22 pending reward points and 14 terminal flags were discarded; another 24 reward
+points and 14 flags remained after the last batch. The reported 103 is stored
+rollout rewards, not all environment-step rewards. The script now labels it
+accordingly and reports each boundary and the final pending values. Fixing this
+trainer defect is a separate task, not a result established by a passing smoke run.
+MinGRU episode resets, long-run binding leaks, GPU behavior, historical opponents,
+and full optimizer/RNG
+checkpoint resume still need focused tests. Do not turn this smoke script into
+the default training entry point without resolving those issues.

requirements/pufferlib4-train.txt

@@ -0,0 +1,6 @@
+# Minimal CPU smoke dependencies; the full upstream package has more extras.
+numpy==1.26.4
+torch==2.9.0
+rich==15.0.0
+rich_argparse==1.8.0
+pybind11==3.0.1

scripts/build_diplomacy_puffer4.py

@@ -27,6 +27,8 @@ def run(command):
 def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--upstream", required=True, type=Path)
+    parser.add_argument("--torch-openmp", action="store_true",
+                        help="on macOS, share the OpenMP runtime shipped with Torch")
     args = parser.parse_args()
     upstream = args.upstream.resolve()
     for name in SOURCES:
@@ -39,13 +41,21 @@ def main():
     includes = [ROOT / "native/pufferlib4", ROOT / "pufferlib/ocean/diplomacy", upstream / "src"]
     common = ["-O2", "-fPIC", "-fvisibility=hidden"]
     omp_link = ["-fopenmp"]
+    torch_omp = None
     if sys.platform == "darwin":
-        omp = Path(os.environ.get("PUFFERWELFARE_LIBOMP", "/opt/homebrew/opt/libomp"))
+        if args.torch_openmp:
+            import torch
+            omp = Path(torch.__file__).resolve().parent
+            torch_omp = omp / "lib/libomp.dylib"
+        else:
+            omp = Path(os.environ.get("PUFFERWELFARE_LIBOMP", "/opt/homebrew/opt/libomp"))
         if not (omp / "include/omp.h").is_file():
             parser.error("libomp is required; set PUFFERWELFARE_LIBOMP to its prefix")
         common += ["-Xpreprocessor", "-fopenmp", f"-I{omp / 'include'}"]
         omp_link = [f"-L{omp / 'lib'}", f"-Wl,-rpath,{omp / 'lib'}", "-lomp"]
     else:
+        if args.torch_openmp:
+            parser.error("--torch-openmp is only needed/supported on macOS")
         # Strict C11 otherwise hides clock_gettime/CLOCK_MONOTONIC on Linux.
         common += ["-fopenmp", "-D_POSIX_C_SOURCE=200809L"]
     cc = shlex.split(os.environ.get("CC", "clang"))
@@ -68,6 +78,12 @@ def main():
     run(cxx + common + ["-std=c++17", "-shared", "-DENV_NAME=diplomacy"]
         + [f"-I{path}" for path in includes + [Path(pybind11.get_include()), Path(sysconfig.get_path("include"))]]
         + [str(upstream / "src/bindings_cpu.cpp")] + objects + omp_link + link + ["-o", str(extension)])
+    if torch_omp is not None:
+        # Torch's bundled library retains a Homebrew install name. Resolve it
+        # through the Torch rpath instead of loading a second Homebrew runtime.
+        install_name = subprocess.check_output(["otool", "-D", str(torch_omp)], text=True).splitlines()[1]
+        run(["install_name_tool", "-change", install_name, "@rpath/libomp.dylib", str(extension)])
+        run(["codesign", "--force", "--sign", "-", str(extension)])
     print(f"Built CPU environment backend: {extension}")
 
 

scripts/smoke_pufferlib4_training.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""Check the real pinned PufferLib 4 CPU trainer, not learning performance.
+
+Run as a separate process: the old package and upstream package share a name.
+The upstream trainer clips rewards and does not mask our incomplete actions.
+Those limitations must be resolved before claiming a useful welfare baseline.
+"""
+import argparse
+import ast
+import configparser
+import importlib.util
+import json
+from pathlib import Path
+import subprocess
+import sys
+import sysconfig
+import tempfile
+
+PIN = "42f70d6932c30ac977736f861006809c50168ba9"
+ROOT = Path(__file__).resolve().parents[1]
+FILES = ["pufferlib/__init__.py", "pufferlib/pufferl.py", "pufferlib/selfplay.py",
+         "pufferlib/torch_pufferl.py", "pufferlib/models.py", "pufferlib/muon.py",
+         "config/default.ini"]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--upstream", type=Path, required=True)
+    cli = parser.parse_args()
+    if sys.version_info < (3, 10):
+        parser.error("the upstream Python trainer requires Python >=3.10")
+    upstream = cli.upstream.resolve()
+    for name in FILES:
+        pinned = subprocess.check_output(["git", "-C", str(upstream), "show", f"{PIN}:{name}"])
+        if (upstream / name).read_bytes() != pinned:
+            parser.error(f"{name} differs from the pinned upstream source")
+    sys.path.insert(0, str(upstream))
+    import pufferlib
+    if Path(pufferlib.__file__).resolve() != upstream / "pufferlib/__init__.py":
+        raise RuntimeError("imported the old package instead of pinned PufferLib 4")
+    extension = ROOT / "build/pufferlib4" / ("_C" + sysconfig.get_config_var("EXT_SUFFIX"))
+    spec = importlib.util.spec_from_file_location("pufferlib._C", extension)
+    backend = importlib.util.module_from_spec(spec)
+    sys.modules["pufferlib._C"] = backend
+    spec.loader.exec_module(backend)
+    pufferlib._C = backend
+    import torch
+    from pufferlib.torch_pufferl import PuffeRL, load_policy
+
+    class WinterOnlySmokePolicy(torch.nn.Module):
+        """A deliberately small policy: HOLD, or unit-zero DISBAND in Winter.
+
+        Mask both evaluation and training identically. Do not exploit the known
+        engine bug allowing voluntary disbands during movement/retreat phases.
+        This exercises the upstream Torch policy hook, not a new trainer.
+        """
+        def __init__(self, base):
+            super().__init__()
+            self.base = base
+
+        def initial_state(self, *args, **kwargs):
+            return self.base.initial_state(*args, **kwargs)
+
+        def mask(self, observations, logits):
+            flat = observations.reshape(-1, observations.shape[-1])
+            allowed = torch.zeros_like(logits, dtype=torch.bool)
+            allowed[:, 0] = True
+            allowed[:, 61] = flat[:, 1624] == 1  # phase one-hot: Winter
+            return logits.masked_fill(~allowed, torch.finfo(logits.dtype).min)
+
+        def forward_eval(self, observations, state):
+            logits, value, state = self.base.forward_eval(observations, state)
+            return self.mask(observations, logits), value, state
+
+        def forward(self, observations):
+            logits, value = self.base(observations)
+            return self.mask(observations, logits), value
+
+    torch.set_num_threads(1)
+    torch.manual_seed(42)
+    cfg = configparser.ConfigParser()
+    cfg.read(upstream / "config/default.ini")
+    args = {}
+    for section in cfg.sections():
+        values = {}
+        for key, value in cfg[section].items():
+            try:
+                values[key] = ast.literal_eval(value)
+            except (SyntaxError, ValueError):
+                values[key] = value
+        if section == "base":
+            args.update(values)
+        else:
+            args[section] = values
+    args.update(env_name="diplomacy", wandb=False, load_id=None, load_model_path=None)
+    args["vec"].update(total_agents=14, num_buffers=1, num_threads=1)
+    args["env"].update(max_years=2, welfare_mode=1)
+    args["policy"].update(hidden_size=32, num_layers=1)
+    args["torch"]["network"] = "MLP"
+    args["train"].update(horizon=12, minibatch_size=168, total_timesteps=336,
+                         replay_ratio=1.0, learning_rate=0.001)
+    vec = backend.create_vec(args, 0)
+    trainer = None
+    try:
+        policy = WinterOnlySmokePolicy(load_policy(args, vec))
+        # Make rare disband actions observable in a tiny smoke run. This is an
+        # explicit initialization bias, not a trained or benchmark policy.
+        with torch.no_grad():
+            policy.base.decoder.decoder.bias[61] = 2
+        trainer = PuffeRL(args, vec, policy, verbose=False)
+        initial = {name: value.detach().clone() for name, value in policy.named_parameters()}
+        reward_total = 0.0
+        clipped_reward_count = 0
+        batch_boundaries = []
+        epochs = []
+        for _ in range(2):
+            pending_rewards = trainer.vec_rewards.clone()
+            pending_terminals = trainer.vec_terminals.clone()
+            trainer.rollouts()
+            batch_boundaries.append({
+                "pending_reward_sum": pending_rewards.sum().item(),
+                "stored_first_row_reward_sum": trainer.rewards[0].sum().item(),
+                "pending_terminal_count": pending_terminals.sum().item(),
+                "stored_first_row_terminal_count": trainer.terminals[0].sum().item(),
+            })
+            is_winter = trainer.observations[:, :, 1624] == 1
+            rollout_actions = trainer.actions[:, :, 0]
+            if not torch.all((rollout_actions == 0) | ((rollout_actions == 61) & is_winter)):
+                raise RuntimeError("smoke policy used an action outside its stated restriction")
+            reward_total += trainer.rewards.sum().item()
+            clipped_reward_count += (trainer.rewards.abs() > 1).sum().item()
+            trainer.train()
+            if not all(torch.isfinite(torch.tensor(value)).item() for value in trainer.losses.values()):
+                raise RuntimeError(f"non-finite losses: {trainer.losses}")
+            epochs.append(dict(trainer.losses))
+        if reward_total <= 0:
+            raise RuntimeError("the real trainer received no positive welfare reward")
+        changed = sum(not torch.equal(initial[name], value) for name, value in policy.named_parameters())
+        if changed == 0:
+            raise RuntimeError("training did not change policy parameters")
+        obs = trainer.vec_obs.clone()
+        state = policy.initial_state(vec.total_agents, device="cpu")
+        with torch.no_grad():
+            expected = policy.forward_eval(obs, state)[:2]
+        with tempfile.TemporaryDirectory(prefix="pufferlib4-checkpoint-") as scratch:
+            path = str(Path(scratch) / "weights.pt")
+            trainer.save_weights(path)
+            with torch.no_grad():
+                for parameter in policy.parameters():
+                    parameter.add_(1)
+            trainer.load_weights(path)
+            with torch.no_grad():
+                actual = policy.forward_eval(obs, state)[:2]
+            for before, after in zip(expected, actual):
+                torch.testing.assert_close(before, after, rtol=0, atol=0)
+        print(json.dumps({
+            "upstream_commit": PIN, "python": sys.version.split()[0],
+            "torch": torch.__version__, "network": "MLP", "device": trainer.device,
+            "steps": trainer.global_step, "stored_rollout_reward_sum": reward_total,
+            "batch_boundaries": batch_boundaries,
+            "final_pending_reward_sum": trainer.vec_rewards.sum().item(),
+            "final_pending_terminal_count": trainer.vec_terminals.sum().item(),
+            "reward_entries_clipped_by_upstream_trainer": clipped_reward_count,
+            "changed_parameter_tensors": changed, "checkpoint_prediction_match": True,
+            "config": {key: args[key] for key in ["vec", "env", "policy", "torch", "train"]},
+            "losses": epochs,
+            "limitations": ["restricted HOLD-or-Winter-DISBAND policy with initial bias, not a learned baseline",
+                "incomplete actions and existing game defects", "upstream reward clipping",
+                "upstream drops pending rewards and terminals between rollout batches",
+                "no recurrent, GPU, opponent-pool or full training-resume check"],
+        }, indent=2))
+    finally:
+        if trainer is not None:
+            trainer.close()
+        else:
+            vec.close()
+
+
+if __name__ == "__main__":
+    main()