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