mirror of
https://github.com/jcreek/CosmicClash.git
synced 2026-09-10 16:04:04 +00:00
Josh Creek
87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
"""Export a trained SB3 checkpoint to the JSON format PolicyNetwork.gd loads.
|
|
|
|
The exported file contains the deterministic policy MLP (obs -> action means);
|
|
the game clamps outputs to [-1, 1] and treats the last value as turbo (> 0).
|
|
A parity self-check compares the JSON forward pass against SB3's own
|
|
deterministic prediction before writing.
|
|
|
|
Example:
|
|
.venv/bin/python export_policy.py checkpoints/smoke/final.zip ../Game/bots/rookie.json
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import pathlib
|
|
|
|
import numpy as np
|
|
import torch
|
|
from stable_baselines3 import PPO
|
|
|
|
|
|
def linear_to_layer(linear: torch.nn.Linear, activation: str) -> dict:
|
|
return {
|
|
"weights": linear.weight.detach().cpu().numpy().tolist(),
|
|
"biases": linear.bias.detach().cpu().numpy().tolist(),
|
|
"activation": activation,
|
|
}
|
|
|
|
|
|
def extract_layers(policy) -> list[dict]:
|
|
# Features extractor must be a passthrough (flatten) for this export to
|
|
# be faithful; it has no parameters for our flat "obs" Box space.
|
|
n_extractor_params = sum(p.numel() for p in policy.features_extractor.parameters())
|
|
assert n_extractor_params == 0, "features extractor has weights; export logic needs updating"
|
|
|
|
layers = []
|
|
modules = list(policy.mlp_extractor.policy_net)
|
|
for i, module in enumerate(modules):
|
|
if isinstance(module, torch.nn.Linear):
|
|
next_is_tanh = i + 1 < len(modules) and isinstance(modules[i + 1], torch.nn.Tanh)
|
|
assert next_is_tanh or i + 1 >= len(modules), (
|
|
f"unsupported activation after layer {i}: {modules[i + 1] if i + 1 < len(modules) else None}"
|
|
)
|
|
layers.append(linear_to_layer(module, "tanh" if next_is_tanh else "linear"))
|
|
elif not isinstance(module, torch.nn.Tanh):
|
|
raise AssertionError(f"unsupported module in policy net: {module}")
|
|
layers.append(linear_to_layer(policy.action_net, "linear"))
|
|
return layers
|
|
|
|
|
|
def json_forward(layers: list[dict], obs: np.ndarray) -> np.ndarray:
|
|
x = obs
|
|
for layer in layers:
|
|
x = np.asarray(layer["weights"]) @ x + np.asarray(layer["biases"])
|
|
if layer["activation"] == "tanh":
|
|
x = np.tanh(x)
|
|
return x
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("checkpoint", help="SB3 checkpoint .zip (e.g. checkpoints/smoke/final.zip)")
|
|
parser.add_argument("output", help="Output JSON path (e.g. ../Game/bots/rookie.json)")
|
|
args = parser.parse_args()
|
|
|
|
model = PPO.load(args.checkpoint, device="cpu")
|
|
policy = model.policy
|
|
layers = extract_layers(policy)
|
|
input_size = model.observation_space["obs"].shape[0]
|
|
|
|
# Parity check: JSON forward pass must match SB3's deterministic action
|
|
rng = np.random.default_rng(0)
|
|
for _ in range(16):
|
|
obs = rng.uniform(-1, 1, input_size).astype(np.float32)
|
|
expected, _ = model.predict({"obs": obs}, deterministic=True)
|
|
actual = np.clip(json_forward(layers, obs), -1.0, 1.0)
|
|
assert np.allclose(actual, expected, atol=1e-5), f"parity check failed: {actual} vs {expected}"
|
|
|
|
output = pathlib.Path(args.output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(output, "w") as f:
|
|
json.dump({"input_size": int(input_size), "layers": layers}, f)
|
|
print(f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, parity OK)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|