强化学习A3C算法
强化学习A3C算法
效果：
a3c.py
import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['font.size'] = 18
matplotlib.rcParams['figure.titlesize'] = 18
matplotlib.rcParams['figure.figsize'] = [9, 7]
matplotlib.rcParams['font.family'] = ['KaiTi']
matplotlib.rcParams['axes.unicode_minus']=False
plt.figure()
import os
import threading
import gym
import multiprocessing
import numpy as np
from queue import Queue
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers,optimizers,losses
# os.environ["CUDA_VISIBLE_DEVICES"] = "0" #使用GPU
# 按需占用GPU显存
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# 设置 GPU 显存占用为按需分配，增长式
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e :
# 异常处理
print(e)
SEED_NUM = 1234
tf.random.set_seed(SEED_NUM)
np.random.seed(SEED_NUM)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
assert tf.__version__.startswith('2.')
# 互斥锁,用于线程同步数据
g_mutex = threading.Lock()
class ActorCritic(keras.Model):
""" Actor-Critic模型 """
def __init__(self, state_size, action_size):
super(ActorCritic, self).__init__()
self.state_size = state_size # 状态向量长度
self.action_size = action_size # 动作数量
# 策略网络Actor
self.dense1 = layers.Dense(128, activation='relu')
self.policy_logits = layers.Dense(action_size)
# V网络Critic
self.dense2 = layers.Dense(128, activation='relu')
self.values = layers.Dense(1)
def call(self, inputs):
# 获得策略分布Pi(a|s)
x = self.dense1(inputs)
logits = self.policy_logits(x)
# 获得v(s)
v = self.dense2(inputs)
values = self.values(v)
return logits, values
def record(episode,
episode_reward,
worker_idx,
global_ep_reward,
result_queue,
total_loss,
num_steps):
""" 统计工具函数 """
if global_ep_reward == 0:
global_ep_reward = episode_reward
else:
global_ep_reward = global_ep_reward * 0.99 + episode_reward * 0.01
print(
f"{episode} | "
f"Average Reward: {int(global_ep_reward)} | "
f"Episode Reward: {int(episode_reward)} | "
f"Loss: {int(total_loss / float(num_steps) * 1000) / 1000} | "
f"Steps: {num_steps} | "
f"Worker: {worker_idx}"
)
result_queue.put(global_ep_reward) # 保存回报，传给主线程
return global_ep_reward
class Memory:
""" 数据 """
def __init__(self):
self.states = []
self.actions = []
self.rewards = []
def store(self, state, action, reward):
self.states.append(state)
self.actions.append(action)
self.rewards.append(reward)
def clear(self):
self.states = []
self.actions = []
self.rewards = []
class Agent:
""" 智能体,包含了中央参数网络server """
def __init__(self):
# 服务模型优化器,client不需要,直接从server拉取参数
self.opt = optimizers.Adam(1e-3)
# 服务模型(状态向量,动作数量)
self.server = ActorCritic(4, 2)
self.server(tf.random.normal((2, 4)))
def train(self):
# 共享队列,线程安全，不需要加锁同步
res_queue = Queue()
# 根据cpu线程数量创建多线程Worker
workers = [Worker(self.server, self.opt, res_queue, i)
for i in range(10)] #multiprocessing.cpu_count()
# 启动多线程Worker
for i, worker in enumerate(workers):
print("Starting worker {}".format(i))
worker.start()
# 统计并绘制总回报曲线
returns = []
while True:
reward = res_queue.get()
if reward is not None:
returns.append(reward)
else: # 结束标志
break
# 等待线程退出
[w.join() for w in workers]
print(returns)
plt.figure()
plt.plot(np.arange(len(returns)), returns)
# plt.plot(np.arange(len(moving_average_rewards)), np.array(moving_average_rewards), 's')
plt.xlabel('回合数')
plt.ylabel('总回报')
plt.savefig('a3c-tf-cartpole.svg')
class Worker(threading.Thread):
def __init__(self, server, opt, result_queue, idx):
super(Worker, self).__init__()
self.result_queue = result_queue # 共享队列
self.server = server # 服务模型
self.opt = opt # 服务优化器
self.client = ActorCritic(4, 2) # 线程私有网络
self.worker_idx = idx # 线程id
self.env = gym.make('CartPole-v1').unwrapped #私有环境
self.ep_loss = 0.0
def run(self):
# 每个worker自己维护一个memory
mem = Memory()
# 1回合最大500步
for epi_counter in range(500):
# 复位client游戏状态
current_state,info = self.env.reset(seed=SEED_NUM)
mem.clear()
ep_reward = 0.0
ep_steps = 0
done = False
while not done:
# 输入AC网络状态获得Pi(a|s),未经softmax
logits, _ = self.client(tf.constant(current_state[None, :],dtype=tf.float32))
# 归一化概率
probs = tf.nn.softmax(logits)
# 随机采样动作
action = np.random.choice(2, p=probs.numpy()[0])
# 交互
new_state, reward, done, truncated, info = self.env.step(action)
# 累加奖励
ep_reward += reward
# 记录
mem.store(current_state, action, reward)
# 计算回合步数
ep_steps += 1
# 刷新状态
current_state = new_state
# 最长500步或者规则结束，回合结束
if ep_steps >= 500 or done:
# 计算当前client上的误差
with tf.GradientTape() as tape:
total_loss = self.compute_loss(done, new_state, mem)
# 计算梯度
grads = tape.gradient(total_loss, self.client.trainable_weights)
# 梯度提交到server,在server上更新梯度
global g_mutex
g_mutex.acquire()
self.opt.apply_gradients(zip(grads,self.server.trainable_weights))
g_mutex.release()
# 从server拉取最新的梯度
g_mutex.acquire()
self.client.set_weights(self.server.get_weights())
g_mutex.release()
# 清空Memory
mem.clear()
# 统计此回合回报
self.result_queue.put(ep_reward)
print(f"thread worker_idx : {self.worker_idx}, episode reward : {ep_reward}")
break
# 线程结束
self.result_queue.put(None)
def compute_loss(self,
done,
new_state,
memory,
gamma=0.99):
if done:
reward_sum = 0. # 终止状态的v(终止)=0
else:
# 私有网络根据新状态计算回报
reward_sum = self.client(tf.constant(new_state[None, :],dtype=tf.float32))[-1].numpy()[0]
# 统计折扣回报
discounted_rewards = []
for reward in memory.rewards[::-1]: # reverse buffer r
reward_sum = reward + gamma * reward_sum
discounted_rewards.append(reward_sum)
discounted_rewards.reverse()
# 输入AC网络环境状态获取 Pi(a|s) v(s) 预测值
logits, values = self.client(tf.constant(np.vstack(memory.states), dtype=tf.float32))
# 计算advantage = R() - v(s) = 真实值 - 预测值
advantage = tf.constant(np.array(discounted_rewards)[:, None], dtype=tf.float32) - values
# Critic网络损失
value_loss = advantage ** 2
# 归一化概率预测值Pi(a|s)
policy = tf.nn.softmax(logits)
# 真实动作a 概率预测值Pi(a|s) 交叉熵
policy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=memory.actions, logits=logits)
# 计算策略网络损失时，并不会计算V网络
policy_loss = policy_loss * tf.stop_gradient(advantage)
# 动作概率测值Pi(a|s) 熵
entropy = tf.nn.softmax_cross_entropy_with_logits(labels=policy, logits=logits)
policy_loss = policy_loss - 0.01 * entropy
# 聚合各个误差
total_loss = tf.reduce_mean((0.5 * value_loss + policy_loss))
return total_loss
if __name__ == '__main__':
master = Agent()
master.train()