Package: torchreinforce
torchreinforceA pythonic implementation of the REINFORCE algorithm that is actually fun to useInstallationYou can install it with pip as you would for any other python packagepip install torchreinforceQuickstartIn order to use the REINFORCE algorithm with your model you only need to do two things:Use theReinforceModuleclass as your base classDecorate yourforwardfunction with@ReinforceModule.forwardThat's it!classModel(ReinforceModule):def__init__(self,**kwargs):super(Model,self).__init__(**kwargs)self.net=torch.nn.Sequential(torch.nn.Linear(20,128),torch.nn.ReLU(),torch.nn.Linear(128,2),torch.nn.Softmax(dim=-1),)@ReinforceModule.forwarddefforward(self,x):returnself.net(x)Your model will now outputReinforceOutputobjects.This objects have two important functionsget()reward(value)You can useoutput.get()to get an actual sample of the overlaying distribution andoutput.reward(value)to set a reward for the specific output.Beingnetyour model you have to do something like thataction=net(observation)observation,reward,done,info=env.step(action.get())action.reward(reward)Wait, did you just said distribution?Yes! As the REINFORCE algorithm states the outputs of your model will be used as parameters for a probability distribution function.Actually you can use whatever probability distribution you want, theReinforceModuleconstructor accepts indeed the following parameters:gammathegammaparameter of the REINFORCE algorithm (default:Categorical)distributioneveryReinforceDistributionorpytorch.distributionsdistribution (default: 0.99)like thatnet=Model(distribution=torch.distributions.Beta,gamma=0.99)Keep in mind that the outputs of yourdecoratedforward(x)outputs will be used as the parameters for thedistribution. If yourdistributionneeds more than one parameters just return a list.I've added the possibility to distribution to have adeterministicbehavior intestingand I've implemented it only for theCategoricaldistribution, if you want to implement your own deterministic logic check the filedistributions/categorical.pyit is pretty straightforwardIf you want to use thetorch.distributions.Betadistribution for example you will need to do something likeclassModel(ReinforceModule):def__init__(self,**kwargs):super(Model,self).__init__(**kwargs)...@ReinforceModule.forwarddefforward(self,x):return[self.net1(x),self.net2(x)]# the Beta distribution accepts two parametersnet=Model(distribution=torch.distributions.Beta,gamma=0.99)action=net(inp)env.step(action.get())Nice! What about training?You can compute the REINFORCE loss by calling theloss()function ofReinforceModuleand than treat it as you would do with any other pytorch loss functionnet=...optmizer=...whiletraining:net.reset()forsteps:....loss=net.loss(normalize=True)optimizer.zero_grad()loss.backward()optmizer.step()Youhave tocall thereset()function ofReinforceModulebeforethe beginning of each episode. You can also pass the argumentnormalizetoloss()if you want to normalize the rewardsPutting all togetherA complete example looks like this:classModel(ReinforceModule):def__init__(self,**kwargs):super(Model,self).__init__(**kwargs)self.net=torch.nn.Sequential(torch.nn.Linear(4,128),torch.nn.ReLU(),torch.nn.Linear(128,2),torch.nn.Softmax(dim=-1),)@ReinforceModule.forwarddefforward(self,x):returnself.net(x)env=gym.make('CartPole-v0')net=Model()optimizer=torch.optim.Adam(net.parameters(),lr=0.001)foriinrange(EPISODES):done=Falsenet.reset()observation=env.reset()whilenotdone:action=net(torch.tensor(observation,dtype=torch.float32))observation,reward,done,info=env.step(action.get())action.reward(reward)loss=net.loss(normalize=False)optimizer.zero_grad()loss.backward()optimizer.step()You can find a running example in theexamples/folder.