A3C 算法的简单实现(ChatGPT)

import gym
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import threading
import multiprocessing
import os

# 定义Actor-Critic模型
class ActorCriticModel(tf.keras.Model):
    def __init__(self, state_size, action_size):
        super(ActorCriticModel, self).__init__()
        self.state_size = state_size
        self.action_size = action_size

        self.dense1 = Dense(128, activation='relu')  # 第一个隐藏层
        self.policy_logits = Dense(action_size)     # 输出动作概率的层
        self.dense2 = Dense(128, activation='relu')  # 第二个隐藏层
        self.values = Dense(1)                      # 输出状态值的层

    def call(self, inputs):
        x = self.dense1(inputs)
        logits = self.policy_logits(x)  # 计算动作概率
        v = self.dense2(inputs)
        values = self.values(v)         # 计算状态值
        return logits, values


# 训练函数
def train(global_model, optimizer, global_step):
    env = gym.make('CartPole-v1')  # 创建CartPole环境，替换为实际环境名称

    max_episodes = 10000  # 最大训练次数
    gamma = 0.99          # 折扣因子
    update_freq = 5       # 更新频率
    num_workers = multiprocessing.cpu_count()  # 并发智能体数

    for episode in range(max_episodes):
        state = env.reset()
        state = np.reshape(state, [1, state_size])

        with tf.GradientTape() as tape:
            total_reward = 0
            num_steps = 0
            done = False

            while not done:
                num_steps += 1
                logits, values = global_model(state)
                probs = tf.nn.softmax(logits)  # 使用softmax计算动作概率
                action = np.random.choice(action_size, p=probs.numpy()[0])  # 根据概率随机选择动作

                next_state, reward, done, _ = env.step(action)
                next_state = np.reshape(next_state, [1, state_size])
                total_reward += reward

                if done or num_steps >= 300:  # 限制最大步数
                    total_reward = -100 if not done else total_reward
                    R = 0
                else:
                    _, R = global_model(next_state)  # 获取下一状态的值

                td_target = reward + gamma * R
                td_error = td_target - values
                actor_loss = -tf.math.log(probs[0, action]) * td_error  # Actor的损失
                critic_loss = tf.square(td_error)                     # Critic的损失

                total_loss = actor_loss + critic_loss  # 综合Actor和Critic的损失
                grads = tape.gradient(total_loss, global_model.trainable_variables)  # 计算梯度
                optimizer.apply_gradients(zip(grads, global_model.trainable_variables))  # 更新参数

                if global_step % update_freq == 0:
                    global_model.set_weights(model.get_weights())  # 更新全局模型参数

                global_step += 1

                state = next_state

        print(f"Episode {episode+1}: Total Reward = {total_reward}")


if __name__ == "__main__":
    state_size = 4  # 状态空间维度，替换为实际状态空间维度
    action_size = 2  # 动作空间维度，替换为实际动作空间维度

    global_model = ActorCriticModel(state_size, action_size)  # 创建全局模型
    optimizer = Adam(learning_rate=0.001)  # Adam优化器，替换为实际学习率
    global_step = 0

    workers = []
    for _ in range(num_workers):
        worker = threading.Thread(target=train, args=(global_model, optimizer, global_step))
        workers.append(worker)

    for worker in workers:
        worker.start()

    for worker in workers:
        worker.join()