pip install torch gym numpy

"""《基于迁移学习的强化学习代码示例》
    时间：2024.09.22
    作者：不去幼儿园
"""
import gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque

# 定义 Q 网络
class QNetwork(nn.Module):
	def __init__(self, state_dim, action_dim):
		super(QNetwork, self).__init__()
		self.fc1 = nn.Linear(state_dim, 128)
		self.fc2 = nn.Linear(128, 128)
		self.fc3 = nn.Linear(128, action_dim)

	def forward(self, x):
		x = torch.relu(self.fc1(x))
		x = torch.relu(self.fc2(x))
		return self.fc3(x)

# Replay buffer，用于存储经验
class ReplayBuffer:
	def __init__(self, capacity):
		self.buffer = deque(maxlen=capacity)

	def add(self, transition):
		self.buffer.append(transition)

	def sample(self, batch_size):
		transitions = random.sample(self.buffer, batch_size)

		# 从 transitions 中解压各个元素
		states, actions, rewards, next_states, dones = zip(*transitions)

		# 将 states 和 next_states 转换为 numpy 数组时保持维度一致，使用 np.stack 而不是 np.array
		states = np.stack(states)  # 将状态堆叠为多维数组
		next_states = np.stack(next_states)  # 保证 next_states 也被堆叠为多维数组

		actions = np.array(actions, dtype=np.int64)  # 动作通常是整数
		rewards = np.array(rewards, dtype=np.float32)  # 奖励通常是浮点数
		dones = np.array(dones, dtype=np.float32)  # dones 是布尔值，用浮点数表示

		return states, actions, rewards, next_states, dones

	def size(self):
		return len(self.buffer)

# 选择动作
def select_action(state, policy_net, epsilon, action_dim):
	if random.random() < epsilon:
		return random.choice(np.arange(action_dim))
	else:
		with torch.no_grad():
			state = torch.FloatTensor(state).unsqueeze(0)
			q_values = policy_net(state)
			return q_values.argmax().item()

# Q-learning 更新
def update_model(policy_net, target_net, optimizer, replay_buffer, batch_size, gamma):
	if replay_buffer.size() < batch_size:
		return

	states, actions, rewards, next_states, dones = replay_buffer.sample(batch_size)

	states = torch.FloatTensor(states)
	actions = torch.LongTensor(actions)
	rewards = torch.FloatTensor(rewards)
	next_states = torch.FloatTensor(next_states)
	dones = torch.FloatTensor(dones)

	q_values = policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)
	next_q_values = target_net(next_states).max(1)[0]
	target_q_values = rewards + gamma * next_q_values * (1 - dones)

	loss = (q_values - target_q_values.detach()).pow(2).mean()

	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

# 训练模型
def train_dqn(env_name, num_episodes=500, gamma=0.99, epsilon_start=1.0, epsilon_end=0.01, epsilon_decay=0.995, batch_size=64):
	env = gym.make(env_name)
	state_dim = env.observation_space.shape[0]
	action_dim = env.action_space.n

	policy_net = QNetwork(state_dim, action_dim)
	target_net = QNetwork(state_dim, action_dim)
	target_net.load_state_dict(policy_net.state_dict())
	target_net.eval()

	optimizer = optim.Adam(policy_net.parameters())
	replay_buffer = ReplayBuffer(10000)
	epsilon = epsilon_start

	for episode in range(num_episodes):
		state, _ = env.reset()  # 只获取 state，忽略 info
		total_reward = 0

		while True:
			action = select_action(state, policy_net, epsilon, action_dim)
			next_state, reward, done, _, __ = env.step(action)  # 只获取 next_state，忽略其他
			replay_buffer.add((state, action, reward, next_state, done))
			state = next_state
			total_reward += reward

			update_model(policy_net, target_net, optimizer, replay_buffer, batch_size, gamma)

			if done:
				print(f"Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}")
				break

		epsilon = max(epsilon_end, epsilon_decay * epsilon)

		if episode % 10 == 0:
			target_net.load_state_dict(policy_net.state_dict())

	return policy_net

# 在源环境（CartPole）中训练模型
policy_net_cartpole = train_dqn(env_name='CartPole-v1')

# 定义函数以匹配网络的结构，进行部分权重的迁移
def transfer_weights(policy_net, target_env_state_dim, target_env_action_dim):
	# 获取预训练的网络
	pretrained_dict = policy_net.state_dict()

	# 创建新网络，适应目标环境的状态和动作维度
	new_policy_net = QNetwork(target_env_state_dim, target_env_action_dim)

	# 获取新网络的权重
	new_dict = new_policy_net.state_dict()

	# 仅保留在预训练网络和新网络中都有的层（即隐藏层的参数）
	pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in new_dict and 'fc1' not in k and 'fc3' not in k}

	# 更新新网络的权重
	new_dict.update(pretrained_dict)

	# 将更新后的字典加载到新模型中
	new_policy_net.load_state_dict(new_dict)

	return new_policy_net


# 微调模型
def fine_tune_dqn(policy_net, env_name, num_episodes=200, gamma=0.99, epsilon_start=0.1, epsilon_end=0.01, epsilon_decay=0.995, batch_size=64):
	env = gym.make(env_name)
	target_state_dim = env.observation_space.shape[0]
	target_action_dim = env.action_space.n

	# 调用 transfer_weights 函数，将权重从 CartPole 模型迁移到新的 MountainCar 模型
	policy_net = transfer_weights(policy_net, target_state_dim, target_action_dim)

	target_net = QNetwork(target_state_dim, target_action_dim)
	target_net.load_state_dict(policy_net.state_dict())
	target_net.eval()

	optimizer = optim.Adam(policy_net.parameters())
	replay_buffer = ReplayBuffer(10000)
	epsilon = epsilon_start

	for episode in range(num_episodes):
		state, _ = env.reset()  # 只获取 state，忽略 info
		total_reward = 0

		while True:
			action = select_action(state, policy_net, epsilon, target_action_dim)
			next_state, reward, done, _, __ = env.step(action)  # 只获取 next_state，忽略其他
			replay_buffer.add((state, action, reward, next_state, done))
			state = next_state
			total_reward += reward

			update_model(policy_net, target_net, optimizer, replay_buffer, batch_size, gamma)

			if done:
				print(f"Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}")
				break

		epsilon = max(epsilon_end, epsilon_decay * epsilon)

		if episode % 10 == 0:
			target_net.load_state_dict(policy_net.state_dict())

	return policy_net

# 微调源环境训练的策略网络到目标环境 MountainCar
fine_tuned_policy_net = fine_tune_dqn(policy_net_cartpole, env_name='MountainCar-v0')

# 测试微调后的策略并显示动画
def test_policy_with_animation(policy_net, env_name, num_episodes=5, max_steps=200):
	# 初始化环境并启用动画显示
	env = gym.make(env_name, render_mode="human")
	total_rewards = []

	for episode in range(num_episodes):
		state, _ = env.reset()  # 获取初始状态
		total_reward = 0

		for step in range(max_steps):
			# 选择动作，不使用探索
			action = select_action(state, policy_net, epsilon=0.0, action_dim=env.action_space.n)
			next_state, reward, done, _, __ = env.step(action)  # 执行动作
			total_reward += reward
			state = next_state  # 更新状态

			if done:
				break  # 结束当前回合

		total_rewards.append(total_reward)
		print(f"Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}")

	# 关闭环境
	env.close()
	print(f"Average Total Reward over {num_episodes} episodes: {np.mean(total_rewards)}")
	return total_rewards


# 测试微调后的策略网络在 MountainCar 环境中的表现
test_rewards = test_policy_with_animation(fine_tuned_policy_net, env_name='MountainCar-v0')