""" GRPOTrainer — Flow-GRPO fine-tuning of MoFlow + SRA on NBA. Subclasses the SFT `Trainer` to reuse its eval / metric / checkpoint machinery (`eval_dataloader`, `compute_*`, `sample_from_denoising_model`, `save_ckpt`), and replaces `train()` with the GRPO loop: rollout (stochastic SDE sampler, no_grad) -> combined reward (accuracy + coherence) -> group-relative advantage over the K modes -> E inner PPO-clip epochs (recompute log-probs w/ grad) + KL to frozen ref -> opt.step, EMA update The model is kept in eval() mode throughout so the SRA backbone builds its pass-1 graph from y_0_prev (self-conditioning), never from GT (teacher forcing). Evaluation uses the inherited *deterministic* sampler (sde_noise -> 0). """ import torch from einops import rearrange from trainer.denoising_model_trainers import Trainer from utils.normalization import unnormalize_min_max, unnormalize_sqrt from grpo.rewards import compute_reward_agentwise, group_advantage class GRPOTrainer(Trainer): def __init__(self, cfg, denoiser, ref_model, train_loader, test_loader, logger=None, tb_log=None, grpo=None): super().__init__(cfg, denoiser, train_loader, test_loader, logger=logger, tb_log=tb_log) g = grpo or {} self.grpo_iters = int(g.get('iters', 2000)) self.rl_steps = g.get('rl_steps', None) # None -> cfg.sampling_steps self.sde_noise = float(g.get('sde_noise', 0.1)) self.eta_min = float(g.get('eta_min', 1e-3)) self.inner_epochs = int(g.get('inner_epochs', 2)) self.clip_eps = float(g.get('clip_eps', 0.2)) self.kl_beta = float(g.get('kl_beta', 0.1)) self.logratio_clip = float(g.get('logratio_clip', 10.0)) self.eval_every = int(g.get('eval_every', 100)) self.adv_eps = float(g.get('adv_eps', 1e-4)) self.reward_kwargs = dict( w_ade=float(g.get('w_ade', 1.0)), w_fde=float(g.get('w_fde', 1.0)), w_jade=float(g.get('w_jade', 0.5)), w_jfde=float(g.get('w_jfde', 0.5)), w_col=float(g.get('w_col', 1.0)), w_kin=float(g.get('w_kin', 0.0)), d_min=float(g.get('d_min', 0.4)), a_max=float(g.get('a_max', 1.0)), ball_idx=g.get('ball_idx', None), ) # eval/checkpoint config self.best_score = float('inf') # best ADE+FDE+JADE+JFDE (accuracy runs) self.best_coll = float('inf') self.ade_tol = float(g.get('ade_tol', 0.71)) # save best only if ADE_min(4s) <= this self.coll_thresholds = (0.2, 0.3, 0.4) self.coll_ball_idx = g.get('ball_idx', 10) self.coll_eval_batches = int(g.get('coll_eval_batches', 5)) # frozen reference policy (KL anchor) self.ref_model = ref_model.to(self.device) self.ref_model.eval() for p in self.ref_model.parameters(): p.requires_grad_(False) self.model_grpo = self.accelerator.unwrap_model(self.denoiser) # ------------------------------------------------------------------ def _pred_to_metric(self, pred_norm): """[B,K,A,Do] normalized -> [B,K,A,T,2] metric scale.""" B, K, A, _ = pred_norm.shape pred = pred_norm.view(B, K, A, self.cfg.future_frames, 2) if self.cfg.get('data_norm', None) == 'min_max': return unnormalize_min_max(pred, self.cfg.fut_traj_min, self.cfg.fut_traj_max, -1, 1) elif self.cfg.get('data_norm', None) == 'sqrt': return unnormalize_sqrt(pred, self.sqrt_a_, self.sqrt_b_) return pred # ------------------------------------------------------------------ @torch.no_grad() def eval_collision(self): """Deterministic player-player collision rate on the val set (a subset of batches for speed). Returns {thresh: rate over all (scene,mode)}.""" self.denoiser.eval() A, T = self.cfg.agents, self.cfg.future_frames pm = ~torch.eye(A, dtype=torch.bool, device=self.device) if self.coll_ball_idx is not None: pm[self.coll_ball_idx, :] = False pm[:, self.coll_ball_idx] = False counts = {th: [0.0, 0.0] for th in self.coll_thresholds} for i, data in enumerate(self.val_loader): if i >= self.coll_eval_batches: break data = {k: v.to(self.device) for k, v in data.items()} B = int(data['batch_size']) pred, _, _, _, _ = self.sample_from_denoising_model(data) # [B*A,K,T,2] metric K = pred.shape[1] pred = pred.reshape(B, A, K, T, 2).permute(0, 2, 1, 3, 4) # [B,K,A,T,2] init = data['past_traj_original_scale'][:, :, -1, 0:2] abs_pred = pred + init[:, None, :, None, :] d = (abs_pred.unsqueeze(3) - abs_pred.unsqueeze(2)).norm(dim=-1) # [B,K,A,A,T] cp = d.min(dim=-1).values.masked_fill(~pm, 1e9).reshape(B, K, -1).min(-1).values # [B,K] for th in self.coll_thresholds: counts[th][0] += (cp < th).float().sum().item() counts[th][1] += cp.numel() return {th: counts[th][0] / max(counts[th][1], 1) for th in self.coll_thresholds} def train(self): self.logger.info('GRPO training start') model, ref = self.model_grpo, self.ref_model model.eval() G = self.cfg.denoising_head_preds for it in range(self.grpo_iters): self.step = it data = {k: v.to(self.device) for k, v in next(self.dl).items()} # ---------------- rollout ---------------- with torch.no_grad(): roll = model.rollout(data, num_trajs=G, sde_noise=self.sde_noise, eta_min=self.eta_min, sampling_steps=self.rl_steps) # ---------------- reward + advantage (per-agent) ---------------- pred_metric = self._pred_to_metric(roll['final_action']) # [B,K,A,T,2] gt_metric = data['fut_traj_original_scale'] # [B,A,T,2] init_pos = data['past_traj_original_scale'][:, :, -1, 0:2] # [B,A,2] reward, info = compute_reward_agentwise(pred_metric, gt_metric, init_pos, **self.reward_kwargs) adv = group_advantage(reward, eps=self.adv_eps) # [B,K,A] # ---------------- PPO-clip update ---------------- num_steps = len(roll['steps']) last_stats = {} for _ in range(self.inner_epochs): self.opt.zero_grad() pg_sum = kl_sum = ratio_sum = clipfrac_sum = 0.0 for step in roll['steps']: logp_new, kl = model.recompute_logp_kl( step, data, ref, self.sde_noise, self.eta_min) logratio = (logp_new - step['logp_old']).clamp( -self.logratio_clip, self.logratio_clip) ratio = logratio.exp() unclipped = ratio * adv clipped = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * adv pg = -torch.min(unclipped, clipped).mean() kl_loss = kl.mean() loss = (pg + self.kl_beta * kl_loss) / num_steps self.accelerator.backward(loss) pg_sum += pg.item() / num_steps kl_sum += kl_loss.item() / num_steps ratio_sum += ratio.mean().item() / num_steps clipfrac_sum += ((ratio - 1.0).abs() > self.clip_eps).float().mean().item() / num_steps self.accelerator.clip_grad_norm_( self.denoiser.parameters(), self.cfg.OPTIMIZATION.GRAD_NORM_CLIP) self.opt.step() if self.accelerator.is_main_process: self.ema.update() last_stats = dict(pg=pg_sum, kl=kl_sum, ratio=ratio_sum, clipfrac=clipfrac_sum) # ---------------- logging ---------------- if it % 10 == 0: self.logger.info( f'[GRPO {it}/{self.grpo_iters}] R={reward.mean().item():.4f} ' f'ADE*={info["ade_bestk"].item():.4f} JADE*={info["jade_bestk"].item():.4f} ' f'coll={info["collision"].mean().item():.4f} collrate={info["collision_rate"].item():.3f} ' f'| pg={last_stats["pg"]:.4f} kl={last_stats["kl"]:.5f} ' f'ratio={last_stats["ratio"]:.3f} clipfrac={last_stats["clipfrac"]:.3f}') if self.tb_log is not None: self.tb_log.add_scalar('grpo/reward', reward.mean().item(), it) self.tb_log.add_scalar('grpo/ade_bestk', info['ade_bestk'].item(), it) self.tb_log.add_scalar('grpo/jade_bestk', info['jade_bestk'].item(), it) self.tb_log.add_scalar('grpo/collision', info['collision'].mean().item(), it) self.tb_log.add_scalar('grpo/collision_rate', info['collision_rate'].item(), it) self.tb_log.add_scalar('grpo/pg_loss', last_stats['pg'], it) self.tb_log.add_scalar('grpo/kl', last_stats['kl'], it) self.tb_log.add_scalar('grpo/ratio', last_stats['ratio'], it) self.tb_log.add_scalar('grpo/clipfrac', last_stats['clipfrac'], it) self.tb_log.add_scalar('grpo/lr', self.opt.param_groups[0]['lr'], it) # ---------------- periodic eval (deterministic sampler) ---------------- if (it + 1) % self.eval_every == 0 and self.accelerator.is_main_process: _, perf, n = self.eval_dataloader(testing_mode=False) pj = self.last_perf_joint ade4 = perf['ADE_min'][3] / n fde4 = perf['FDE_min'][3] / n jade4 = pj['JADE_min'][3] / n jfde4 = pj['JFDE_min'][3] / n score = ade4 + fde4 + jade4 + jfde4 # combined accuracy (lower=better) self.logger.info( f'[GRPO eval @ {it}] ADE={ade4:.5f} FDE={fde4:.5f} ' f'JADE={jade4:.5f} JFDE={jfde4:.5f} | sum={score:.5f} (best={self.best_score:.5f})') if self.tb_log is not None: self.tb_log.add_scalar('eval/ADE', ade4, it) self.tb_log.add_scalar('eval/FDE', fde4, it) self.tb_log.add_scalar('eval/JADE', jade4, it) self.tb_log.add_scalar('eval/JFDE', jfde4, it) if score < self.best_score: self.best_score = score self.logger.info(f' new best sum={score:.5f} ' f'(ADE={ade4:.5f} JADE={jade4:.5f}) -> checkpoint_best') self.save_ckpt('checkpoint_best') self.save_ckpt('checkpoint_last') model.eval() self.save_ckpt('checkpoint_last') self.logger.info('GRPO training complete')