A3Cでcarttpole-v1の学習がうまくいっていない件

実現したいこと

numyとthreadingでOpen AI gym のcartpole-v1を学習させたい

前提

https://qiita.com/s05319ss/items/2fe9bfe562fea1707e79

上の記事を参考にA3CでOpen AI gym のcartpole-v1を学習させているが、うまく学習ができない

発生している問題・エラーメッセージ

獲得報酬が20ほどであり、最大報酬である300(私が報酬をいじっています)を獲得できるよううまく学習ができていない

該当のソースコード

Python
1import numpy as np
2import gym
3import matplotlib.pylab as plt
4import random
5import threading
6import copy
7
8class Environment:
9    def __init__(self,batchsize=1):
10        self.batchsize=batchsize
11        self.state=np.zeros((batchsize, 4))
12        self.step=np.zeros(batchsize)
13        self.score=np.zeros(batchsize)
14        self.flag=np.ones(batchsize).astype(bool)
15        self.environment=gym.make('CartPole-v1', render_mode='rgb_array')
16        initial_state, info = self.environment.reset()
17        for i in range(4):
18            self.state[0][i] = initial_state[i]
19            
20    def transition(self,action):
21        a = np.copy(action)
22        for i in range(self.batchsize):
23            next_state, reward, terminated, truncated, info = self.environment.step(int(a[i]))
24            self.flag[i] = not terminated
25            if self.flag[i]==True:
26                self.score[i]+= reward
27                self.step[i]+=1
28                if self.step[i]>=200:
29                    self.flag[i]=False
30                    self.score[i]+= 100
31            else:
32                self.score[i]+= (-1)
33            for j in range(4):
34                self.state[i][j] = next_state[j]
35        
36
37    def initialize(self, flag=None):
38        if flag is None:
39            flag = np.ones(self.state.shape[0], dtype=bool)
40        self.state[flag]*=0
41        self.step[flag]*=0
42        self.score[flag]*=0
43        self.flag[flag]=np.ones(np.sum(flag)).astype(bool)
44        initial_state, info = self.environment.reset()
45        for i in range(4):
46            self.state[0][i] = initial_state[i]
47
48
49# Agent to play the above RL task with finite, discrete action space.
50class Agent_Ex1_Discrete:
51    def __init__(self):
52        self.sharedbody=SharedBody()
53        self.vhead=Affine()
54        self.pihead=Affine()
55        self.V=None
56        self.probdist=None
57        self.a=None
58        self.action_kinds=None
59        self.H=None
60        self.action_number = 2
61        
62    def action(self,state,params):
63        s=copy.deepcopy(state)
64        latentvect=self.sharedbody.forward(s,params['sharedbody'])
65        self.V=self.vhead.forward(latentvect,params['vhead'])
66        preprobdist=self.pihead.forward(latentvect,params['pihead'])
67        self.probdist=softmax(preprobdist)
68        self.action_kinds = np.zeros(self.probdist.shape[0])
69        self.a=np.zeros([self.probdist.shape[0],self.action_number])
70        for i in range(self.action_kinds.shape[0]):
71            self.action_kinds[i] = np.random.choice(a=self.action_number,p=self.probdist[i, :])
72            self.a[i][int(self.action_kinds[i])] = 1
73        self.H=-np.sum(self.probdist*np.log(self.probdist),axis=1).reshape(-1,1)
74        return self.a
75  #def gradsは元の記事と同じ
76　def initialparams(self):
77        params={}
78        params['sharedbody']={}
79        params['vhead']={}
80        params['pihead']={}
81        params['sharedbody']['W1']=np.random.randn(4,10)* 0.01
82        params['sharedbody']['b1']=np.zeros(10)
83        params['sharedbody']['W2']=np.random.randn(10,10)/np.sqrt(10)
84        params['sharedbody']['b2']=np.zeros(10)
85        params['vhead']['W']=np.random.randn(10,1)/np.sqrt(10)
86        params['vhead']['b']=np.zeros(1)
87        params['pihead']['W']=np.random.randn(10,2)/np.sqrt(10)
88        params['pihead']['b']=np.zeros(2)
89        return params
90
91#def softmax,class SharedBody,class Affine, class Sigmoid, class Tanhは元の記事と同じ
92
93    
94# Global computation in the A3C algorithm.
95class A3C:
96    def __init__(self,t_max,T_max,agents,environments,params=None,gamma=0.99,alpha=0.99,learningrate=7e-4,vweight=0.5,entropyweight=0.01,epsilon=0.1,epoch=10):
97        self.t_max=t_max
98        self.T_max=T_max
99        self.N=len(agents)
100        self.gamma=gamma
101        self.lr=learningrate
102        self.kappa=vweight
103        self.beta=entropyweight
104        self.epsilon=epsilon
105        self.c=epoch
106        self.scorehistory=np.array([])
107        if params==None:
108            self.params=agents[0].initialparams()
109        else:
110            self.params=params.copy
111        self.g={}
112        for key1 in self.params.keys():
113            self.g[key1]={}
114            for key2 in self.params[key1].keys():
115                self.g[key1][key2]=np.zeros_like(self.params[key1][key2])
116        self.alpha=alpha
117        self.T=0
118        self.threadinstances={}
119        for i in range(self.N):
120            self.threadinstances[i]=A3Cthread(agents[i],environments[i])
121        self.accumulatedscores=0.0
122        self.numberofsamples=0
123    def update(self,threadinstance):
124        while self.T<self.T_max:
125            eta=-self.T*self.lr/float(self.T_max)+self.lr
126            threadgrads,cnt,ifterminated,threadscore=threadinstance.grads(self.params,self.gamma,self.t_max,self.kappa,self.beta)
127            for key1 in self.params.keys():
128                for key2 in self.params[key1].keys():
129                    self.g[key1][key2]*=self.alpha
130                    self.g[key1][key2]+=(1-self.alpha)*threadgrads[key1][key2]**2
131                    self.params[key1][key2]-=eta*threadgrads[key1][key2]/np.sqrt(self.g[key1][key2]+self.epsilon)
132            self.T+=cnt
133            self.accumulatedscores+=ifterminated*threadscore
134            self.numberofsamples+=ifterminated
135            if self.numberofsamples>self.c:
136                self.scorehistory=np.append(self.scorehistory,self.accumulatedscores/float(max(1,self.numberofsamples)))
137                self.accumulatedscores*=0
138                self.numberofsamples*=0
139                print(str(self.T)+" steps done with average score:")
140                print(self.scorehistory[-1])
141    def train(self):
142        threads=[]
143        for i in range(self.N):
144            threads.append(threading.Thread(target=self.update,args=(self.threadinstances[i],)))
145        for i in range(self.N):
146            threads[i].start()
147        for i in range(self.N):
148            threads[i].join()
149        return self.params
150
151# Local computation for each A3C thread.
152class A3Cthread:
153    def __init__(self,agent,environment):
154        self.agent=agent
155        self.environment=environment
156        self.environment.initialize()
157        self.t=0
158        self.vlast=None
159        self.alast=None
160        
161    def grads(self,globalparams,gamma,t_max,kappa=0.5,beta=0.01):
162        params={}
163        grads={}
164        for key1 in globalparams.keys():
165            params[key1]={}
166            grads[key1]={}
167            for key2 in globalparams[key1].keys():
168                params[key1][key2]=globalparams[key1][key2].copy()
169                grads[key1][key2]=np.zeros_like(globalparams[key1][key2])
170        gradstoaccumulate={}
171        rewards={}
172        vvalues={}
173        t_start=self.t
174        while bool((self.t-t_start<t_max)*(np.prod(self.environment.flag))):
175            rewards[self.t]=-self.environment.score.copy()
176            if bool((self.t>0)*(self.t==t_start)):
177                a=self.alast.copy()
178                vvalues[self.t]=self.vlast.copy()
179            else:
180                a=self.agent.action(self.environment.state,params)
181                vvalues[self.t]=self.agent.V.copy()
182            gradstoaccumulate[self.t]=self.agent.grads(1,2*kappa,beta)
183            self.environment.transition(self.agent.action_kinds)
184            rewards[self.t]+=self.environment.score.copy()
185            self.t+=1
186        self.alast=self.agent.action(self.environment.state,params)
187        self.vlast=self.agent.V.copy()
188        R=self.vlast
189        for i in range(self.t-t_start):
190            R=gamma*R+rewards[self.t-i-1]
191            for key1 in globalparams.keys():
192                for key2 in globalparams[key1].keys():
193                    grads[key1][key2]-=np.sum(R-vvalues[self.t-i-1])*gradstoaccumulate[self.t-i-1][0][key1][key2]+gradstoaccumulate[self.t-i-1][1][key1][key2]
194        threadscore=np.sum(self.environment.score.copy())
195        ifterminated=not(bool(np.prod(self.environment.flag)))
196        if ifterminated:
197            self.environment.initialize()
198        return grads,self.t-t_start,ifterminated,threadscore
199
200# Showing how the agent with finite, discrete action space performs before training.
201testagent=Agent_Ex1_Discrete()
202untrainedparams=testagent.initialparams()
203testenvironment=Environment()
204while np.sum(testenvironment.flag)>0:
205    a=testagent.action(testenvironment.state,untrainedparams)
206    testenvironment.transition(testagent.action_kinds)
207print("Average performance of the agent with finite, discrete action space before training:")
208print(np.sum(testenvironment.score))
209# Training the agent.
210agents=[Agent_Ex1_Discrete()]
211environments=[Environment()]
212for i in range(7):
213    agents.append(Agent_Ex1_Discrete())
214    environments.append(Environment())
215trainer=A3C(t_max=5,T_max=90000,agents=agents,environments=environments)
216print("Training the agent by the A3C algorithm:")
217trainedparams=trainer.train()
218scorehistory_discrete=trainer.scorehistory.copy()
219# Showing how the A3C-trained agent with finite, discrete action space performs.
220testagent=Agent_Ex1_Discrete()
221testenvironment=Environment()
222while np.sum(testenvironment.flag)>0:
223    a=testagent.action(testenvironment.state,trainedparams)
224    testenvironment.transition(testagent.action_kinds)
225print("Average performance of the A3C-trained agent with finite, discrete action space:")
226print(np.sum(testenvironment.score))