113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
"""Compare iter100 vs iter150 checkpoints: action_mean diff and refine_mask equality."""
|
|
import numpy as np
|
|
import torch
|
|
from torch_geometric.data import Batch
|
|
|
|
from src.network import create_model
|
|
from src.utils import load_checkpoint, setup_helmholtz_config
|
|
|
|
|
|
def load_config():
|
|
from src.utils import load_config as _lc
|
|
from pathlib import Path
|
|
cfg_path = Path(__file__).resolve().parent / "src" / "config.yaml"
|
|
return _lc(str(cfg_path))
|
|
|
|
|
|
def compare_checkpoints(ckpt_a, ckpt_b, label_a="iter100", label_b="iter150"):
|
|
config = load_config()
|
|
setup_helmholtz_config(config)
|
|
algo = config.get("algorithm", {})
|
|
|
|
from environment.mesh_refinement import MeshRefinement
|
|
|
|
env = MeshRefinement(
|
|
environment_config=config.get("environment", {}).get("mesh_refinement", {}),
|
|
seed=99,
|
|
)
|
|
|
|
# ── Load both models ──
|
|
model_a = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
|
load_checkpoint(model_a, ckpt_a)
|
|
model_a.eval()
|
|
|
|
model_b = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
|
load_checkpoint(model_b, ckpt_b)
|
|
model_b.eval()
|
|
|
|
# ── Get same initial observation ──
|
|
env.reset()
|
|
obs = env.reset() # second reset ensures same state
|
|
|
|
with torch.no_grad():
|
|
batch = Batch.from_data_list([obs])
|
|
|
|
# Model A
|
|
shared_a, batch_a = model_a._encode(batch)
|
|
latent_pi_a = model_a.policy_mlp(shared_a)
|
|
action_mean_a = model_a.action_out(latent_pi_a).cpu().numpy().flatten()
|
|
dist_a = model_a._make_distribution(latent_pi_a)
|
|
actions_a = dist_a.get_actions(deterministic=True).cpu().numpy().flatten()
|
|
|
|
# Model B
|
|
shared_b, batch_b = model_b._encode(batch)
|
|
latent_pi_b = model_b.policy_mlp(shared_b)
|
|
action_mean_b = model_b.action_out(latent_pi_b).cpu().numpy().flatten()
|
|
dist_b = model_b._make_distribution(latent_pi_b)
|
|
actions_b = dist_b.get_actions(deterministic=True).cpu().numpy().flatten()
|
|
|
|
# ── Compare action_mean ──
|
|
diff = action_mean_a - action_mean_b
|
|
print(f"\n{'='*60}")
|
|
print(f" 1. action_mean comparison")
|
|
print(f"{'='*60}")
|
|
print(f" {label_a} action_mean: min={action_mean_a.min():.6f} max={action_mean_a.max():.6f} mean={action_mean_a.mean():.6f} std={action_mean_a.std():.6f}")
|
|
print(f" {label_b} action_mean: min={action_mean_b.min():.6f} max={action_mean_b.max():.6f} mean={action_mean_b.mean():.6f} std={action_mean_b.std():.6f}")
|
|
print(f" ---")
|
|
print(f" |diff|: min={np.abs(diff).min():.8f} max={np.abs(diff).max():.8f} mean={np.abs(diff).mean():.8f}")
|
|
print(f" diff = 0 exactly: {int(np.sum(diff == 0))} / {len(diff)} ({100 * np.sum(diff == 0) / len(diff):.2f}%)")
|
|
print(f" |diff| < 1e-6: {int(np.sum(np.abs(diff) < 1e-6))} / {len(diff)}")
|
|
print(f" |diff| < 1e-4: {int(np.sum(np.abs(diff) < 1e-4))} / {len(diff)}")
|
|
print(f" cos similarity: {np.dot(action_mean_a, action_mean_b) / (np.linalg.norm(action_mean_a) * np.linalg.norm(action_mean_b) + 1e-12):.8f}")
|
|
|
|
# ── Compare refine_mask (action > 0) ──
|
|
mask_a = actions_a > 0.0
|
|
mask_b = actions_b > 0.0
|
|
mask_equal = np.array_equal(mask_a, mask_b)
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f" 2. refine_mask comparison")
|
|
print(f"{'='*60}")
|
|
print(f" {label_a} refine_mask: sum={mask_a.sum()} / {len(mask_a)} ({100 * mask_a.sum() / len(mask_a):.1f}%)")
|
|
print(f" {label_b} refine_mask: sum={mask_b.sum()} / {len(mask_b)} ({100 * mask_b.sum() / len(mask_b):.1f}%)")
|
|
print(f" refine_mask exactly equal: {mask_equal}")
|
|
print(f" mask XOR sum: {(mask_a ^ mask_b).sum()} / {len(mask_a)}")
|
|
|
|
if not mask_equal:
|
|
diff_idx = np.where(mask_a != mask_b)[0]
|
|
print(f" First 20 differing indices: {diff_idx[:20].tolist()}")
|
|
print(f" At those indices, {label_a} action_mean: {action_mean_a[diff_idx[:10]]}")
|
|
print(f" At those indices, {label_b} action_mean: {action_mean_b[diff_idx[:10]]}")
|
|
|
|
# ── 3. Parameter-level diff ──
|
|
print(f"\n{'='*60}")
|
|
print(f" 3. Model parameter weight diff (L2 norm)")
|
|
print(f"{'='*60}")
|
|
sd_a = torch.load(ckpt_a, map_location="cpu")["model_state_dict"]
|
|
sd_b = torch.load(ckpt_b, map_location="cpu")["model_state_dict"]
|
|
for k in sorted(sd_a.keys()):
|
|
w_a = sd_a[k].float()
|
|
w_b = sd_b[k].float()
|
|
l2 = torch.norm(w_a - w_b).item()
|
|
rel = l2 / (torch.norm(w_a).item() + 1e-12)
|
|
print(f" {k:55s} |Δ|₂={l2:.6e} rel={rel:.6e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
d1 = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/model_iter0100.pt"
|
|
d2 = sys.argv[2] if len(sys.argv) > 2 else "checkpoints/model_iter0150.pt"
|
|
l1 = sys.argv[3] if len(sys.argv) > 3 else "iter100"
|
|
l2 = sys.argv[4] if len(sys.argv) > 4 else "iter150"
|
|
compare_checkpoints(d1, d2, l1, l2)
|