Spaces:
Running
Running
| """Small dense networks and Adam, shared by PPO and its finite-difference checks.""" | |
| import json | |
| import numpy as np | |
| class Dense: | |
| def __init__(self,model=None,sizes=None,mean=None,std=None,rng=None): | |
| rng=rng or np.random.default_rng(1) | |
| if model: | |
| self.meta=model;self.mean=np.array(model['inputMean'],dtype='f');self.std=np.array(model['inputStd'],dtype='f');self.w=[];self.b=[];width=model['inputDim'] | |
| for layer in model['layers']: | |
| n=len(layer['bias']);self.w.append(np.array(layer['weight'],dtype='f').reshape(n,width).T.copy());self.b.append(np.array(layer['bias'],dtype='f'));width=n | |
| else: | |
| self.meta={};self.mean=mean;self.std=std;self.w=[(rng.standard_normal((a,b))*np.sqrt(1/a)).astype('f') for a,b in zip(sizes,sizes[1:])];self.b=[np.zeros(n,'f') for n in sizes[1:]] | |
| self.p=[v for pair in zip(self.w,self.b) for v in pair];self.m=[np.zeros_like(v) for v in self.p];self.v=[np.zeros_like(v) for v in self.p];self.steps=0 | |
| def forward(self,x,cache=False): | |
| a=(x-self.mean)/self.std;acts=[a] | |
| for i,(w,b) in enumerate(zip(self.w,self.b)):a=a@w+b;a=np.tanh(a) if i<len(self.w)-1 else a;acts.append(a) | |
| return (a,acts) if cache else a | |
| def gradients(self,acts,grad): | |
| result=[] | |
| for i in range(len(self.w)-1,-1,-1): | |
| result[0:0]=[acts[i].T@grad,grad.sum(0)] | |
| if i:grad=(grad@self.w[i].T)*(1-acts[i]**2) | |
| return result | |
| def update(self,grads,rate): | |
| self.steps+=1;norm=np.sqrt(sum(float((g*g).sum()) for g in grads));scale=min(1.,.5/max(1e-8,norm)) | |
| for i,(p,g) in enumerate(zip(self.p,grads)): | |
| g=g*scale;self.m[i]=.9*self.m[i]+.1*g;self.v[i]=.999*self.v[i]+.001*g*g;p-=rate*(self.m[i]/(1-.9**self.steps))/(np.sqrt(self.v[i]/(1-.999**self.steps))+1e-8) | |
| def export(self,path,name,metadata): | |
| model={**self.meta,'name':name,'training':'PPO reinforcement learning after behavior cloning','rl':metadata,'layers':[{'weight':w.T.flatten().tolist(),'bias':b.tolist()} for w,b in zip(self.w,self.b)]} | |
| path.write_text(json.dumps(model,separators=(',',':'))) | |
| def sigmoid(x):return 1/(1+np.exp(-np.clip(x,-30,30))) | |
| def log_prob(mean,z,fire,sd): | |
| p=sigmoid(mean[:,4]);return -.5*(((z-mean[:,:4])/sd)**2+2*np.log(sd)+np.log(2*np.pi)).sum(1)+fire*np.log(p+1e-8)+(1-fire)*np.log(1-p+1e-8) | |
| def ppo_gradient(mean,z,fire,oldlog,adv,sd,clip=.2): | |
| logp=log_prob(mean,z,fire,sd);ratio=np.exp(np.clip(logp-oldlog,-15,15));active=np.where(adv>=0,ratio<=1+clip,ratio>=1-clip);factor=-adv*ratio*active/len(adv) | |
| grad=np.zeros_like(mean);grad[:,:4]=factor[:,None]*(z-mean[:,:4])/(sd**2);grad[:,4]=factor*(fire-sigmoid(mean[:,4]));return grad,float(np.mean(oldlog-logp)),float(np.mean(np.minimum(ratio*adv,np.clip(ratio,1-clip,1+clip)*adv))) | |
| def gae(rewards,values,next_value,dones,gamma=.99,lam=.95): | |
| adv=np.zeros_like(rewards);last=np.zeros_like(next_value) | |
| for t in range(len(rewards)-1,-1,-1): | |
| following=next_value if t==len(rewards)-1 else values[t+1];nonterminal=1-dones[t];delta=rewards[t]+gamma*following*nonterminal-values[t];last=delta+gamma*lam*nonterminal*last;adv[t]=last | |
| return adv,adv+values | |
| MACRO_HEADS=[(0,11),(11,4),(15,7),(22,9),(31,5),(38,3)] | |
| def macro_probabilities(raw,group_masks): | |
| result=[] | |
| for start,width in MACRO_HEADS: | |
| logits=raw[:,start:start+width].copy() | |
| if start==22:logits=np.where(group_masks,logits,-1e9) | |
| p=np.exp(logits-logits.max(1,keepdims=True));p/=p.sum(1,keepdims=True);result.append(p) | |
| return result | |
| def macro_log_prob(raw,categories,z,masks,sd=.3): | |
| probabilities=macro_probabilities(raw,masks);log=np.zeros(len(raw),dtype=raw.dtype) | |
| for i,p in enumerate(probabilities):log+=np.log(p[np.arange(len(p)),categories[:,i]]+1e-8) | |
| return log-.5*(((z-raw[:,36:38])/sd)**2+2*np.log(sd)+np.log(2*np.pi)).sum(1) | |
| def macro_ppo_gradient(raw,categories,z,masks,oldlog,adv,sd=.3,clip=.2): | |
| log=macro_log_prob(raw,categories,z,masks,sd);ratio=np.exp(np.clip(log-oldlog,-15,15));active=np.where(adv>=0,ratio<=1+clip,ratio>=1-clip);factor=-adv*ratio*active/len(raw);grad=np.zeros_like(raw) | |
| for h,((start,width),p) in enumerate(zip(MACRO_HEADS,macro_probabilities(raw,masks))): | |
| dp=-p;dp[np.arange(len(p)),categories[:,h]]+=1;grad[:,start:start+width]=dp*factor[:,None] | |
| grad[:,36:38]=factor[:,None]*(z-raw[:,36:38])/sd**2 | |
| return grad,float(np.mean(oldlog-log)) | |