16importnumpyasnp17importtorch1819fromlabmlimporttracker,experiment,logger,monit20fromlabml.internal.configs.dynamic\_hyperparamimportFloatDynamicHyperParam21fromlabml\_helpers.scheduleimportPiecewise22fromlabml\_nn.rl.dqnimportQFuncLoss23fromlabml\_nn.rl.dqn.modelimportModel24fromlabml\_nn.rl.dqn.replay\_bufferimportReplayBuffer25fromlabml\_nn.rl.gameimportWorker

28iftorch.cuda.is\_available():29device=torch.device("cuda:0")30else:31device=torch.device("cpu")

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

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

39classTrainer:

44def\_\_init\_\_(self,\*,45updates:int,epochs:int,46n\_workers:int,worker\_steps:int,mini\_batch\_size:int,47update\_target\_model:int,48learning\_rate:FloatDynamicHyperParam,49):

51self.n\_workers=n\_workers

53self.worker\_steps=worker\_steps

55self.train\_epochs=epochs

58self.updates=updates

60self.mini\_batch\_size=mini\_batch\_size

63self.update\_target\_model=update\_target\_model

66self.learning\_rate=learning\_rate

69self.exploration\_coefficient=Piecewise(70[71(0,1.0),72(25\_000,0.1),73(self.updates/2,0.01)74],outside\_value=0.01)

77self.prioritized\_replay\_beta=Piecewise(78[79(0,0.4),80(self.updates,1)81],outside\_value=1)

84self.replay\_buffer=ReplayBuffer(2\*\*14,0.6)

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

89self.target\_model=Model().to(device)

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

95self.obs=np.zeros((self.n\_workers,4,84,84),dtype=np.uint8)

98forworkerinself.workers:99worker.child.send(("reset",None))

102fori,workerinenumerate(self.workers):103self.obs[i]=worker.child.recv()

106self.loss\_func=QFuncLoss(0.99)

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

110def\_sample\_action(self,q\_value:torch.Tensor,exploration\_coefficient:float):

120withtorch.no\_grad():

122greedy\_action=torch.argmax(q\_value,dim=-1)

124random\_action=torch.randint(q\_value.shape[-1],greedy\_action.shape,device=q\_value.device)

126is\_choose\_rand=torch.rand(greedy\_action.shape,device=q\_value.device)\<exploration\_coefficient

128returntorch.where(is\_choose\_rand,random\_action,greedy\_action).cpu().numpy()

130defsample(self,exploration\_coefficient:float):

134withtorch.no\_grad():

136fortinrange(self.worker\_steps):

138q\_value=self.model(obs\_to\_torch(self.obs))

140actions=self.\_sample\_action(q\_value,exploration\_coefficient)

143forw,workerinenumerate(self.workers):144worker.child.send(("step",actions[w]))

147forw,workerinenumerate(self.workers):

149next\_obs,reward,done,info=worker.child.recv()

152self.replay\_buffer.add(self.obs[w],actions[w],reward,next\_obs,done)

158ifinfo:159tracker.add('reward',info['reward'])160tracker.add('length',info['length'])

163self.obs[w]=next\_obs

165deftrain(self,beta:float):

169for\_inrange(self.train\_epochs):

171samples=self.replay\_buffer.sample(self.mini\_batch\_size,beta)

173q\_value=self.model(obs\_to\_torch(samples['obs']))

177withtorch.no\_grad():

179double\_q\_value=self.model(obs\_to\_torch(samples['next\_obs']))

181target\_q\_value=self.target\_model(obs\_to\_torch(samples['next\_obs']))

184td\_errors,loss=self.loss\_func(q\_value,185q\_value.new\_tensor(samples['action']),186double\_q\_value,target\_q\_value,187q\_value.new\_tensor(samples['done']),188q\_value.new\_tensor(samples['reward']),189q\_value.new\_tensor(samples['weights']))

192new\_priorities=np.abs(td\_errors.cpu().numpy())+1e-6

194self.replay\_buffer.update\_priorities(samples['indexes'],new\_priorities)

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

200self.optimizer.zero\_grad()

202loss.backward()

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

206self.optimizer.step()

208defrun\_training\_loop(self):

214tracker.set\_queue('reward',100,True)215tracker.set\_queue('length',100,True)

218self.target\_model.load\_state\_dict(self.model.state\_dict())219220forupdateinmonit.loop(self.updates):

222exploration=self.exploration\_coefficient(update)223tracker.add('exploration',exploration)

225beta=self.prioritized\_replay\_beta(update)226tracker.add('beta',beta)

229self.sample(exploration)

232ifself.replay\_buffer.is\_full():

234self.train(beta)

237ifupdate%self.update\_target\_model==0:238self.target\_model.load\_state\_dict(self.model.state\_dict())

241tracker.save()

243if(update+1)%1\_000==0:244logger.log()

246defdestroy(self):

251forworkerinself.workers:252worker.child.send(("close",None))

255defmain():

257experiment.create(name='dqn')

260configs={

262'updates':1\_000\_000,

264'epochs':8,

266'n\_workers':8,

268'worker\_steps':4,

270'mini\_batch\_size':32,

272'update\_target\_model':250,

274'learning\_rate':FloatDynamicHyperParam(1e-4,(0,1e-3)),275}

278experiment.configs(configs)

281m=Trainer(\*\*configs)

283withexperiment.start():284m.run\_training\_loop()

286m.destroy()

290if\_\_name\_\_=="\_\_main\_\_":291main()