276 lines
11 KiB
Python
276 lines
11 KiB
Python
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch_geometric.data import Batch
|
||
from torch_scatter import scatter_add
|
||
|
||
|
||
class RolloutBuffer:
|
||
|
||
def __init__(self, buffer_size: int,
|
||
gae_lambda: float,
|
||
discount_factor: float,
|
||
device=None,
|
||
):
|
||
self.buffer_size = buffer_size
|
||
self.gae_lambda = gae_lambda
|
||
self.discount_factor = discount_factor
|
||
self.device = device
|
||
self.reset()
|
||
|
||
def reset(self):
|
||
self.observations = []
|
||
self.actions = []
|
||
self.log_probs = []
|
||
self.rewards = [] # per-agent rewards (list of tensors, varying shapes)
|
||
self.values = [] # per-agent values (list of tensors, varying shapes)
|
||
self.dones = []
|
||
self.agent_mappings = [] # mapping from new → old agent indices per step
|
||
self.pos = 0
|
||
|
||
def add(
|
||
self, observation, actions, reward, done, value, log_probs,
|
||
agent_mapping=None,
|
||
):
|
||
dev = self.device
|
||
self.observations.append(observation.to(dev))
|
||
self.actions.append(actions.to(dev))
|
||
self.log_probs.append(log_probs.to(dev))
|
||
self.rewards.append(torch.as_tensor(reward, dtype=torch.float32, device=dev).flatten())
|
||
self.values.append(value.flatten().to(dev))
|
||
self.dones.append(float(done))
|
||
self.agent_mappings.append(
|
||
torch.as_tensor(agent_mapping, dtype=torch.long, device=dev).flatten()
|
||
)
|
||
self.pos += 1
|
||
|
||
def compute_returns_and_advantage(self, last_value):
|
||
"""Single-path GAE: potential-shaped per-agent reward with scatter_add for mesh refinement."""
|
||
last_value = last_value.to(self.device).flatten()
|
||
n = self.buffer_size
|
||
|
||
dones = torch.as_tensor(self.dones, device=self.device)
|
||
|
||
# ---- 0. Normalize rewards to unit scale ----
|
||
all_rews = torch.cat([r.flatten() for r in self.rewards])
|
||
rew_mean = all_rews.mean()
|
||
rew_std = all_rews.std()
|
||
if rew_std > 1e-8:
|
||
self.rewards = [(r - rew_mean) / rew_std for r in self.rewards]
|
||
|
||
# ---- 1. Per-agent GAE (scatter_add for mesh refinement) ----
|
||
advantages = [None] * n
|
||
deltas = []
|
||
next_values = self.values[1:] + [last_value]
|
||
|
||
for step in range(n):
|
||
if dones[step]:
|
||
next_val = self.values[step]
|
||
else:
|
||
next_val = scatter_add(next_values[step], self.agent_mappings[step], dim=0)
|
||
delta = self.rewards[step] + (0 if dones[step] else self.discount_factor * next_val) - self.values[step]
|
||
deltas.append(delta)
|
||
|
||
last_gae = torch.zeros_like(self.agent_mappings[-1], dtype=torch.float32, device=self.device)
|
||
for step in reversed(range(n)):
|
||
if dones[step]:
|
||
last_gae = deltas[step]
|
||
else:
|
||
last_gae = deltas[step] + self.discount_factor * self.gae_lambda * scatter_add(last_gae, self.agent_mappings[step], dim=0)
|
||
advantages[step] = last_gae
|
||
|
||
self.returns = [adv + val for adv, val in zip(advantages, self.values)]
|
||
|
||
# ---- 2. Normalize advantages (per-batch, zero-mean unit-std) ----
|
||
all_advs = torch.cat([a.flatten() for a in advantages])
|
||
adv_mean = all_advs.mean()
|
||
adv_std = all_advs.std()
|
||
if adv_std > 1e-8:
|
||
advantages = [(a - adv_mean) / adv_std for a in advantages]
|
||
# NOTE: returns and values keep their original scale — no unit-scale normalization,
|
||
# so the value network sees a stable regression target across iterations.
|
||
|
||
self.advantages = [ret - val for ret, val in zip(self.returns, self.values)]
|
||
|
||
def get(self, batch_size: int):
|
||
"""Yield random minibatches from the buffer."""
|
||
indices = np.random.permutation(self.buffer_size)
|
||
start = 0
|
||
while start < self.buffer_size:
|
||
batch_idx = indices[start : start + batch_size]
|
||
start += batch_size
|
||
|
||
obs_batch = Batch.from_data_list([self.observations[i] for i in batch_idx])
|
||
acts = torch.cat([self.actions[i] for i in batch_idx], dim=0)
|
||
lps = torch.cat([self.log_probs[i].flatten() for i in batch_idx], dim=0)
|
||
vals = torch.cat([self.values[i].flatten() for i in batch_idx], dim=0)
|
||
advs = torch.cat([self.advantages[i].flatten() for i in batch_idx], dim=0)
|
||
rets = torch.cat([self.returns[i].flatten() for i in batch_idx], dim=0)
|
||
|
||
obs_batch, acts, lps, vals, advs, rets = (
|
||
x.to(self.device) for x in (obs_batch, acts, lps, vals, advs, rets)
|
||
)
|
||
yield obs_batch, acts, lps, vals, advs, rets
|
||
|
||
@property
|
||
def full(self):
|
||
return self.pos >= self.buffer_size
|
||
|
||
@property
|
||
def explained_variance(self):
|
||
all_vals = torch.cat([v.flatten() for v in self.values])
|
||
all_rets = torch.cat([r.flatten() for r in self.returns])
|
||
var_ret = torch.var(all_rets)
|
||
if var_ret < 1e-12:
|
||
return 0.0
|
||
return float(1.0 - torch.var(all_rets - all_vals) / var_ret)
|
||
|
||
|
||
# ── PPO losses ────────────────────────────────────────────
|
||
def policy_loss(advantages: torch.Tensor, ratio: torch.Tensor, clip_range: float) -> torch.Tensor:
|
||
"""Clipped PPO policy loss."""
|
||
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
||
loss1 = advantages * ratio
|
||
loss2 = advantages * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
|
||
return -torch.min(loss1, loss2).mean()
|
||
|
||
|
||
def value_loss(
|
||
returns: torch.Tensor, values: torch.Tensor,
|
||
old_values: torch.Tensor, clip_range: float,
|
||
) -> torch.Tensor:
|
||
"""Clipped value function loss."""
|
||
vf_loss = F.mse_loss(returns, values)
|
||
if clip_range > 0:
|
||
v_clipped = old_values + (values - old_values).clamp(-clip_range, clip_range)
|
||
vf_loss = torch.max(vf_loss, F.mse_loss(returns, v_clipped))
|
||
return vf_loss
|
||
|
||
|
||
def entropy_loss(entropy) -> torch.Tensor:
|
||
"""Entropy bonus for exploration."""
|
||
return -torch.mean(entropy)
|
||
|
||
|
||
class PPOTrainer:
|
||
|
||
def __init__(self, actor_critic, environment, config: dict, device=None):
|
||
self.policy = actor_critic
|
||
self.env = environment
|
||
self.device = device
|
||
|
||
ppo_cfg = config.get("ppo", {})
|
||
self.num_rollout_steps = ppo_cfg.get("num_rollout_steps", 256)
|
||
self.epochs_per_iteration = ppo_cfg.get("epochs_per_iteration", 5)
|
||
self.batch_size = config.get("batch_size", 32)
|
||
self.clip_range = ppo_cfg.get("clip_range", 0.2)
|
||
self.max_grad_norm = ppo_cfg.get("max_grad_norm", 0.5)
|
||
self.entropy_coef = ppo_cfg.get("entropy_coefficient", 0.0)
|
||
self.vf_coef = ppo_cfg.get("value_function_coefficient", 0.5)
|
||
self.vf_clip_range = ppo_cfg.get("value_function_clip_range", 0.2)
|
||
self.gae_lambda = ppo_cfg.get("gae_lambda", 0.95)
|
||
self.discount_factor = config.get("discount_factor", 1.0)
|
||
|
||
self.buffer = RolloutBuffer(
|
||
buffer_size=self.num_rollout_steps,
|
||
gae_lambda=self.gae_lambda,
|
||
discount_factor=self.discount_factor,
|
||
device=device,
|
||
)
|
||
|
||
def collect_rollouts(self):
|
||
self.policy.eval()
|
||
self.buffer.reset()
|
||
obs = self.env.reset()
|
||
step_rewards, step_num_agents = [], []
|
||
_rho_keys = ("rho_int_mean", "rho_jump_mean", "rho_sbc_mean",
|
||
"w_rho_int", "w_rho_jump", "w_rho_sbc")
|
||
rho_accum = {k: 0.0 for k in _rho_keys}
|
||
diag_keys = ("neg_action_ratio", "eligible_ratio", "selected_count",
|
||
"dorfler_tail_ratio", "dorfler_floor_active")
|
||
diag_accum = {k: 0.0 for k in diag_keys}
|
||
diag_steps = 0
|
||
|
||
for _ in range(self.num_rollout_steps):
|
||
with torch.no_grad():
|
||
actions, values, log_probs = self.policy(
|
||
Batch.from_data_list([obs]), deterministic=False
|
||
)
|
||
values = values.flatten()
|
||
next_obs, reward, done, info = self.env.step(actions.cpu().numpy())
|
||
step_rewards.append(float(np.sum(reward)))
|
||
step_num_agents.append(int(len(reward)))
|
||
for k in _rho_keys:
|
||
if k in info:
|
||
rho_accum[k] += float(info[k])
|
||
for k in diag_keys:
|
||
if k in info:
|
||
diag_accum[k] += float(info[k])
|
||
diag_steps += 1
|
||
|
||
self.buffer.add(
|
||
observation=obs, actions=actions, reward=reward,
|
||
done=float(done), value=values, log_probs=log_probs,
|
||
agent_mapping=self.env.agent_mapping,
|
||
)
|
||
obs = self.env.reset() if done else next_obs
|
||
|
||
with torch.no_grad():
|
||
_, last_value, _ = self.policy(Batch.from_data_list([obs]), deterministic=True)
|
||
last_value = last_value.squeeze(-1).flatten()
|
||
self.buffer.compute_returns_and_advantage(last_value)
|
||
|
||
n = max(1, self.num_rollout_steps)
|
||
metrics = {
|
||
"num_agents": step_num_agents[-1], "reward": step_rewards[-1],
|
||
"avg_agents": np.mean(step_num_agents),
|
||
"avg_reward": np.mean(step_rewards),
|
||
"min_reward": np.min(step_rewards),
|
||
"max_reward": np.max(step_rewards),
|
||
"sum_reward": np.sum(step_rewards),
|
||
}
|
||
# rho diagnostics for weight calibration (averaged over rollout)
|
||
for k in _rho_keys:
|
||
metrics[k] = rho_accum[k] / n
|
||
# score-based refinement diagnostics
|
||
n_diag = max(1, diag_steps)
|
||
for k in diag_keys:
|
||
metrics[k] = diag_accum[k] / n_diag
|
||
return metrics
|
||
|
||
def train_step(self):
|
||
self.policy.train()
|
||
total_losses = []
|
||
for _ in range(self.epochs_per_iteration):
|
||
for obs_batch, acts, old_lp, old_vals, advs, rets in self.buffer.get(self.batch_size):
|
||
values, log_probs, entropy = self.policy.evaluate_actions(obs_batch, acts)
|
||
values = values.squeeze(-1)
|
||
|
||
ratio = torch.exp(log_probs - old_lp)
|
||
|
||
pl = policy_loss(advs, ratio, self.clip_range)
|
||
vl = self.vf_coef * value_loss(rets, values, old_vals, self.vf_clip_range)
|
||
el = self.entropy_coef * entropy_loss(entropy)
|
||
loss = pl + vl + el
|
||
|
||
self.policy.optimizer.zero_grad()
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
|
||
self.policy.optimizer.step()
|
||
if self.policy.log_std is not None:
|
||
self.policy.log_std.data.clamp_(-3.0, -1.0) # σ ∈ [0.05, 0.37]
|
||
total_losses.append(loss.item())
|
||
|
||
if self.policy.lr_scheduler is not None:
|
||
self.policy.lr_scheduler.step()
|
||
|
||
return {
|
||
"loss": np.mean(total_losses) if total_losses else 0.0,
|
||
"explained_variance": self.buffer.explained_variance,
|
||
}
|
||
|
||
def fit_iteration(self):
|
||
metrics = self.collect_rollouts()
|
||
metrics.update(self.train_step())
|
||
return metrics
|