| from typing import Tuple | |
| import copy | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from .config import CFG | |
| class Actor(nn.Module): | |
| def __init__(self, state_dim: int, action_dim: int, max_action: float = 1.0, hidden: int = 256): | |
| super().__init__() | |
| self.max_action = max_action | |
| self.net = nn.Sequential( | |
| nn.Linear(state_dim, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, action_dim), | |
| nn.Tanh() | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.max_action * self.net(x) | |
| class TwinCritic(nn.Module): | |
| def __init__(self, state_dim: int, action_dim: int, hidden: int = 256): | |
| super().__init__() | |
| self.q1 = nn.Sequential( | |
| nn.Linear(state_dim + action_dim, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, 1) | |
| ) | |
| self.q2 = nn.Sequential( | |
| nn.Linear(state_dim + action_dim, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, hidden), | |
| nn.ReLU(), | |
| nn.Linear(hidden, 1) | |
| ) | |
| def forward(self, state: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| sa = torch.cat([state, action], dim=-1) | |
| return self.q1(sa), self.q2(sa) | |
| class TD3Agent: | |
| def __init__(self, state_dim: int, action_dim: int, max_action: float = 1.0, device: str = "cpu"): | |
| self.device = device | |
| self.actor = Actor(state_dim, action_dim, max_action).to(self.device) | |
| self.actor_target = copy.deepcopy(self.actor).to(self.device) | |
| self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=3e-4) | |
| self.critic = TwinCritic(state_dim, action_dim).to(self.device) | |
| self.critic_target = copy.deepcopy(self.critic).to(self.device) | |
| self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=3e-4) | |
| self.max_action = max_action | |
| self.policy_noise = CFG.td3_policy_noise | |
| self.noise_clip = CFG.td3_noise_clip | |
| self.policy_delay = CFG.td3_policy_delay | |
| self.tau = CFG.td3_tau | |
| self.gamma = CFG.td3_gamma | |
| self.total_it = 0 | |
| def select_action(self, state: np.ndarray, explore: bool = True) -> np.ndarray: | |
| state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| action = self.actor(state_t).cpu().numpy().flatten() | |
| if explore: | |
| action = (action + np.random.normal(0, self.max_action * 0.1, size=action.shape)).clip(-self.max_action, self.max_action) | |
| return action | |
| def train_step(self, replay_buffer, batch_size: int = 256): | |
| self.total_it += 1 | |
| state, action, reward, next_state, done = replay_buffer.sample(batch_size) | |
| state = torch.FloatTensor(state).to(self.device) | |
| action = torch.FloatTensor(action).to(self.device) | |
| reward = torch.FloatTensor(reward).unsqueeze(1).to(self.device) | |
| next_state = torch.FloatTensor(next_state).to(self.device) | |
| done = torch.FloatTensor(done).unsqueeze(1).to(self.device) | |
| with torch.no_grad(): | |
| noise = (torch.randn_like(action) * self.policy_noise).clamp(-self.noise_clip, self.noise_clip) | |
| next_action = (self.actor_target(next_state) + noise).clamp(-self.max_action, self.max_action) | |
| q1_next, q2_next = self.critic_target(next_state, next_action) | |
| q_next = torch.min(q1_next, q2_next) | |
| target_q = reward + (1 - done) * self.gamma * q_next | |
| q1, q2 = self.critic(state, action) | |
| critic_loss = F.mse_loss(q1, target_q) + F.mse_loss(q2, target_q) | |
| self.critic_optimizer.zero_grad() | |
| critic_loss.backward() | |
| self.critic_optimizer.step() | |
| if self.total_it % self.policy_delay == 0: | |
| actor_loss = -self.critic.q1(torch.cat([state, self.actor(state)], dim=-1)).mean() if hasattr(self.critic, "q1") else -self.critic(state, self.actor(state))[0].mean() | |
| self.actor_optimizer.zero_grad() | |
| actor_loss.backward() | |
| self.actor_optimizer.step() | |
| # soft update | |
| for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): | |
| target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) | |
| for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): | |
| target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) | |
| return {"critic_loss": critic_loss.item()} | |
Xet Storage Details
- Size:
- 4.75 kB
- Xet hash:
- 27aca9866af0abe210923cf11ee3a498ebcc34dbb3b3d90813d488b2f750338b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.