16fromtypingimportDict1718importnumpyasnp19importtorch20fromtorchimportnn21fromtorchimportoptim22fromtorch.distributionsimportCategorical2324fromlabmlimportmonit,tracker,logger,experiment25fromlabml.configsimportFloatDynamicHyperParam,IntDynamicHyperParam26fromlabml\_helpers.moduleimportModule27fromlabml\_nn.rl.gameimportWorker28fromlabml\_nn.rl.ppoimportClippedPPOLoss,ClippedValueFunctionLoss29fromlabml\_nn.rl.ppo.gaeimportGAE

32iftorch.cuda.is\_available():33device=torch.device("cuda:0")34else:35device=torch.device("cpu")

38classModel(Module):

43def\_\_init\_\_(self):44super().\_\_init\_\_()

48self.conv1=nn.Conv2d(in\_channels=4,out\_channels=32,kernel\_size=8,stride=4)

52self.conv2=nn.Conv2d(in\_channels=32,out\_channels=64,kernel\_size=4,stride=2)

56self.conv3=nn.Conv2d(in\_channels=64,out\_channels=64,kernel\_size=3,stride=1)

61self.lin=nn.Linear(in\_features=7\*7\*64,out\_features=512)

64self.pi\_logits=nn.Linear(in\_features=512,out\_features=4)

67self.value=nn.Linear(in\_features=512,out\_features=1)

70self.activation=nn.ReLU()

72defforward(self,obs:torch.Tensor):73h=self.activation(self.conv1(obs))74h=self.activation(self.conv2(h))75h=self.activation(self.conv3(h))76h=h.reshape((-1,7\*7\*64))7778h=self.activation(self.lin(h))7980pi=Categorical(logits=self.pi\_logits(h))81value=self.value(h).reshape(-1)8283returnpi,value

86defobs\_to\_torch(obs:np.ndarray)-\>torch.Tensor:

88returntorch.tensor(obs,dtype=torch.float32,device=device)/255.

91classTrainer:

96def\_\_init\_\_(self,\*,97updates:int,epochs:IntDynamicHyperParam,98n\_workers:int,worker\_steps:int,batches:int,99value\_loss\_coef:FloatDynamicHyperParam,100entropy\_bonus\_coef:FloatDynamicHyperParam,101clip\_range:FloatDynamicHyperParam,102learning\_rate:FloatDynamicHyperParam,103):

107self.updates=updates

109self.epochs=epochs

111self.n\_workers=n\_workers

113self.worker\_steps=worker\_steps

115self.batches=batches

117self.batch\_size=self.n\_workers\*self.worker\_steps

119self.mini\_batch\_size=self.batch\_size//self.batches120assert(self.batch\_size%self.batches==0)

123self.value\_loss\_coef=value\_loss\_coef

125self.entropy\_bonus\_coef=entropy\_bonus\_coef

128self.clip\_range=clip\_range

130self.learning\_rate=learning\_rate

135self.workers=[Worker(47+i)foriinrange(self.n\_workers)]

138self.obs=np.zeros((self.n\_workers,4,84,84),dtype=np.uint8)139forworkerinself.workers:140worker.child.send(("reset",None))141fori,workerinenumerate(self.workers):142self.obs[i]=worker.child.recv()

145self.model=Model().to(device)

148self.optimizer=optim.Adam(self.model.parameters(),lr=2.5e-4)

151self.gae=GAE(self.n\_workers,self.worker\_steps,0.99,0.95)

154self.ppo\_loss=ClippedPPOLoss()

157self.value\_loss=ClippedValueFunctionLoss()

159defsample(self)-\>Dict[str,torch.Tensor]:

164rewards=np.zeros((self.n\_workers,self.worker\_steps),dtype=np.float32)165actions=np.zeros((self.n\_workers,self.worker\_steps),dtype=np.int32)166done=np.zeros((self.n\_workers,self.worker\_steps),dtype=np.bool)167obs=np.zeros((self.n\_workers,self.worker\_steps,4,84,84),dtype=np.uint8)168log\_pis=np.zeros((self.n\_workers,self.worker\_steps),dtype=np.float32)169values=np.zeros((self.n\_workers,self.worker\_steps+1),dtype=np.float32)170171withtorch.no\_grad():

173fortinrange(self.worker\_steps):

176obs[:,t]=self.obs

179pi,v=self.model(obs\_to\_torch(self.obs))180values[:,t]=v.cpu().numpy()181a=pi.sample()182actions[:,t]=a.cpu().numpy()183log\_pis[:,t]=pi.log\_prob(a).cpu().numpy()

186forw,workerinenumerate(self.workers):187worker.child.send(("step",actions[w,t]))188189forw,workerinenumerate(self.workers):

191self.obs[w],rewards[w,t],done[w,t],info=worker.child.recv()

196ifinfo:197tracker.add('reward',info['reward'])198tracker.add('length',info['length'])

201\_,v=self.model(obs\_to\_torch(self.obs))202values[:,self.worker\_steps]=v.cpu().numpy()

205advantages=self.gae(done,rewards,values)

208samples={209'obs':obs,210'actions':actions,211'values':values[:,:-1],212'log\_pis':log\_pis,213'advantages':advantages214}

218samples\_flat={}219fork,vinsamples.items():220v=v.reshape(v.shape[0]\*v.shape[1],\*v.shape[2:])221ifk=='obs':222samples\_flat[k]=obs\_to\_torch(v)223else:224samples\_flat[k]=torch.tensor(v,device=device)225226returnsamples\_flat

228deftrain(self,samples:Dict[str,torch.Tensor]):

238for\_inrange(self.epochs()):

240indexes=torch.randperm(self.batch\_size)

243forstartinrange(0,self.batch\_size,self.mini\_batch\_size):

245end=start+self.mini\_batch\_size246mini\_batch\_indexes=indexes[start:end]247mini\_batch={}248fork,vinsamples.items():249mini\_batch[k]=v[mini\_batch\_indexes]

252loss=self.\_calc\_loss(mini\_batch)

255forpginself.optimizer.param\_groups:256pg['lr']=self.learning\_rate()

258self.optimizer.zero\_grad()

260loss.backward()

262torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=0.5)

264self.optimizer.step()

266@staticmethod267def\_normalize(adv:torch.Tensor):

269return(adv-adv.mean())/(adv.std()+1e-8)

271def\_calc\_loss(self,samples:Dict[str,torch.Tensor])-\>torch.Tensor:

277sampled\_return=samples['values']+samples['advantages']

283sampled\_normalized\_advantage=self.\_normalize(samples['advantages'])

287pi,value=self.model(samples['obs'])

290log\_pi=pi.log\_prob(samples['actions'])

293policy\_loss=self.ppo\_loss(log\_pi,samples['log\_pis'],sampled\_normalized\_advantage,self.clip\_range())

299entropy\_bonus=pi.entropy()300entropy\_bonus=entropy\_bonus.mean()

303value\_loss=self.value\_loss(value,samples['values'],sampled\_return,self.clip\_range())

308loss=(policy\_loss309+self.value\_loss\_coef()\*value\_loss310-self.entropy\_bonus\_coef()\*entropy\_bonus)

313approx\_kl\_divergence=.5\*((samples['log\_pis']-log\_pi)\*\*2).mean()

316tracker.add({'policy\_reward':-policy\_loss,317'value\_loss':value\_loss,318'entropy\_bonus':entropy\_bonus,319'kl\_div':approx\_kl\_divergence,320'clip\_fraction':self.ppo\_loss.clip\_fraction})321322returnloss

324defrun\_training\_loop(self):

330tracker.set\_queue('reward',100,True)331tracker.set\_queue('length',100,True)332333forupdateinmonit.loop(self.updates):

335samples=self.sample()

338self.train(samples)

341tracker.save()

343if(update+1)%1\_000==0:344logger.log()

346defdestroy(self):

351forworkerinself.workers:352worker.child.send(("close",None))

355defmain():

357experiment.create(name='ppo')

359configs={

361'updates':10000,

365'epochs':IntDynamicHyperParam(8),

367'n\_workers':8,

369'worker\_steps':128,

371'batches':4,

375'value\_loss\_coef':FloatDynamicHyperParam(0.5),

379'entropy\_bonus\_coef':FloatDynamicHyperParam(0.01),

381'clip\_range':FloatDynamicHyperParam(0.1),

385'learning\_rate':FloatDynamicHyperParam(1e-3,(0,1e-3)),386}387388experiment.configs(configs)

391m=Trainer(\*\*configs)

394withexperiment.start():395m.run\_training\_loop()

397m.destroy()

401if\_\_name\_\_=="\_\_main\_\_":402main()