| | |
| | import torch |
| | import torch.nn as nn |
| | import random |
| | import numpy as np |
| | from torch.distributions import Normal |
| | from torch.amp import autocast |
| | from torch.cuda.amp import GradScaler |
| |
|
| | |
| | if torch.cuda.is_available(): |
| | device = torch.device("cuda") |
| | print("Using CUDA (NVIDIA GPU)") |
| | else: |
| | device = torch.device("cpu") |
| | print("Using CPU") |
| |
|
| | def set_global_seed(seed: int): |
| | random.seed(seed) |
| | np.random.seed(seed) |
| | torch.manual_seed(seed) |
| |
|
| | if torch.cuda.is_available(): |
| | torch.cuda.manual_seed_all(seed) |
| | torch.backends.cudnn.deterministic = False |
| | torch.backends.cudnn.benchmark = True |
| |
|
| | SEED = 42 |
| | set_global_seed(SEED) |
| |
|
| | class MLP(nn.Module): |
| | def __init__(self, input_dim, hidden_dims, output_dim): |
| | super().__init__() |
| | layers = [] |
| | last_dim = input_dim |
| | for h in hidden_dims: |
| | layers += [nn.Linear(last_dim, h), nn.ReLU()] |
| | last_dim = h |
| | layers.append(nn.Linear(last_dim, output_dim)) |
| | self.net = nn.Sequential(*layers) |
| |
|
| | def forward(self, x): |
| | return self.net(x) |
| |
|
| | class Actor(nn.Module): |
| | def __init__(self, obs_dim, act_dim, hidden=(64,64)): |
| | super().__init__() |
| | self.net = MLP(obs_dim, hidden, act_dim) |
| | self.log_std = nn.Parameter(torch.zeros(act_dim)) |
| |
|
| | def forward(self, x): |
| | mean = self.net(x) |
| | std = torch.exp(self.log_std) |
| | return mean, std |
| |
|
| | class Critic(nn.Module): |
| | def __init__(self, state_dim, hidden=(128,128)): |
| | super().__init__() |
| | self.net = MLP(state_dim, hidden, 1) |
| |
|
| | def forward(self, x): |
| | return self.net(x).squeeze(-1) |
| |
|
| | class MeanField: |
| | def __init__( |
| | self, |
| | n_agents, |
| | local_dim, |
| | global_dim, |
| | act_dim, |
| | lr=3e-4, |
| | gamma=0.99, |
| | lam=0.95, |
| | clip_eps=0.2, |
| | k_epochs=10, |
| | batch_size=1024, |
| | episode_len=96 |
| | ): |
| | self.n_agents = n_agents |
| | self.local_dim = local_dim |
| | self.global_dim = global_dim |
| | self.act_dim = act_dim |
| | self.gamma = gamma |
| | self.lam = lam |
| | self.clip_eps = clip_eps |
| | self.k_epochs = k_epochs |
| | self.batch_size = batch_size |
| | self.episode_len = episode_len |
| |
|
| | self.actor = Actor(local_dim + global_dim, act_dim).to(device) |
| | self.critic = Critic(global_dim).to(device) |
| |
|
| | self.opt_a = torch.optim.Adam(self.actor.parameters(), lr=lr) |
| | self.opt_c = torch.optim.Adam(self.critic.parameters(), lr=lr) |
| |
|
| | print("MeanField CUDA AMP is disabled for stability.") |
| | |
| | self.init_buffer() |
| |
|
| | def init_buffer(self): |
| | self.ls_buf = np.zeros((self.episode_len, self.n_agents, self.local_dim), dtype=np.float32) |
| | self.gs_buf = np.zeros((self.episode_len, self.global_dim), dtype=np.float32) |
| | self.ac_buf = np.zeros((self.episode_len, self.n_agents, self.act_dim), dtype=np.float32) |
| | self.lp_buf = np.zeros((self.episode_len, self.n_agents), dtype=np.float32) |
| | self.rw_buf = np.zeros((self.episode_len, self.n_agents), dtype=np.float32) |
| | self.done_buf = np.zeros((self.episode_len, self.n_agents), dtype=np.float32) |
| | self.next_gs_buf = np.zeros((self.episode_len, self.global_dim), dtype=np.float32) |
| | self.step_idx = 0 |
| |
|
| | @torch.no_grad() |
| | def select_action(self, local_obs, global_obs): |
| | l = torch.from_numpy(local_obs).float().to(device) |
| | g = torch.from_numpy(global_obs).float().to(device).unsqueeze(0).expand(self.n_agents, -1) |
| | input_x = torch.cat([l, g], dim=-1) |
| | mean, std = self.actor(input_x) |
| | dist = Normal(mean, std) |
| | a = dist.sample() |
| | return a.cpu().numpy(), dist.log_prob(a).sum(-1).cpu().numpy() |
| |
|
| | def store(self, local_obs, global_obs, action, logp, reward, done, next_global_obs): |
| | if self.step_idx < self.episode_len: |
| | self.ls_buf[self.step_idx] = local_obs |
| | self.gs_buf[self.step_idx] = global_obs |
| | self.ac_buf[self.step_idx] = action |
| | self.lp_buf[self.step_idx] = logp |
| | self.rw_buf[self.step_idx] = reward |
| | self.done_buf[self.step_idx] = done |
| | self.next_gs_buf[self.step_idx] = next_global_obs |
| | self.step_idx += 1 |
| |
|
| | def compute_gae(self, T, vals): |
| | """ |
| | Computes Generalized Advantage Estimation (GAE). |
| | """ |
| | N = self.n_agents |
| | adv_buf = np.zeros_like(self.rw_buf[:T]) |
| |
|
| |
|
| | if not self.done_buf[T-1].all(): |
| | with torch.no_grad(): |
| | v_last = self.critic( |
| | torch.from_numpy(self.next_gs_buf[T-1]).float().to(device) |
| | ).cpu().numpy() |
| | else: |
| | v_last = 0.0 |
| | vals_agent = vals.unsqueeze(1).expand(-1, N).cpu().numpy() |
| | rewards = self.rw_buf[:T] |
| | masks = 1.0 - self.done_buf[:T] |
| | gae = 0 |
| | for t in reversed(range(T)): |
| | v_next = vals_agent[t+1] if t < T - 1 else v_last |
| | delta = rewards[t] + self.gamma * v_next * masks[t] - vals_agent[t] |
| | adv_buf[t] = gae = delta + self.gamma * self.lam * masks[t] * gae |
| | ret_buf = adv_buf + vals_agent |
| | adv_flat = torch.from_numpy(adv_buf.flatten()).float().to(device) |
| | ret_flat = torch.from_numpy(ret_buf.flatten()).float().to(device) |
| | return adv_flat, ret_flat |
| |
|
| | def update(self): |
| | T = self.step_idx |
| | if T == 0: return |
| |
|
| | gs_tensor = torch.from_numpy(self.gs_buf[:T]).float().to(device) |
| | ls_tensor = torch.from_numpy(self.ls_buf[:T]).float().to(device).view(T * self.n_agents, -1) |
| | ac_tensor = torch.from_numpy(self.ac_buf[:T]).float().to(device).view(T * self.n_agents, -1) |
| | lp_tensor = torch.from_numpy(self.lp_buf[:T]).float().to(device).view(-1) |
| | |
| | with torch.no_grad(): |
| | vals = self.critic(gs_tensor) |
| |
|
| | adv_flat, ret_flat = self.compute_gae(T, vals) |
| | adv_flat = (adv_flat - adv_flat.mean()) / (adv_flat.std() + 1e-8) |
| |
|
| | gs_for_batch = gs_tensor.unsqueeze(1).expand(-1, self.n_agents, -1).reshape(T * self.n_agents, self.global_dim) |
| |
|
| | dataset = torch.utils.data.TensorDataset(ls_tensor, gs_for_batch, ac_tensor, lp_tensor, adv_flat, ret_flat) |
| | gen = torch.Generator() |
| | gen.manual_seed(SEED) |
| | loader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True, generator=gen) |
| |
|
| | for _ in range(self.k_epochs): |
| | for b_ls, b_gs, b_ac, b_lp, b_adv, b_ret in loader: |
| | input_a = torch.cat([b_ls, b_gs], dim=-1) |
| | mean, std = self.actor(input_a) |
| | dist = Normal(mean, std) |
| |
|
| | entropy = dist.entropy().mean() |
| |
|
| | lp_new = dist.log_prob(b_ac).sum(-1) |
| | ratio = torch.exp(lp_new - b_lp) |
| | surr1 = ratio * b_adv |
| | surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * b_adv |
| |
|
| | actor_loss = -torch.min(surr1, surr2).mean() - 0.01 * entropy |
| |
|
| | self.opt_a.zero_grad() |
| | actor_loss.backward() |
| | nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=0.5) |
| | self.opt_a.step() |
| |
|
| |
|
| | val_pred = self.critic(b_gs) |
| | critic_loss = nn.MSELoss()(val_pred, b_ret) |
| |
|
| | self.opt_c.zero_grad() |
| | critic_loss.backward() |
| | nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=0.5) |
| | self.opt_c.step() |
| |
|
| | self.step_idx = 0 |
| | |
| | def save(self, path): |
| | torch.save({'actor': self.actor.state_dict(), |
| | 'critic': self.critic.state_dict()}, path) |
| |
|
| | def load(self, path): |
| | data = torch.load(path, map_location=device) |
| | self.actor.load_state_dict(data['actor']) |
| | self.critic.load_state_dict(data['critic']) |
