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