import torch, math
def find_sd(c, prefer='act_net_state_dict'):
    if isinstance(c, dict):
        if prefer in c and isinstance(c[prefer], dict): return c[prefer]
        for k in ['act_net_state_dict','model_state_dict','state_dict','model','policy','nets']:
            if k in c and isinstance(c[k], dict): return c[k]
        ten = {k:v for k,v in c.items() if torch.is_tensor(v)}
        if ten: return ten
    return c
BC='/data0/hongyi/act_runs/open_laptop_100demos_chunk25/policy_epoch_4000_seed_0.ckpt'
G='/data0/hongyi/act_sim_ppo_v6_gtsim/run_openlaptop_grpo_bc_start/checkpoints'
bc=find_sd(torch.load(BC,map_location='cpu',weights_only=False))
g64=find_sd(torch.load(f'{G}/iter_00064.pt',map_location='cpu',weights_only=False))
g256=find_sd(torch.load(f'{G}/iter_00256.pt',map_location='cpu',weights_only=False))
print('bc keys:',len(bc),'| g64 keys:',len(g64))
common=set(bc)&set(g64)
print('common keys:',len(common))
if len(common) < 0.4*len(g64):
    print('  bc sample:',list(bc)[:3]); print('  g64 sample:',list(g64)[:3])
    # try stripping a leading prefix from one side
    import os
def cmp(a,b,label):
    bb_w=bb_d=tr_w=tr_d=0.0; nbb=ntr=0; biggest=[]
    for k in common:
        if not torch.is_tensor(a[k]) or not a[k].is_floating_point(): continue
        nw=a[k].float().norm().item(); d=(b[k].float()-a[k].float()).norm().item()
        rel=d/(nw+1e-9)
        if 'backbone' in k.lower(): bb_w+=nw*nw; bb_d+=d*d; nbb+=1
        else:
            tr_w+=nw*nw; tr_d+=d*d; ntr+=1; biggest.append((rel,k,nw,d))
    print(f'  {label}')
    print(f'    frozen backbone:  |Δ|={math.sqrt(bb_d):8.4f}  |W|={math.sqrt(bb_w):8.1f}  rel={math.sqrt(bb_d/(bb_w+1e-9)):.5f}  (n={nbb})')
    print(f'    transf+heads:     |Δ|={math.sqrt(tr_d):8.4f}  |W|={math.sqrt(tr_w):8.1f}  rel={math.sqrt(tr_d/(tr_w+1e-9)):.5f}  (n={ntr})')
    biggest.sort(reverse=True)
    for rel,k,nw,d in biggest[:4]: print(f'      top-Δ {k:40s} rel={rel:.3f} |W|={nw:.2f} |Δ|={d:.3f}')
cmp(bc,g64,'BC -> ep64  (peak 85%)')
cmp(bc,g256,'BC -> ep256 (degraded 70%)')