76 lines
3.3 KiB
Python
76 lines
3.3 KiB
Python
import torch
|
||
import torch.optim as optim
|
||
import torch.nn.functional as F
|
||
import numpy as np
|
||
from networks_cont import ContActor, ContQCritic
|
||
|
||
class DPACAgent:
|
||
def __init__(self, state_dim, action_dim, max_action, device, actor_lr=0.001, critic_lr=0.002, gamma=0.99):
|
||
self.device = device
|
||
self.gamma = gamma
|
||
self.max_action = max_action
|
||
|
||
self.actor = ContActor(state_dim, action_dim, max_action).to(self.device)
|
||
self.critic = ContQCritic(state_dim, action_dim).to(self.device)
|
||
|
||
self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)
|
||
self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)
|
||
|
||
def select_action(self, state):
|
||
# 行为策略：在确定性目标策略的基础上加入高斯噪声，用于探索
|
||
with torch.no_grad():
|
||
state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
|
||
action = self.actor(state_tensor).cpu().data.numpy().flatten()
|
||
|
||
# 探索噪声 (行为策略 beta 与目标策略 mu 的区别就在这里)
|
||
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, state, action, reward, next_state, next_action, done):
|
||
state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
|
||
next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0).to(self.device)
|
||
reward_tensor = torch.FloatTensor([reward]).unsqueeze(0).to(self.device)
|
||
# 注意：连续动作直接是浮点数，不需要转为索引
|
||
action_tensor = torch.FloatTensor(action).unsqueeze(0).to(self.device)
|
||
|
||
# --- Critic 更新 (Algorithm 10.4 核心逻辑) ---
|
||
# 1. 目标策略 mu 在下一个状态的理想输出
|
||
next_mu_action = self.actor(next_state_tensor).detach()
|
||
# 2. 评估这个理想动作的 Q 值
|
||
next_q_value = self.critic(next_state_tensor, next_mu_action).detach()
|
||
# 3. 计算 TD 目标
|
||
td_target = reward_tensor + self.gamma * next_q_value * (1 - int(done))
|
||
|
||
# 4. 当前实际采取动作的 Q 值
|
||
current_q_value = self.critic(state_tensor, action_tensor)
|
||
|
||
critic_loss = F.mse_loss(current_q_value, td_target)
|
||
|
||
self.critic_optimizer.zero_grad()
|
||
critic_loss.backward()
|
||
self.critic_optimizer.step()
|
||
|
||
# --- Actor 更新 (链式法则) ---
|
||
# 1. 目标策略当前状态的输出
|
||
mu_action = self.actor(state_tensor)
|
||
|
||
# 2. 拿到这个动作去问 Critic："给我打分"。
|
||
# 为了让 Q 值最大化，我们加负号转化为梯度下降
|
||
actor_loss = -self.critic(state_tensor, mu_action).mean()
|
||
|
||
self.actor_optimizer.zero_grad()
|
||
actor_loss.backward()
|
||
self.actor_optimizer.step()
|
||
|
||
def save(self, path):
|
||
torch.save({
|
||
'actor': self.actor.state_dict(),
|
||
'critic': self.critic.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'])
|