MapleL
/
ctm-dqn
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

手动train以避免logger问题

This commit is contained in:
Zihan Ye 2026-04-01 01:50:56 +08:00
parent 2cbaa27b8b
commit 5471098e77
1 changed files with 36 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
td3_agent.py
View File

@ -3,6 +3,7 @@ TD3 Agent using Stable-Baselines3
适配 MultiDiscrete 动作空间的 VSL 控制
"""
import numpy as np
import torch
from stable_baselines3 import TD3
from stable_baselines3.common.noise import NormalActionNoise
import gymnasium as gym
@ -108,7 +109,41 @@ class TD3Agent:
if self.model.replay_buffer.size() < self.learning_starts:
return {}
self.model.train(gradient_steps=1)
# 手动更新而不是调用train()
self.model._n_updates += 1
gradient_steps = 1
for _ in range(gradient_steps):
self.model._update_learning_rate(self.model.actor.optimizer)
self.model._update_learning_rate(self.model.critic.optimizer)
replay_data = self.model.replay_buffer.sample(self.model.batch_size)
with torch.no_grad():
noise = replay_data.actions.clone().data.normal_(0, self.model.target_policy_noise)
noise = noise.clamp(-self.model.target_noise_clip, self.model.target_noise_clip)
next_actions = (self.model.actor_target(replay_data.next_observations) + noise).clamp(-1, 1)
next_q_values = torch.cat(self.model.critic_target(replay_data.next_observations, next_actions), dim=1)
next_q_values, _ = torch.min(next_q_values, dim=1, keepdim=True)
target_q_values = replay_data.rewards + (1 - replay_data.dones) * self.model.gamma * next_q_values
current_q_values = self.model.critic(replay_data.observations, replay_data.actions)
critic_loss = sum(torch.nn.functional.mse_loss(current_q, target_q_values) for current_q in current_q_values)
self.model.critic.optimizer.zero_grad()
critic_loss.backward()
self.model.critic.optimizer.step()
if self.model._n_updates % self.model.policy_delay == 0:
actor_loss = -self.model.critic.q1_forward(replay_data.observations, self.model.actor(replay_data.observations)).mean()
self.model.actor.optimizer.zero_grad()
actor_loss.backward()
self.model.actor.optimizer.step()
self.model._polyak_update(self.model.critic.parameters(), self.model.critic_target.parameters(), self.model.tau)
self.model._polyak_update(self.model.actor.parameters(), self.model.actor_target.parameters(), self.model.tau)
return {"actor_loss": 0.0, "critic_loss": 0.0}
def save(self, path: str):