MoFlow/trainer/grpo_trainer.py

"""
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')