| """current sign-aware hierarchical gravity expert; learns robot design response. |
| |
| Only robot expert is changed. Legacy 11 manufacturing operations untouched. |
| Features: 13 original joint + 10 signed gravity/topology per joint; 18 original |
| physical aggregate + 8 signed gravity global; 12 history probes. |
| """ |
| import numpy as np |
| import torch |
| from torch import nn |
| from robotics_core import Prepared as PreviousPrepared |
|
|
| class Prepared51(PreviousPrepared): |
| def __init__(self,data,mean=None,std=None): |
| super().__init__(data,mean,std) |
| rawj=np.asarray(data['joint_gravity'],np.float32) |
| rawg=np.asarray(data['global_gravity'],np.float32) |
| if mean is None: |
| active=rawj[data['mask']>0] |
| self.mean['joint_gravity']=active.mean(axis=0) |
| self.std['joint_gravity']=np.maximum(active.std(axis=0),.03) |
| self.mean['global_gravity']=rawg.mean(axis=0) |
| self.std['global_gravity']=np.maximum(rawg.std(axis=0),.03) |
| nj=np.clip((rawj-self.mean['joint_gravity'])/self.std['joint_gravity'],-8,8)*data['mask'][:,:,None] |
| ng=np.clip((rawg-self.mean['global_gravity'])/self.std['global_gravity'],-8,8) |
| self.norm['joint']=np.concatenate([self.norm['joint'],nj],axis=2).astype(np.float32) |
| self.norm['global_features']=np.concatenate([self.norm['global_features'],ng],axis=1).astype(np.float32) |
|
|
| class RoboticsMoE51(nn.Module): |
| """Version-scaled sign-aware hierarchical robot model; 1-15 axes only.""" |
| def __init__(self,joint_dim=23,global_dim=26,history_dim=12,hidden=None): |
| super().__init__() |
| from config import choice |
| hidden=choice(24,48,72) if hidden is None else hidden |
| mc1,mc2=choice((24,16),(40,28),(64,40)) |
| co1,co2=choice((48,24),(72,40),(96,64)) |
| gr1,gr2=choice((24,16),(40,24),(64,40)) |
| he1,he2=choice((48,24),(96,48),(128,64)) |
| self.joint_expert=nn.Sequential(nn.Linear(joint_dim,hidden),nn.GELU(),nn.Linear(hidden,hidden),nn.GELU()) |
| self.machine_expert=nn.Sequential(nn.Linear(history_dim,mc1),nn.GELU(),nn.Linear(mc1,mc2),nn.Tanh()) |
| self.coupling_expert=nn.Sequential(nn.Linear(hidden*6+global_dim,co1),nn.GELU(),nn.Linear(co1,co2),nn.GELU()) |
| self.gravity_expert=nn.Sequential(nn.Linear(8+mc2,gr1),nn.GELU(),nn.Linear(gr1,gr2),nn.GELU()) |
| self.motion_energy_head=nn.Sequential(nn.Linear(co2+mc2+gr2+global_dim,he1),nn.GELU(),nn.Linear(he1,he2),nn.GELU(),nn.Linear(he2,2)) |
| def forward(self,joint,mask,global_x,history): |
| x=self.joint_expert(joint)*mask.unsqueeze(-1) |
| count=mask.sum(dim=1,keepdim=True).clamp(min=1) |
| pooled=x.sum(dim=1)/count |
| mx=x.masked_fill(mask.unsqueeze(-1)==0,-1e4).max(dim=1).values |
| arm1=(joint[:,:,10]>.5).float()*mask |
| torso=(joint[:,:,11]>.5).float()*mask |
| arm0=(1.-arm1)*(1.-torso)*mask |
| def pool(w):return (x*w.unsqueeze(-1)).sum(dim=1)/w.sum(dim=1,keepdim=True).clamp(min=1) |
| sum_scale=x.sum(dim=1)/count.sqrt() |
| coupling=self.coupling_expert(torch.cat([pooled,mx,pool(arm0),pool(arm1),pool(torso),sum_scale,global_x],dim=1)) |
| machine=self.machine_expert(history) |
| gravity=self.gravity_expert(torch.cat([global_x[:,-8:],machine],dim=1)) |
| return self.motion_energy_head(torch.cat([coupling,machine,gravity,global_x],dim=1)) |
|
|
| def prediction(model,ds,batch=512,device=None,use_history=True): |
| model.eval();device=device or next(model.parameters()).device;j,g,h,mask=ds.tensors(device) |
| if not use_history:h=torch.zeros_like(h) |
| parts=[] |
| with torch.no_grad(): |
| for k in range(0,len(ds.y),batch): |
| parts.append(model(j[k:k+batch],mask[k:k+batch],g[k:k+batch],h[k:k+batch]).cpu().numpy()) |
| ratio=np.exp(np.clip(np.concatenate(parts),-1.5,1.5)) |
| return ratio*ds.prior |
|
|
