Actor-Critic Reinforcement Learning,network can't learning - reinforcement-learning

Question

I have some Actor-Critic reinforcement learning questions want to ask.In my code,I try to use two network(actor and critic) to play Pong-game,but I don't know my "learn()" function have problem?The "Loss1" is actor_network loss(-logP(a|s) * TDerror),the "Loss2" is critic_network loss((pre - TDtarget) ** 2 / 2).
class Actor_Critic():
def __init__(self,input_dim=6400,n_actions=2):
self.input_dim = input_dim
self.n_action = n_actions//Pong game action space
self.lr = 0.0005
self.gamma = 0.99 // discounted factor
self.critic_network = self.__build_critic_network()
self.actor_network = self.__build_actor_network()
def __build_critic_network(self):
model_input = layers.Input(shape=(self.input_dim + self.n_action,))
layer1 = layers.Dense(128, activation='relu')(model_input)
layer2 = layers.Dense(32, activation='relu')(layer1)
model_output = layers.Dense(1, activation=None)(layer2)
model = Model(model_input,model_output)
model.compile(optimizer=Adam(learning_rate=self.lr))
return model
def __build_actor_network(self):
model_input = layers.Input(shape=(self.input_dim,))
layer1 = layers.Dense(128, activation='relu')(model_input)
layer2 = layers.Dense(32, activation='relu')(layer1)
model_output = layers.Dense(self.n_action, activation='softmax')(layer2)
model = Model(model_input, model_output)
model.compile(optimizer=Adam(learning_rate=self.lr))
return model
def choose_action(self,state):
state = tf.convert_to_tensor([state],dtype=tf.float16)
probability = self.actor_network(state)
action_probs = tfp.distributions.Categorical(probs=probability)
action = action_probs.sample()
return action_probs,action.numpy()[0] // return action probability and action
def action_hot_code(self, action):
action_hot_code = np.zeros(self.n_action, dtype=np.float16)
action_hot_code[action] = 1.0
action_hot_code = tf.convert_to_tensor([action_hot_code], dtype=tf.float16)
return action_hot_code
def learn(self,state,reward,next_state,done):
with tf.GradientTape(persistent=True) as tape:
action_probs, action = self.choose_action(state)
next_action_prob, next_action = self.choose_action(next_state)
action_hot_code = self.action_hot_code(action)
next_action_hot_code = self.action_hot_code(next_action)
state = tf.convert_to_tensor([state], dtype=tf.float16)
next_state = tf.convert_to_tensor([next_state], dtype=tf.float16)
state_action = tf.concat([state, action_hot_code], axis=1)
next_state_action = tf.concat([next_state, next_action_hot_code], axis=1)
state_value = self.critic_network(state_action)
next_state_value = self.critic_network(next_state_action)
TDtarget = reward + self.gamma * next_state_value * (1-done)
TDerror = state_value - TDtarget
Loss1 = -TDerror * action_probs.log_prob(action)
Loss2 = (TDerror ** 2) / 2
Loss1 = tf.squeeze(Loss1)
Loss2 = tf.squeeze(Loss2)
gradient1 = tape.gradient(Loss1, self.actor_network.trainable_variables)
gradient2 = tape.gradient(Loss2,self.critic_network.trainable_variables)
self.actor_network.optimizer.apply_gradients(zip(gradient1, self.actor_network.trainable_variables))
self.critic_network.optimizer.apply_gradients(zip(gradient2,self.critic_network.trainable_variables))