| |
| """Train frozen architecture families from scratch on external Gymnasium controls.""" |
| from __future__ import annotations |
| import argparse,csv,json,sys |
| from pathlib import Path |
| import numpy as np, torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| sys.path.insert(0,str(Path(__file__).resolve().parents[1])) |
| from spectral_world_models.models import build_model |
| from spectral_world_models.v14_external_control import ExternalConfig,collect_external_npz,ExternalTransitionDataset,ExternalSequenceDataset,load_meta,SUPPORTED_ENVS |
| from spectral_world_models.metrics import count_parameters |
|
|
| MODELS=("neural_operator","swm_structured_cf_v9","no_spectral_transition") |
|
|
| def loss_fn(model,b,vocab_size): |
| o=model(b['image'],b['tokens'],b['action']) |
| with torch.no_grad(): tz=model.encode(b['next_image'],b['next_tokens']) |
| im=F.mse_loss(o['image_pred'],b['next_image']); lat=F.mse_loss(o['z_pred'],tz) |
| txt=F.cross_entropy(o['text_logits'].reshape(-1,vocab_size),b['next_tokens'].reshape(-1),ignore_index=0) |
| return im+.15*lat+.05*txt+.01*model.stability_penalty() |
|
|
| def rollout_metrics(model,loader,device,horizons): |
| sums={h:[] for h in horizons} |
| model.eval() |
| with torch.no_grad(): |
| for b in loader: |
| imgs=b['images'].to(device); toks=b['tokens'].to(device); acts=b['actions'].to(device); z=model.encode(imgs[:,0],toks[:,0]) |
| step=[] |
| for t in range(acts.shape[1]): |
| z=model.transition(z,acts[:,t]); p=model.image_decoder(z); step.append(((p-imgs[:,t+1])**2).flatten(1).mean(1).cpu()) |
| for h in horizons: |
| hh=min(h,len(step)); sums[h].append(torch.stack(step[:hh],1).mean(1)) |
| return {f'rollout_mse_h{h}':float(torch.cat(sums[h]).mean()) for h in horizons} |
|
|
| def train(env_id,model_name,seed,args,out): |
| data=out/'data'/f'{env_id}.npz' |
| if not data.exists(): |
| collect_external_npz(data,ExternalConfig(env_id=env_id,seq_len=max(args.horizons)+1,train_episodes=args.train_episodes,val_episodes=args.val_episodes,test_episodes=args.test_episodes,seed=args.data_seed)) |
| meta=load_meta(data); device=torch.device(args.device) |
| torch.manual_seed(seed) |
| model=build_model(model_name,image_size=32,vocab_size=meta['vocab_size'],max_text_len=8,action_dim=meta['action_dim']).to(device) |
| tr=DataLoader(ExternalTransitionDataset(data,'train'),batch_size=args.batch_size,shuffle=True); te=DataLoader(ExternalSequenceDataset(data,'test'),batch_size=args.batch_size) |
| opt=torch.optim.AdamW(model.parameters(),lr=args.lr) |
| model.train() |
| for _ in range(args.epochs): |
| for b in tr: |
| b={k:v.to(device) for k,v in b.items()}; opt.zero_grad(set_to_none=True); L=loss_fn(model,b,meta['vocab_size']); L.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),1.0); opt.step() |
| m=rollout_metrics(model,te,device,args.horizons) |
| m.update(env_id=env_id,model=model_name,seed=seed,params=count_parameters(model)) |
| ck=out/'checkpoints'/env_id; ck.mkdir(parents=True,exist_ok=True); torch.save({'model':model.state_dict(),'meta':meta,'model_name':model_name,'seed':seed},ck/f'{model_name}_seed{seed}.pt') |
| return m |
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('--envs',nargs='+',default=list(SUPPORTED_ENVS)); p.add_argument('--models',nargs='+',default=list(MODELS)); p.add_argument('--seeds',nargs='+',type=int,default=[0,1,2,3,4]); p.add_argument('--horizons',nargs='+',type=int,default=[5,10,20,30]); p.add_argument('--epochs',type=int,default=5); p.add_argument('--batch-size',type=int,default=64); p.add_argument('--lr',type=float,default=2e-3); p.add_argument('--train-episodes',type=int,default=256); p.add_argument('--val-episodes',type=int,default=64); p.add_argument('--test-episodes',type=int,default=64); p.add_argument('--data-seed',type=int,default=1400); p.add_argument('--device',default='cpu'); p.add_argument('--out',default='results/v14_external_control'); a=p.parse_args() |
| out=Path(a.out); out.mkdir(parents=True,exist_ok=True); rows=[] |
| for e in a.envs: |
| for m in a.models: |
| for s in a.seeds: |
| print(f'[V14] {e} {m} seed={s}',flush=True); rows.append(train(e,m,s,a,out)) |
| fields=list(rows[0]); |
| with open(out/'metrics_by_seed.csv','w',newline='') as f: w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(rows) |
| |
| summary=[] |
| for e in a.envs: |
| for m in a.models: |
| rr=[r for r in rows if r['env_id']==e and r['model']==m]; row={'env_id':e,'model':m,'params':rr[0]['params'],'n_seeds':len(rr)} |
| for h in a.horizons: |
| vals=np.array([r[f'rollout_mse_h{h}'] for r in rr]); row[f'rollout_mse_h{h}_mean']=vals.mean(); row[f'rollout_mse_h{h}_std']=vals.std(ddof=1) if len(vals)>1 else 0. |
| summary.append(row) |
| sf=list(summary[0]); |
| with open(out/'metrics_summary.csv','w',newline='') as f: w=csv.DictWriter(f,fieldnames=sf); w.writeheader(); w.writerows(summary) |
| (out/'run_config.json').write_text(json.dumps(vars(a),indent=2)); print(out/'metrics_summary.csv') |
| if __name__=='__main__': main() |
|
|
