File size: 7,096 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """
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
|