139 lines
5.9 KiB
Python
139 lines
5.9 KiB
Python
import torch
|
|
import torch.optim as optim
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import random
|
|
from networks_cont import ContActor, ContQCritic
|
|
|
|
# --- 经验回放池 (Replay Buffer) 补丁 ---
|
|
class ReplayBuffer:
|
|
def __init__(self, capacity):
|
|
self.capacity = capacity
|
|
self.buffer = []
|
|
self.position = 0
|
|
|
|
def add(self, state, action, reward, next_state, done):
|
|
if len(self.buffer) < self.capacity:
|
|
self.buffer.append(None)
|
|
self.buffer[self.position] = (state, action, reward, next_state, done)
|
|
self.position = (self.position + 1) % self.capacity
|
|
|
|
def sample(self, batch_size):
|
|
batch = random.sample(self.buffer, batch_size)
|
|
# 解包数据，转换为 numpy 数组以提高效率
|
|
state, action, reward, next_state, done = map(np.stack, zip(*batch))
|
|
return state, action, reward, next_state, done
|
|
|
|
def __len__(self):
|
|
return len(self.buffer)
|
|
|
|
# --- DDPG 算法代理 ---
|
|
class DDPGAgent:
|
|
def __init__(self, state_dim, action_dim, max_action, device,
|
|
actor_lr=1e-4, critic_lr=1e-3, gamma=0.99, tau=0.005, buffer_size=100000):
|
|
self.device = device
|
|
self.gamma = gamma
|
|
self.tau = tau # 软更新系数
|
|
self.max_action = max_action
|
|
|
|
# --- 网络补丁：Online 与 Target 网络 ---
|
|
# 1. 创建 Online 网络 (负责被优化器更新)
|
|
self.actor = ContActor(state_dim, action_dim, max_action).to(self.device)
|
|
self.critic = ContQCritic(state_dim, action_dim).to(self.device)
|
|
|
|
# 2. 创建影子 Target 网络 (不参与梯度下降)
|
|
self.actor_target = ContActor(state_dim, action_dim, max_action).to(self.device)
|
|
self.critic_target = ContQCritic(state_dim, action_dim).to(self.device)
|
|
|
|
# 初始化影子网络的权重与主网络一致
|
|
self.actor_target.load_state_dict(self.actor.state_dict())
|
|
self.critic_target.load_state_dict(self.critic.state_dict())
|
|
|
|
# 优化器
|
|
self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)
|
|
self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)
|
|
|
|
# 经验回放池
|
|
self.replay_buffer = ReplayBuffer(buffer_size)
|
|
|
|
def select_action(self, state):
|
|
with torch.no_grad():
|
|
state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
|
|
# 动作由 Online Actor 输出
|
|
action = self.actor(state_tensor).cpu().data.numpy().flatten()
|
|
|
|
# 探索噪声：为了让行为策略具有探索性，加入高斯噪声
|
|
noise = np.random.normal(0, 0.1 * self.max_action, size=action.shape)
|
|
action = np.clip(action + noise, -self.max_action, self.max_action)
|
|
return action
|
|
|
|
def update(self, batch_size=64):
|
|
# 如果缓冲池数据不够，不进行更新
|
|
if len(self.replay_buffer) < batch_size:
|
|
return
|
|
|
|
# 1. 从回放池中随机抽取一批打乱的数据 (彻底打破时序相关性)
|
|
state, action, reward, next_state, done = self.replay_buffer.sample(batch_size)
|
|
|
|
# 转换为 Tensor 维度
|
|
state_t = torch.FloatTensor(state).to(self.device)
|
|
action_t = torch.FloatTensor(action).to(self.device)
|
|
reward_t = torch.FloatTensor(reward).unsqueeze(1).to(self.device)
|
|
next_state_t = torch.FloatTensor(next_state).to(self.device)
|
|
done_t = torch.FloatTensor(done).unsqueeze(1).to(self.device)
|
|
|
|
# --- Critic 更新逻辑：计算稳定目标值 ---
|
|
# 1. 使用延迟反馈的影子网络 Target Actor 预测下一个状态的最理想动作
|
|
next_mu_action = self.actor_target(next_state_t)
|
|
# 2. 使用影子网络 Target Critic 评估这个理想动作的 Q 值 (提供平滑参考)
|
|
next_q_value = self.critic_target(next_state_t, next_mu_action.detach())
|
|
|
|
# 计算稳定的 TD 目标
|
|
td_target = reward_t + self.gamma * next_q_value * (1 - done_t)
|
|
|
|
# 当前 Online Critic 的评估值
|
|
current_q_value = self.critic(state_t, action_t)
|
|
|
|
critic_loss = F.mse_loss(current_q_value, td_target.detach())
|
|
|
|
self.critic_optimizer.zero_grad()
|
|
critic_loss.backward()
|
|
self.critic_optimizer.step()
|
|
|
|
# --- Actor 更新逻辑：只更新 Online Actor ---
|
|
mu_action = self.actor(state_t)
|
|
|
|
# Actor 损失：让 Online Critic 给这个动作打分，越大越好
|
|
actor_loss = -self.critic(state_t, mu_action).mean()
|
|
|
|
self.actor_optimizer.zero_grad()
|
|
actor_loss.backward()
|
|
self.actor_optimizer.step()
|
|
|
|
# --- 软更新 (Soft Update) 补丁 ---
|
|
# 影子网络缓慢向主网络靠近，解决目标乱动问题
|
|
self._soft_update(self.actor_target, self.actor)
|
|
self._soft_update(self.critic_target, self.critic)
|
|
|
|
def _soft_update(self, target_model, online_model):
|
|
"""影子网络参数 = tau * 主网络参数 + (1 - tau) * 影子网络参数"""
|
|
for target_param, online_param in zip(target_model.parameters(), online_model.parameters()):
|
|
target_param.data.copy_(
|
|
target_param.data * (1.0 - self.tau) + online_param.data * self.tau
|
|
)
|
|
|
|
def save(self, path):
|
|
torch.save({
|
|
'actor': self.actor.state_dict(),
|
|
'critic': self.critic.state_dict(),
|
|
'actor_target': self.actor_target.state_dict(),
|
|
'critic_target': self.critic_target.state_dict(),
|
|
}, path)
|
|
|
|
def load(self, path):
|
|
checkpoint = torch.load(path, map_location=self.device)
|
|
self.actor.load_state_dict(checkpoint['actor'])
|
|
self.critic.load_state_dict(checkpoint['critic'])
|
|
self.actor_target.load_state_dict(checkpoint['actor_target'])
|
|
self.critic_target.load_state_dict(checkpoint['critic_target'])
|