手动train以避免logger问题

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):