From 5471098e77786d9272612c858120a72932ec4059 Mon Sep 17 00:00:00 2001 From: Maple-YZ Date: Wed, 1 Apr 2026 01:50:56 +0800 Subject: [PATCH] =?UTF-8?q?=E6=89=8B=E5=8A=A8train=E4=BB=A5=E9=81=BF?= =?UTF-8?q?=E5=85=8Dlogger=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- td3_agent.py | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/td3_agent.py b/td3_agent.py index cd46c32..5c05779 100644 --- a/td3_agent.py +++ b/td3_agent.py @@ -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):