MoFlow/models/flow_matching_grpo.py

"""
FlowMatcherGRPO — stochastic (SDE) sampler exposing per-step Gaussian log-probs
and closed-form transition KL, for GRPO fine-tuning of MoFlow + SRA graph.

Policy transition (Gaussian around the deterministic Euler mean):

    mu_t   = y_t + v_theta(y_t, t) * dt          # existing deterministic update
    eta_t  = sde_noise * (1 - t)                 # exploration std, -> 0 as t -> 1
    y_{t+dt} ~ N(mu_t, eta_t^2 I)

`eta_t -> 0` recovers MoFlow's deterministic ODE sampler exactly, so evaluation
still uses the inherited deterministic `sample()`.  The velocity `v_theta` comes
from the *unchanged* two-pass graph forward (`model_predictions`), so the SRA
graph / uncertainty machinery is untouched — we only wrap the update rule.

The MODEL MUST BE IN eval() MODE during both rollout and update, so the backbone
builds its pass-1 graph from `y_0_prev` (self-conditioning) rather than from the
GT future (teacher forcing, which only happens when self.training is True).
Gradient is controlled by the surrounding torch.no_grad()/enable_grad context.

Log-probs are summed over the agent and feature dims -> one scalar per (scene,
mode) = per group member, which is the granularity GRPO's advantage/ratio use.
"""

import math
import numpy as np
import torch

from models.flow_matching import FlowMatcher


class FlowMatcherGRPO(FlowMatcher):

    # ------------------------------------------------------------------
    # Time grid (euler only)
    # ------------------------------------------------------------------
    def _build_time_grid(self, sampling_steps=None):
        assert self.solver == 'euler', 'GRPO sampler supports the euler solver only'
        n = int(sampling_steps or self.sampling_steps)
        dt = 1.0 / n
        t_ls = (dt * np.arange(n)).tolist()
        dt_ls = (dt * np.ones(n)).tolist()
        return t_ls, dt_ls

    @staticmethod
    def _eta(cur_t, sde_noise, eta_min):
        # exploration std; decays to ~0 near t=1 so final samples stay on-manifold
        return max(sde_noise * (1.0 - float(cur_t)), float(eta_min))

    # ------------------------------------------------------------------
    # One stochastic step (sample, or evaluate log-prob at a given action)
    # ------------------------------------------------------------------
    def policy_step(self, y_t, cur_t, cur_dt, x_data, y_0_prev,
                    sde_noise, eta_min, action=None, deterministic=False):
        """Returns (action, logp[B,K], pred_data, pred_score, mean, eta).

        action is None  -> sample a fresh action from the Gaussian policy.
        action given    -> evaluate logp of that action under the current policy
                           (used by the PPO recompute pass; keep grad enabled).
        """
        B, K, A = y_t.shape[0], y_t.shape[1], y_t.shape[2]
        batched_t = torch.full((B,), float(cur_t), device=y_t.device, dtype=torch.float)
        preds = self.model_predictions(y_t, x_data, batched_t,
                                        flag_print=False, y_0_prev=y_0_prev)
        mean = y_t + preds.pred_vel * cur_dt                       # [B, K, A, Do]
        eta = self._eta(cur_t, sde_noise, eta_min)

        if deterministic or eta <= 0.0:
            out_action = mean if action is None else action
            logp = torch.zeros(B, K, A, device=y_t.device)
            return out_action, logp, preds.pred_data, preds.pred_score, mean, eta

        var = eta * eta
        if action is None:
            action = mean + eta * torch.randn_like(mean)
        # Gaussian log-density; sum over the FEATURE dim only -> per (B, K, A).
        # The diagonal Gaussian factorizes over agents, so per-agent log-probs are
        # valid, enabling per-agent credit assignment for marginal (ADE/FDE) metrics.
        logp = -0.5 * (((action - mean) ** 2) / var + math.log(2.0 * math.pi * var))
        logp = logp.sum(dim=3)                                     # [B, K, A]
        return action, logp, preds.pred_data, preds.pred_score, mean, eta

    # ------------------------------------------------------------------
    # Rollout: sample K trajectories/scene, recording GRPO bookkeeping
    # ------------------------------------------------------------------
    @torch.no_grad()
    def rollout(self, x_data, num_trajs, sde_noise, eta_min=1e-3, sampling_steps=None):
        """Model must be in eval() mode.  Returns per-step states/actions/logp_old
        plus the final data-space prediction (normalized) [B, K, A, Do]."""
        B = int(x_data['batch_size'])
        num_agents = (x_data['agent_mask'].shape[1]
                      if isinstance(x_data, dict) and 'agent_mask' in x_data
                      else self.num_agents)
        device = self.device
        y_t = torch.randn((B, num_trajs, num_agents, self.out_dim), device=device)

        t_ls, dt_ls = self._build_time_grid(sampling_steps)
        steps, y_0_prev = [], None
        for cur_t, cur_dt in zip(t_ls, dt_ls):
            action, logp, pred_data, _, _, eta = self.policy_step(
                y_t, cur_t, cur_dt, x_data, y_0_prev, sde_noise, eta_min)
            steps.append({
                'y_t': y_t.detach(),
                'action': action.detach(),
                'logp_old': logp.detach(),
                't': float(cur_t), 'dt': float(cur_dt), 'eta': float(eta),
                'y_0_prev': (None if y_0_prev is None else y_0_prev.detach()),
            })
            y_t = action
            y_0_prev = pred_data.detach()

        # final_action = the last sample the policy actually emitted (what eval
        # scores); final_pred = the model's data prediction at the last step.
        return {'steps': steps, 'final_action': y_t, 'final_pred': y_0_prev,
                'B': B, 'A': num_agents, 'K': num_trajs}

    # ------------------------------------------------------------------
    # Recompute logp under current theta (grad) + closed-form KL to reference
    # ------------------------------------------------------------------
    def recompute_logp_kl(self, step, x_data, ref_model, sde_noise, eta_min=1e-3):
        """Returns (logp_new[B,K], kl[B,K]).  Call with grad enabled; model in eval()."""
        y_t, action = step['y_t'], step['action']
        cur_t, cur_dt, y_0_prev = step['t'], step['dt'], step['y_0_prev']

        _, logp_new, _, _, mean_new, eta = self.policy_step(
            y_t, cur_t, cur_dt, x_data, y_0_prev, sde_noise, eta_min, action=action)

        with torch.no_grad():
            B = y_t.shape[0]
            batched_t = torch.full((B,), float(cur_t), device=y_t.device, dtype=torch.float)
            ref_preds = ref_model.model_predictions(
                y_t, x_data, batched_t, flag_print=False, y_0_prev=y_0_prev)
            mean_ref = y_t + ref_preds.pred_vel * cur_dt

        var = eta * eta
        # KL( N(mean_new, var) || N(mean_ref, var) ) = ||dmu||^2 / (2 var);
        # sum over the feature dim only -> per (B, K, A) (matches logp granularity).
        kl = (0.5 * ((mean_new - mean_ref) ** 2) / var).sum(dim=3)         # [B, K, A]
        return logp_new, kl