| """From-scratch shared 1--15 DOF robot motion/energy MoE training and reporting.""" |
| from __future__ import annotations |
| import argparse,json,csv,math,random,time,copy,sys |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
| from robotics_calibrated import RoboticsMoECalibrated as RoboticsMoE,PreparedCalibrated as Prepared,prediction |
| ROOT=Path(__file__).resolve().parents[1] |
|
|
| def load(split): |
| with np.load(ROOT/'data'/'robotics'/f'{split}.npz') as data: |
| |
| d={k:data[k] for k in data.files} |
| with np.load(ROOT/'data'/'robotics'/f'{split}_gravity.npz') as extra: |
| d.update({k:extra[k] for k in extra.files}) |
| with np.load(ROOT/'data'/'robotics'/f'{split}_past_calibration.npz') as past: |
| d.update({k:past[k] for k in past.files}) |
| return d |
|
|
| def scores(true,pred): |
| true=np.asarray(true,float);pred=np.asarray(pred,float) |
| err=pred-true |
| return dict(MAPE_pct=float(np.mean(abs(err)/np.maximum(abs(true),1e-8))*100), |
| WAPE_pct=float(abs(err).sum()/max(abs(true).sum(),1e-8)*100), |
| MSE=float(np.mean(err*err)),MAE=float(np.mean(abs(err))), |
| R2=float(1-np.sum(err**2)/max(np.sum((true-true.mean())**2),1e-9))) |
|
|
| def grouped(y,p,axes,top,split): |
| rows=[] |
| for typ,labels in [('all',['all']),('axes',list(range(1,16))),('topology',sorted(set(top))),('axes_topology',[(n,t) for n in range(1,16) for t in sorted(set(top)) if np.any((axes==n)&(top==t))])]: |
| for lab in labels: |
| sel=np.ones(len(axes),bool) if typ=='all' else (axes==lab if typ=='axes' else ((axes==lab[0])&(top==lab[1]) if typ=='axes_topology' else top==lab)) |
| for k,name,unit in ((0,'time_s','s'),(1,'energy_Wh','Wh')): |
| rows.append({'split':split,'group_type':typ,'group':f'{lab[0]:02d}_{lab[1]}' if typ=='axes_topology' else lab,'metric_target':name,'unit':unit, |
| 'n_programs':int(sel.sum()),**scores(y[sel,k],p[sel,k])}) |
| return rows |
|
|
| def run(args): |
| t0=time.time();random.seed(args.seed);np.random.seed(args.seed);torch.manual_seed(args.seed) |
| torch.set_num_threads(1) |
| device=('cuda' if torch.cuda.is_available() else 'cpu') if args.device=='auto' else args.device |
| if device=='cuda' and not torch.cuda.is_available():raise RuntimeError('CUDA unavailable') |
| train=Prepared(load('train')) |
| sets={'train':train} |
| for split in ('val','test','ood'): |
| sets[split]=Prepared(load(split),train.mean,train.std) |
| m=RoboticsMoE().to(device);optim=torch.optim.AdamW(m.parameters(),lr=args.lr,weight_decay=.00008) |
| step=0;best=float('inf');best_state=None;best_ep=0;history=[];wait=0 |
| J,G,H,M=train.tensors(device) |
| Y=torch.from_numpy(train.target).to(device) |
| for ep in range(1,args.epochs+1): |
| m.train();perm=torch.randperm(len(Y));losses=[] |
| for ix in perm.split(args.batch): |
| optim.zero_grad() |
| o=m(J[ix],M[ix],G[ix],H[ix]);loss=torch.nn.functional.smooth_l1_loss(o,Y[ix],beta=.025) |
| loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),3) |
| optim.step();losses.append(float(loss.detach())) |
| v=sets['val'];yp=prediction(m,v) |
| valerr=np.abs(yp[:,1]-v.y[:,1])/np.maximum(v.y[:,1],1e-8)*100 |
| val_time=np.abs(yp[:,0]-v.y[:,0])/np.maximum(v.y[:,0],1e-8)*100 |
| by_axes=[float(np.mean(valerr[v.axes==n])) for n in range(1,16)] |
| by_configuration=[float(np.mean(valerr[(v.axes==n)&(v.topology==t)])) for n in range(1,16) for t in set(v.topology) if np.any((v.axes==n)&(v.topology==t))] |
| criterion=float(np.mean(by_configuration)+.25*np.mean(val_time)) |
| history.append({'epoch':ep,'train_loss':float(np.mean(losses)),'val_mean_axis_energy_MAPE':np.mean(by_axes),'val_mean_configuration_energy_MAPE':np.mean(by_configuration),'val_time_MAPE':np.mean(val_time),'val_criterion':criterion}) |
| if criterion<best-1e-6: |
| best=criterion;best_state=copy.deepcopy(m.state_dict());best_ep=ep;wait=0 |
| else:wait+=1 |
| if ep==1 or ep%10==0: |
| print(f'ROBOT EPOCH {ep} train={np.mean(losses):.5f} val_energy={np.mean(by_configuration):.3f}% val_time={np.mean(val_time):.3f}% best_ep={best_ep}',flush=True) |
| if wait>=args.patience:break |
| m.load_state_dict(best_state) |
| out=ROOT/'reports'/'robotics';out.mkdir(parents=True,exist_ok=True) |
| results={};allrows=[] |
| for split,d in sets.items(): |
| p=prediction(m,d);cold=prediction(m,d,use_history=False) |
| rr=grouped(d.y,p,d.axes,d.topology,split) |
| allrows+=rr |
| results[split]={'duration_s':scores(d.y[:,0],p[:,0]),'energy_Wh':scores(d.y[:,1],p[:,1]), |
| 'cold_energy_Wh':scores(d.y[:,1],cold[:,1]), |
| 'energy_by_axis_MAPE_pct':{str(n):scores(d.y[d.axes==n,1],p[d.axes==n,1])['MAPE_pct'] for n in range(1,16)}} |
| if split in ('val','test','ood'): |
| with (out/f'{split}_predictions.csv').open('w',newline='') as f: |
| w=csv.writer(f);w.writerow(['robot_id','n_axes','topology','actual_time_s','predicted_time_s','actual_energy_Wh','predicted_energy_Wh','nominal_energy_Wh','cold_energy_Wh']) |
| for j in range(len(p)): |
| w.writerow([d.ids[j],int(d.axes[j]),d.topology[j],float(d.y[j,0]),float(p[j,0]),float(d.y[j,1]),float(p[j,1]),float(d.prior[j,1]),float(cold[j,1])]) |
| with (out/'all_metrics.csv').open('w',newline='') as f: |
| w=csv.DictWriter(f,fieldnames=list(allrows[0]));w.writeheader();w.writerows(allrows) |
| report={'checkpoint_selection':'minimum validation mean-per-axis-topology energy MAPE + 0.25 time MAPE; test and OOD never used for selection', |
| 'architecture':'signed gravity+subtree context; MCE latent from six independent past sensor trajectories; physically calibrated baseline with learnable residual', |
| 'initialized_from_scratch':True,'best_epoch':best_ep,'runtime_s':time.time()-t0, |
| 'trainable_parameters':sum(p.numel() for p in m.parameters()),'seed':args.seed, |
| 'results':results,'history':history,'notes':['No true machine efficiency coefficients in model inputs.','Each robot ID belongs to exactly one split.','OOD tests deliberately altered mass and link length distribution.','Synthetic Newton-Euler and electrical surrogate, not MuJoCo benchmark.']} |
| (out/'metrics.json').write_text(json.dumps(report,indent=2,default=lambda x: x.item() if hasattr(x,'item') else str(x))) |
| torch.save({'model_type':'RoboticsMoECalibrated','robotics':m.cpu().state_dict(), |
| 'train_normalizer':{'mean':train.mean,'std':train.std},'hyperparameters':vars(args), |
| 'best_epoch':best_ep,'data_manifest':json.loads((ROOT/'data/robotics/manifest.json').read_text())},out/'robotics_model.pt') |
| print('FINISHED',json.dumps({'epochs':best_ep,'results':results,'time_s':round(time.time()-t0,1)},indent=2),flush=True) |
|
|
| if __name__=='__main__': |
| pa=argparse.ArgumentParser();pa.add_argument('--epochs',type=int,default=55);pa.add_argument('--patience',type=int,default=16) |
| pa.add_argument('--batch',type=int,default=168);pa.add_argument('--lr',type=float,default=.0017) |
| pa.add_argument('--seed',type=int,default=20261018);pa.add_argument('--device',choices=['auto','cpu','cuda'],default='auto');run(pa.parse_args()) |
|
|
