From 141900c0139f900bb5f6ea6480f8c2cb6f2c57ea Mon Sep 17 00:00:00 2001 From: Shintaro Sakoda Date: Tue, 21 Jul 2026 21:27:53 +0900 Subject: [PATCH 1/3] Revert RTC Signed-off-by: Shintaro Sakoda --- .../diffusion_planner/grpo_utils.py | 18 +- .../diffusion_utils/dpm_solver_pytorch.py | 158 ++++++++++++++---- .../model/diffusion_utils/sde.py | 10 ++ .../model/guidance/composer.py | 15 +- .../diffusion_planner/model/module/decoder.py | 77 +++------ .../diffusion_planner/model/module/dit.py | 110 ++++++++---- .../diffusion_planner/utils/onnx_export.py | 6 +- .../diffusion_planner/validate_model.py | 11 +- rlvr/grpo_logprob_loss.py | 14 +- rlvr/grpo_loss.py | 90 ++-------- rlvr/grpo_sft_trainer.py | 19 +-- 11 files changed, 267 insertions(+), 261 deletions(-) diff --git a/diffusion_planner/diffusion_planner/grpo_utils.py b/diffusion_planner/diffusion_planner/grpo_utils.py index ed0c0ad86..a7edb5804 100644 --- a/diffusion_planner/diffusion_planner/grpo_utils.py +++ b/diffusion_planner/diffusion_planner/grpo_utils.py @@ -22,8 +22,6 @@ supported; the helpers raise a clear error otherwise. """ -import random - import torch from diffusion_planner.dimensions import MAX_NUM_AGENTS, OUTPUT_T, POSE_DIM @@ -34,7 +32,6 @@ loss_func, ) from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear -from diffusion_planner.model.module.decoder import generate_prefix_mask from diffusion_planner.utils.unicycle_accel_curvature import smoothing_future_trajectory @@ -80,7 +77,6 @@ def sample_group( inference_inputs["sampled_trajectories"] = ( torch.randn(B, MAX_NUM_AGENTS, OUTPUT_T + 1, POSE_DIM, device=device) * per_row_scale ) - inference_inputs["delay"] = torch.zeros(B, dtype=torch.float32, device=device) _, outputs = model(inference_inputs) ego_world = outputs["prediction"][:, 0].detach() # [B*N, T, 4] @@ -269,7 +265,6 @@ def compute_grpo_loss( ego_target = ego_pseudo_gt.detach() B, Pn, T, _ = neighbors_future.shape - P = 1 + Pn device = ego_pseudo_gt.device ego_current = norm_inputs["ego_current_state"][:, :4] @@ -287,30 +282,21 @@ def compute_grpo_loss( eps = 1e-3 t = torch.rand(B, device=device) * (1 - eps) + eps - t = t.view(B, 1, 1, 1).expand(B, P, T + 1, 1) z = torch.randn_like(gt_future) - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (B,), device=device) - prefix_mask = generate_prefix_mask(delay, P, T + 1) # [B, P, T+1, 1] - mask_coeff = random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t * mask_coeff, torch.tensor(eps, device=device)) - t = torch.where(prefix_mask, curr_mask_time, t) - all_gt = torch.cat([current_states[:, :, None, :], norm(gt_future)], dim=2) # [B, P, T+1, 4] all_gt[:, 1:][neighbor_mask] = 0.0 - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t[..., 1:, :]) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * z xT = torch.cat([all_gt[:, :, :1, :], xT], dim=2) - xT = torch.where(prefix_mask, all_gt, xT) merged_inputs = { **norm_inputs, "gt_trajectories": all_gt, "sampled_trajectories": xT, "diffusion_time": t, - "prefix_mask": prefix_mask, } _, decoder_output = model(merged_inputs) model_output = decoder_output["model_output"][:, :, 1:, :] # [B, P, T, 4] diff --git a/diffusion_planner/diffusion_planner/model/diffusion_utils/dpm_solver_pytorch.py b/diffusion_planner/diffusion_planner/model/diffusion_utils/dpm_solver_pytorch.py index 6ab9c62bf..8053ea394 100644 --- a/diffusion_planner/diffusion_planner/model/diffusion_utils/dpm_solver_pytorch.py +++ b/diffusion_planner/diffusion_planner/model/diffusion_utils/dpm_solver_pytorch.py @@ -181,7 +181,6 @@ def model_fn(x, t_continuous) -> noise: def noise_pred_fn(x, t_continuous, cond=None): if cond is None: - x = x.reshape(x.shape[0], x.shape[1], -1, 4) output = model(x, t_continuous, **model_kwargs) else: output = model(x, t_continuous, cond, **model_kwargs) @@ -192,7 +191,7 @@ def noise_pred_fn(x, t_continuous, cond=None): noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous), ) - return (x - alpha_t * output) / sigma_t + return (x - expand_dims(alpha_t, x.dim()) * output) / expand_dims(sigma_t, x.dim()) elif model_type == "v": alpha_t, sigma_t = ( noise_schedule.marginal_alpha(t_continuous), @@ -223,11 +222,7 @@ def model_fn(x, t_continuous): cond_grad = cond_grad_fn(x, t_continuous) sigma_t = noise_schedule.marginal_std(t_continuous) noise = noise_pred_fn(x, t_continuous) - # Reshape cond_grad to match noise shape for compatibility. - # cond_grad may be 3D [B,P,T*4] while noise is 4D [B,P,T,4] - # when prefix_constraint flattens x between solver steps. - cond_grad = cond_grad.reshape(noise.shape) - return noise - guidance_scale * expand_dims(sigma_t, noise.dim()) * cond_grad + return noise - guidance_scale * expand_dims(sigma_t, x.dim()) * cond_grad elif guidance_type == "classifier-free": if guidance_scale == 1.0 or unconditional_condition is None: return noise_pred_fn(x, t_continuous, cond=condition) @@ -250,12 +245,17 @@ def __init__( self, model_fn, noise_schedule, + correcting_x0_fn=None, correcting_xt_fn=None, + thresholding_max_val=1.0, + dynamic_thresholding_ratio=0.995, ): """Construct a DPM-Solver. We support only DPM-Solver++. + We also support the "dynamic thresholding" method in Imagen[1]. For pixel-space diffusion models, you + can set `correcting_x0_fn="dynamic_thresholding"` to use the dynamic thresholding. The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales. Note that the thresholding method is **unsuitable** for latent-space DPMs (such as stable-diffusion). @@ -271,6 +271,20 @@ def model_fn(x, t_continuous): `` The shape of `x` is `(batch_size, **shape)`, and the shape of `t_continuous` is `(batch_size,)`. noise_schedule: A noise schedule object, such as NoiseScheduleVP. + correcting_x0_fn: A `str` or a function with the following format: + ``` + def correcting_x0_fn(x0, t): + x0_new = ... + return x0_new + ``` + This function is to correct the outputs of the data prediction model at each sampling step. e.g., + ``` + x0_pred = data_pred_model(xt, t) + if correcting_x0_fn is not None: + x0_pred = correcting_x0_fn(x0_pred, t) + xt_1 = update(x0_pred, xt, t) + ``` + If `correcting_x0_fn="dynamic_thresholding"`, we use the dynamic thresholding proposed in Imagen[1]. correcting_xt_fn: A function with the following format: ``` def correcting_xt_fn(xt, t, step): @@ -282,25 +296,50 @@ def correcting_xt_fn(xt, t, step): xt = ... xt = correcting_xt_fn(xt, t, step) ``` + thresholding_max_val: A `float`. The max value for thresholding. + Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. + dynamic_thresholding_ratio: A `float`. The ratio for dynamic thresholding (see Imagen[1] for details). + Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. + [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b. """ - self.model = model_fn + self.model = lambda x, t: model_fn(x, t.expand((x.shape[0]))) self.noise_schedule = noise_schedule + if correcting_x0_fn == "dynamic_thresholding": + self.correcting_x0_fn = self.dynamic_thresholding_fn + else: + self.correcting_x0_fn = correcting_x0_fn self.correcting_xt_fn = correcting_xt_fn + self.dynamic_thresholding_ratio = dynamic_thresholding_ratio + self.thresholding_max_val = thresholding_max_val + + def dynamic_thresholding_fn(self, x0, t): + """ + The dynamic thresholding method. + """ + dims = x0.dim() + p = self.dynamic_thresholding_ratio + s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1) + s = expand_dims( + torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims + ) + x0 = torch.clamp(x0, -s, s) / s + return x0 def data_prediction_fn(self, x, t): """ Return the data prediction model (with corrector). """ noise = self.model(x, t) - x = x.reshape(x.shape[0], x.shape[1], -1, 4) alpha_t, sigma_t = ( self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t), ) x0 = (x - sigma_t * noise) / alpha_t + if self.correcting_x0_fn is not None: + x0 = self.correcting_x0_fn(x0, t) return x0 def model_fn(self, x, t): @@ -348,6 +387,72 @@ def get_time_steps(self, skip_type, t_T, t_0, N, device): ) ) + def get_orders_and_timesteps_for_singlestep_solver( + self, steps, order, skip_type, t_T, t_0, device + ): + """ + Get the order of each step for sampling by the singlestep DPM-Solver. + + We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". + Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is: + - If order == 1: + We take `steps` of DPM-Solver-1 (i.e. DDIM). + - If order == 2: + - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. + - If steps % 2 == 0, we use K steps of DPM-Solver-2. + - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. + - If order == 3: + - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. + - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. + - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. + - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2. + + ============================================ + Args: + order: A `int`. The max order for the solver (2 or 3). + steps: A `int`. The total number of function evaluations (NFE). + skip_type: A `str`. The type for the spacing of the time steps. We support three types: + - 'logSNR': uniform logSNR for the time steps. + - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) + - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) + t_T: A `float`. The starting time of the sampling (default is T). + t_0: A `float`. The ending time of the sampling (default is epsilon). + device: A torch device. + Returns: + orders: A list of the solver order of each step. + """ + if order == 3: + K = steps // 3 + 1 + if steps % 3 == 0: + orders = [3] * (K - 2) + [2, 1] + elif steps % 3 == 1: + orders = [3] * (K - 1) + [1] + else: + orders = [3] * (K - 1) + [2] + elif order == 2: + if steps % 2 == 0: + K = steps // 2 + orders = [2] * K + else: + K = steps // 2 + 1 + orders = [2] * (K - 1) + [1] + elif order == 1: + K = steps + orders = [1] * steps + else: + raise ValueError("'order' must be '1' or '2' or '3'.") + if skip_type == "logSNR": + # To reproduce the results in DPM-Solver paper + timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device) + else: + timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[ + torch.cumsum( + torch.tensor([0] + orders), + 0, + ).to(device) + ] + return timesteps_outer, orders + def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False): """ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. @@ -372,7 +477,6 @@ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=Fal phi_1 = torch.expm1(-h) if model_s is None: model_s = self.model_fn(x, s) - model_s = model_s.reshape(x.shape) x_t = sigma_t / sigma_s * x - alpha_t * phi_1 * model_s if return_intermediate: return x_t, {"model_s": model_s} @@ -411,7 +515,6 @@ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t) r0 = h_0 / h D1_0 = (1.0 / r0) * (model_prev_0 - model_prev_1) phi_1 = torch.expm1(-h) - x = x.reshape(x.shape[0], x.shape[1], -1, 4) x_t = ( (sigma_t / sigma_prev_0) * x - (alpha_t * phi_1) * model_prev_0 @@ -439,7 +542,7 @@ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order) else: raise ValueError("Solver order must be 1 or 2, got {}".format(order)) - def sample(self, x, steps, prefix_mask, skip_type="time_uniform"): + def sample(self, x, steps, skip_type="time_uniform"): """ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`. @@ -476,9 +579,9 @@ def sample(self, x, steps, prefix_mask, skip_type="time_uniform"): ===================================================== Args: - x: (B, P, T, D) + x: A pytorch tensor. The initial value at time `t_start` + e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution. steps: A `int`. The total number of function evaluations (NFE). - prefix_mask: (B, P, T, 1) skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'. Returns: x_end: A pytorch tensor. The approximated solution at time `t_end`. @@ -499,45 +602,30 @@ def sample(self, x, steps, prefix_mask, skip_type="time_uniform"): "Cannot use adaptive solver when correcting_xt_fn is not None" ) device = x.device - T = 81 - x = x.reshape(x.shape[0], x.shape[1], T, -1) - t_shape = (x.shape[0], x.shape[1], T, 1) with torch.no_grad(): assert steps >= order timesteps = self.get_time_steps( skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device ) - timesteps_masked = self.get_time_steps( - skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device - ) assert timesteps.shape[0] - 1 == steps # Init the initial values. step = 0 t = timesteps[step] - t_masked = timesteps_masked[step] - t_BPT1 = t.reshape((1, 1, 1, 1)).expand(t_shape) - t_BPT1 = torch.where(prefix_mask, t_masked, t_BPT1) t_prev_list = [t] - model_prev_list = [self.model_fn(x, t_BPT1)] + model_prev_list = [self.model_fn(x, t)] if self.correcting_xt_fn is not None: x = self.correcting_xt_fn(x, t, step) # Init the first `order` values by lower order multistep DPM-Solver. for step in range(1, order): t = timesteps[step] - t_masked = timesteps_masked[step] - t_BPT1 = t.reshape((1, 1, 1, 1)).expand(t_shape) - t_BPT1 = torch.where(prefix_mask, t_masked, t_BPT1) x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, t, step) if self.correcting_xt_fn is not None: x = self.correcting_xt_fn(x, t, step) t_prev_list.append(t) - model_prev_list.append(self.model_fn(x, t_BPT1)) + model_prev_list.append(self.model_fn(x, t)) # Compute the remaining values by `order`-th order multistep DPM-Solver. for step in range(order, steps + 1): t = timesteps[step] - t_masked = timesteps_masked[step] - t_BPT1 = t.reshape((1, 1, 1, 1)).expand(t_shape) - t_BPT1 = torch.where(prefix_mask, t_masked, t_BPT1) # We only use lower order for steps < 10 if steps < 10: step_order = min(order, steps + 1 - step) @@ -552,12 +640,10 @@ def sample(self, x, steps, prefix_mask, skip_type="time_uniform"): t_prev_list[-1] = t # We do not need to evaluate the final model value. if step < steps: - model_prev_list[-1] = self.model_fn(x, t_BPT1) + model_prev_list[-1] = self.model_fn(x, t) if denoise_to_zero: t = torch.ones((1,)).to(device) * t_0 - t_BPT1 = t.reshape((1, 1, 1, 1)).expand(t_shape) - t_BPT1 = torch.where(prefix_mask, t_0, t_BPT1) - x = self.data_prediction_fn(x, t_BPT1) + x = self.data_prediction_fn(x, t) if self.correcting_xt_fn is not None: x = self.correcting_xt_fn(x, t, step + 1) return x @@ -578,6 +664,4 @@ def expand_dims(v, dims): Returns: a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`. """ - if v.dim() >= dims: - return v return v[(...,) + (None,) * (dims - 1)] diff --git a/diffusion_planner/diffusion_planner/model/diffusion_utils/sde.py b/diffusion_planner/diffusion_planner/model/diffusion_utils/sde.py index 4dee92f8c..f2b6f8f11 100644 --- a/diffusion_planner/diffusion_planner/model/diffusion_utils/sde.py +++ b/diffusion_planner/diffusion_planner/model/diffusion_utils/sde.py @@ -79,6 +79,11 @@ def sde(self, x, t): drift = $-\frac{\beta(t)}{2} x$ diffusion = $\sqrt{\beta(t)}$ """ + shape = x.shape + reshape = [-1] + [ + 1, + ] * (len(shape) - 1) + t = t.reshape(reshape) beta_t = (self._beta_max - self._beta_min) * t + self._beta_min drift = -0.5 * beta_t * x @@ -90,6 +95,11 @@ def marginal_prob(self, x, t): """ Parameters to determine the marginal distribution of the SDE, $p_t(x)$. """ + shape = x.shape + reshape = [-1] + [ + 1, + ] * (len(shape) - 1) + t = t.reshape(reshape) mean_log_coeff = -0.25 * t**2 * (self._beta_max - self._beta_min) - 0.5 * self._beta_min * t mean = torch.exp(mean_log_coeff) * x diff --git a/diffusion_planner/diffusion_planner/model/guidance/composer.py b/diffusion_planner/diffusion_planner/model/guidance/composer.py index f51b11ea8..3f78364be 100644 --- a/diffusion_planner/diffusion_planner/model/guidance/composer.py +++ b/diffusion_planner/diffusion_planner/model/guidance/composer.py @@ -38,20 +38,21 @@ def __call__(self, x_in, t_input, cond, *args, **kwargs): model = kwargs["model"] model_condition = kwargs["model_condition"] - # x_in may be 3D [B,P,T*4] (after prefix_constraint) or 4D [B,P,T,4]. - # The DiT model requires 4D, so reshape for the x_start correction. + # The DiT model takes a flattened 3D trajectory [B, P, T*4] and a scalar + # diffusion time per batch element. x_in from the solver is 3D [B, P, T*4]. + # Extract one scalar t per batch element (also used for time-gating below). + t_scalar = t_input.reshape(B, -1)[:, 0] if t_input.dim() > 1 else t_input + + x_3d = x_in.reshape(B, P, -1) + model_out = model(x_3d, t_scalar, **model_condition) # [B, P, T*4] x_4d = x_in.reshape(B, P, -1, 4) - t_4d = t_input if t_input.dim() == 4 else t_input - x_fix = model(x_4d, t_4d, **model_condition).detach() - x_4d.detach() + x_fix = model_out.reshape(B, P, -1, 4).detach() - x_4d.detach() x_fix[:, :, 0] = 0.0 x_corrected = x_4d + x_fix x_phys = state_normalizer.inverse(x_corrected.detach()) inputs = observation_normalizer.inverse(kwargs["inputs"]) - # Extract one scalar t per batch element for time-gating in BaseGuidance.energy() - t_scalar = t_input.reshape(B, -1)[:, 0] if t_input.dim() > 1 else t_input - # Compute guidance gradient on detached 4D trajectory, then use # surrogate energy = dot(grad, x_in) so autograd returns a gradient # matching x_in's shape (3D or 4D), compatible with the DPM solver. diff --git a/diffusion_planner/diffusion_planner/model/module/decoder.py b/diffusion_planner/diffusion_planner/model/module/decoder.py index cc5b6e102..95778f0bf 100644 --- a/diffusion_planner/diffusion_planner/model/module/decoder.py +++ b/diffusion_planner/diffusion_planner/model/module/decoder.py @@ -1,4 +1,3 @@ -import random from argparse import Namespace from functools import partial @@ -27,30 +26,6 @@ from diffusion_planner.utils.normalizer import ObservationNormalizer, StateNormalizer -def generate_prefix_mask(delay: torch.Tensor, num_agents: int, max_len: int) -> torch.Tensor: - """Generates a prefix mask based on a delay tensor. - - Args: - delay: A 1D tensor of shape (B,) with delay values. - num_agents: The number of agents (P). - max_len: The maximum length of the sequence (T+1 or T_plus_1). - - Returns: - A 4D boolean tensor of shape (B, num_agents, max_len, 1) where mask[i, :, j, 0] is True if j <= delay[i]. - """ - # Create steps tensor (1, 1, max_len, 1) - steps = torch.arange(max_len, device=delay.device).view(1, 1, -1, 1) - # Reshape delay to (B, 1, 1, 1) for broadcasting - reshaped_delay = delay.reshape(delay.shape[0], 1, 1, 1) - # Perform the comparison, result is (B, 1, max_len, 1) - mask = steps <= reshaped_delay - ego_mask = mask.expand(-1, 1, -1, -1) - neighbor_mask = torch.zeros( - (delay.shape[0], num_agents - 1, max_len, 1), dtype=torch.bool, device=delay.device - ) - return torch.cat([ego_mask, neighbor_mask], dim=1) - - def replace_current_state(x: torch.Tensor, current_states: torch.Tensor) -> torch.Tensor: """Return a trajectory tensor with the first timestep replaced.""" return torch.cat([current_states[:, :, None, :], x[:, :, 1:, :]], dim=2) @@ -78,7 +53,6 @@ def compute_training_loss( neighbors_future_valid = ~neighbor_future_mask # [B, Pn, V] B, Pn, T, _ = neighbors_future.shape - P = 1 + Pn ego_current, neighbors_current = ( inputs["ego_current_state"][:, :4], inputs["neighbor_agents_past"][:, :Pn, -1, :4], @@ -96,17 +70,8 @@ def compute_training_loss( eps = 1e-3 t = torch.rand(B, device=gt_future.device) * (1 - eps) + eps # [B,] - t = t.view(B, 1, 1, 1) - t = t.expand(B, P, T + 1, 1) z = torch.randn_like(gt_future, device=gt_future.device) # [B, P, T, 4] - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (B,), device=gt_future.device) # [B,] - prefix_mask = generate_prefix_mask(delay, 1 + Pn, T + 1) # (B, P, T+1, 1) - mask_coeff = random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t * mask_coeff, torch.tensor(eps, device=gt_future.device)) - t = torch.where(prefix_mask, curr_mask_time, t) - if use_velocity: full_traj = torch.cat([current_states[:, :, None, :], gt_future], dim=2) # [B, P, T+1, 4] gt_velocity = waypoints_to_velocity(full_traj) # [B, P, T, 4] @@ -116,19 +81,17 @@ def compute_training_loss( all_gt[:, 1:][neighbor_mask] = 0.0 if model_type == "x_start": - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t[..., 1:, :]) - # mean([B, P, T, D]), std([B, 1, T, 1]), z([B, P, T, D]) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * z xT = torch.cat([all_gt[:, :, :1, :], xT], dim=2) - xT = torch.where(prefix_mask, all_gt, xT) # [B, P, 1 + T, 4] merged_inputs = { **inputs, "gt_trajectories": all_gt, "sampled_trajectories": xT, "diffusion_time": t, - "prefix_mask": prefix_mask, } _, decoder_output = model(merged_inputs) # [B, P, 1 + T, 4] model_output = decoder_output["model_output"][:, :, 1:, :] # [B, P, T, 4] @@ -185,7 +148,6 @@ def compute_training_loss( "gt_trajectories": all_gt, "sampled_trajectories": xT, "diffusion_time": t, - "prefix_mask": prefix_mask, } _, decoder_output = model(merged_inputs) # [B, P, 1 + T, 4] model_output = decoder_output["model_output"][:, :, 1:, :] # [B, P, T, 4] @@ -316,7 +278,18 @@ def _basic_init(m): self.apply(_basic_init) + # Initialize timestep embedding MLP: + nn.init.normal_(self.dit.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.dit.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + for block in self.dit.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + # Zero-out output layers: + nn.init.constant_(self.dit.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.dit.final_layer.adaLN_modulation[-1].bias, 0) nn.init.constant_(self.dit.final_layer.proj[-1].weight, 0) nn.init.constant_(self.dit.final_layer.proj[-1].bias, 0) @@ -373,7 +346,7 @@ def _forward_training(self, encoding, inputs, neighbor_current_mask, encoding_po P = 1 + self._predicted_neighbor_num sampled_trajectories = inputs["sampled_trajectories"].reshape( - B, P, (1 + self._future_len), 4 + B, P, (1 + self._future_len) * 4 ) diffusion_time = inputs["diffusion_time"] @@ -461,19 +434,14 @@ def _inference_x_start( B = encoding.shape[0] P = 1 + self._predicted_neighbor_num - action_prefix = sampled_trajectories.reshape(B, P, 1 + self._future_len, 4) - action_prefix = replace_current_state(action_prefix, current_states) - xT = action_prefix.reshape(B, P, (1 + self._future_len) * 4) - - B, P, T_plus_1, D = action_prefix.shape + xT = sampled_trajectories.reshape(B, P, 1 + self._future_len, 4) + xT = replace_current_state(xT, current_states) + xT = xT.reshape(B, P, (1 + self._future_len) * 4) - delay = inputs["delay"].to(device=action_prefix.device) - mask = generate_prefix_mask(delay, P, T_plus_1) # (B, P, T_plus_1, 1) - - def prefix_constraint(xt, t, step): + def initial_state_constraint(xt, t, step): xt = xt.reshape(B, P, 1 + self._future_len, 4) xt = replace_current_state(xt, current_states) - return xt + return xt.reshape(B, P, -1) model_wrapper_params = { "classifier_fn": self._guidance_fn, @@ -504,9 +472,11 @@ def prefix_constraint(xt, t, step): **model_wrapper_params, ) - dpm_solver = dpm.DPM_Solver(model_fn, noise_schedule, correcting_xt_fn=prefix_constraint) + dpm_solver = dpm.DPM_Solver( + model_fn, noise_schedule, correcting_xt_fn=initial_state_constraint + ) - x0 = dpm_solver.sample(xT, steps=10, prefix_mask=mask, skip_type="logSNR") + x0 = dpm_solver.sample(xT, steps=10, skip_type="logSNR") x0 = x0.reshape(B, P, (1 + self._future_len), 4) ego_trajectory = x0[:, 0, 1::10, :2].reshape(B, 2 * (self._future_len // 10)) @@ -576,7 +546,6 @@ def forward(self, encoding, inputs): "neighbor_agent_past": past and current neighbor states, "sampled_trajectories": sampled current-future ego & neighbor states, [B, P, 1 + self._future_len, 4] - "delay": number of initial steps to keep fixed (>=0), [training-only] "diffusion_time": timestep of diffusion process $t \in [0, 1]$, [B] ... } diff --git a/diffusion_planner/diffusion_planner/model/module/dit.py b/diffusion_planner/diffusion_planner/model/module/dit.py index eb5b2f61f..1014386d8 100644 --- a/diffusion_planner/diffusion_planner/model/module/dit.py +++ b/diffusion_planner/diffusion_planner/model/module/dit.py @@ -5,11 +5,67 @@ from timm.models.layers import Mlp -def modulate(x, shift, scale): - x = x * (1 + scale) + shift +def modulate(x, shift, scale, only_first=False): + if only_first: + x_first, x_rest = x[:, :1], x[:, 1:] + x = torch.cat([x_first * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), x_rest], dim=1) + else: + x = x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + return x + + +def scale(x, scale, only_first=False): + if only_first: + x_first, x_rest = x[:, :1], x[:, 1:] + x = torch.cat([x_first * (1 + scale.unsqueeze(1)), x_rest], dim=1) + else: + x = x * (1 + scale.unsqueeze(1)) + return x +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + class DiTBlock(nn.Module): """ A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning for ego and Cross-Attention. @@ -37,12 +93,12 @@ def __init__(self, dim=192, heads=6, dropout=0.1, mlp_ratio=4.0): def forward(self, x, cross_c, y, attn_mask): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation( y - ).chunk(6, dim=2) + ).chunk(6, dim=1) modulated_x = modulate(self.norm1(x), shift_msa, scale_msa) x = ( x - + gate_msa + + gate_msa.unsqueeze(1) * self.attn( modulated_x, modulated_x, @@ -53,7 +109,7 @@ def forward(self, x, cross_c, y, attn_mask): ) modulated_x = modulate(self.norm2(x), shift_mlp, scale_mlp) - x = x + gate_mlp * self.mlp1(modulated_x) + x = x + gate_mlp.unsqueeze(1) * self.mlp1(modulated_x) x = x + self.cross_attn(self.norm3(x), cross_c, cross_c, need_weights=False)[0] x = x + self.mlp2(self.norm4(x)) @@ -84,7 +140,7 @@ def __init__(self, hidden_size, output_size): def forward(self, x, y): B, P, _ = x.shape - shift, scale = self.adaLN_modulation(y).chunk(2, dim=2) + shift, scale = self.adaLN_modulation(y).chunk(2, dim=1) x = modulate(self.norm_final(x), shift, scale) x = self.proj(x) return x @@ -102,23 +158,15 @@ def __init__( ): super().__init__() - T = 81 - D = 4 self.agent_embedding = nn.Embedding(2, hidden_dim) self.preproj = Mlp( - in_features=T * D, - hidden_features=512, - out_features=hidden_dim, - act_layer=nn.GELU, - drop=0.0, - ) - self.t_embedder = Mlp( - in_features=T, + in_features=output_dim, hidden_features=512, out_features=hidden_dim, act_layer=nn.GELU, drop=0.0, ) + self.t_embedder = TimestepEmbedder(hidden_dim) self.blocks = nn.ModuleList( [DiTBlock(hidden_dim, heads, dropout, mlp_ratio) for i in range(depth)] ) @@ -127,20 +175,13 @@ def __init__( def forward(self, x, t, cross_c, neighbor_current_mask): """ Forward pass of DiT. - x: (B, P, T, D) -> Embedded out of DiT - t: (B, P, T, 1) + x: (B, P, output_dim) -> Embedded out of DiT + t: (B,) cross_c: (B, N, D) -> Cross-Attention context """ - assert x.dim() == 4, f"{x.dim()=}" - assert t.dim() == 4, f"{t.dim()=}" - assert x.shape[2] == t.shape[2], f"{x.shape[2]=} {t.shape[2]=}" - B, P, T, D = x.shape - - x = x.reshape(B, P, T * D) # (B, P, T*D) - t = t.reshape(B, P, T) # (B, P, T) + B, P, _ = x.shape - x = self.preproj(x) # (B, P, hidden_dim) - t = self.t_embedder(t) # (B, P, hidden_dim) + x = self.preproj(x) x_embedding = torch.cat( [ @@ -148,16 +189,17 @@ def forward(self, x, t, cross_c, neighbor_current_mask): self.agent_embedding.weight[1][None, :].expand(P - 1, -1), ], dim=0, - ) # (P, hidden_dim) - x_embedding = x_embedding[None, :, :].expand(B, -1, -1) # (B, P, hidden_dim) + ) # (P, D) + x_embedding = x_embedding[None, :, :].expand(B, -1, -1) # (B, P, D) x = x + x_embedding - ego_mask = torch.zeros((B, 1), dtype=torch.bool, device=x.device) - attn_mask = torch.cat([ego_mask, neighbor_current_mask], dim=1) + y = self.t_embedder(t) + + attn_mask = torch.zeros((B, P), dtype=torch.bool, device=x.device) + attn_mask[:, 1:] = neighbor_current_mask for block in self.blocks: - x = block(x, cross_c, t, attn_mask) + x = block(x, cross_c, y, attn_mask) - x = self.final_layer(x, t) # (B, P, output_dim) - x = x.reshape(B, P, T, D) + x = self.final_layer(x, y) return x diff --git a/diffusion_planner/diffusion_planner/utils/onnx_export.py b/diffusion_planner/diffusion_planner/utils/onnx_export.py index 291438862..c28f3dd80 100644 --- a/diffusion_planner/diffusion_planner/utils/onnx_export.py +++ b/diffusion_planner/diffusion_planner/utils/onnx_export.py @@ -41,7 +41,6 @@ "goal_pose", "ego_shape", "turn_indicators", - "delay", ] ENCODER_INPUT_NAMES = [ @@ -244,7 +243,6 @@ def forward( goal_pose: torch.Tensor, ego_shape: torch.Tensor, turn_indicators: torch.Tensor, - delay: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = { "sampled_trajectories": sampled_trajectories, @@ -263,7 +261,6 @@ def forward( "goal_pose": goal_pose, "ego_shape": ego_shape, "turn_indicators": turn_indicators, - "delay": delay, } _, decoder_outputs = self.model(inputs) return decoder_outputs["prediction"], decoder_outputs["turn_indicator_logit"] @@ -303,7 +300,6 @@ def build_dummy_inputs() -> TensorDict: inputs["goal_pose"] = torch.randn(1, POSE_DIM, dtype=torch.float32) inputs["ego_shape"] = torch.tensor([[2.75, 4.34, 1.70]], dtype=torch.float32) inputs["turn_indicators"] = torch.randint(0, 3, (1, INPUT_T + 1), dtype=torch.float32) - inputs["delay"] = torch.zeros(1, 1, dtype=torch.float32) return inputs @@ -311,7 +307,7 @@ def build_decoder_inputs(inputs: TensorDict, encoding: torch.Tensor) -> TensorDi return { "encoding": encoding, "sampled_trajectories": inputs["sampled_trajectories"], - "diffusion_time": torch.ones(1, MAX_NUM_AGENTS, OUTPUT_T + 1, 1, dtype=torch.float32), + "diffusion_time": torch.ones(1, dtype=torch.float32), "neighbor_agents_past": inputs["neighbor_agents_past"], } diff --git a/diffusion_planner/diffusion_planner/validate_model.py b/diffusion_planner/diffusion_planner/validate_model.py index 389dacc92..d368a8d89 100644 --- a/diffusion_planner/diffusion_planner/validate_model.py +++ b/diffusion_planner/diffusion_planner/validate_model.py @@ -46,7 +46,7 @@ class _PreparedValidationBatch: turn_indicator_seq: torch.Tensor -def _prepare_validation_inputs(inputs, args, device, delay=0) -> _PreparedValidationBatch: +def _prepare_validation_inputs(inputs, args, device) -> _PreparedValidationBatch: inputs = {key: value.to(device) for key, value in inputs.items()} batch_size = inputs["ego_current_state"].shape[0] turn_indicator_seq = inputs["turn_indicators"] @@ -58,7 +58,6 @@ def _prepare_validation_inputs(inputs, args, device, delay=0) -> _PreparedValida dtype=torch.float32, device=device, ) - inputs["delay"] = torch.full((batch_size,), delay, dtype=torch.float32, device=device) inputs["ego_agent_past"] = heading_to_cos_sin(inputs["ego_agent_past"]) inputs["goal_pose"] = heading_to_cos_sin(inputs["goal_pose"]) ego_future = heading_to_cos_sin(inputs["ego_agent_future"]) @@ -81,8 +80,8 @@ def _prepare_validation_inputs(inputs, args, device, delay=0) -> _PreparedValida ) -def _predict_ego_for_temporal_metrics(model, inputs, args, device, delay=0): - batch = _prepare_validation_inputs(inputs, args, device, delay) +def _predict_ego_for_temporal_metrics(model, inputs, args, device): + batch = _prepare_validation_inputs(inputs, args, device) _, outputs = model(batch.inputs) return outputs["prediction"][:, 0], batch.ego_future @@ -293,8 +292,6 @@ def validate_model(model, val_loader, args, return_pred=False) -> tuple[float, f turn_indicator_change_correct = 0.0 turn_indicator_change_total = 0 - delay = 0 - # Progress is driven by the SLOWEST rank: every `progress_sync_every` batches all # ranks rendezvous on an all-reduce(MIN) of their completed-batch count and rank 0 # displays that minimum, so the bar reaches 100% only when every rank is done. The @@ -303,7 +300,7 @@ def validate_model(model, val_loader, args, return_pred=False) -> tuple[float, f total_batches = len(val_loader) pbar = tqdm(total=total_batches, desc="validate (slowest rank)", disable=ddp.get_rank() != 0) for step, inputs in enumerate(val_loader): - prepared = _prepare_validation_inputs(inputs, args, device, delay) + prepared = _prepare_validation_inputs(inputs, args, device) inputs = prepared.inputs ego_future = prepared.ego_future neighbors_future = prepared.neighbors_future diff --git a/rlvr/grpo_logprob_loss.py b/rlvr/grpo_logprob_loss.py index 6dbcebee9..edd5bc53f 100644 --- a/rlvr/grpo_logprob_loss.py +++ b/rlvr/grpo_logprob_loss.py @@ -16,7 +16,6 @@ import torch import torch.nn.functional as F from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear -from diffusion_planner.model.module.decoder import generate_prefix_mask from rlvr.vpsde_logprob import ( compute_discount_weights, @@ -103,12 +102,8 @@ def _build_model_inputs( neighbor_slice = all_gt[:, 1 : 1 + nf_pn] neighbor_slice.masked_fill_(full_neighbor_mask.unsqueeze(-1).expand_as(neighbor_slice), 0.0) - # Build t tensor [N, P, T+1, 1] - t_4d = torch.full((N, P, future_len + 1, 1), t_value, device=device) - - # Prefix mask (no random delay for logprob — use fixed delay=0) - delay = torch.zeros(N, dtype=torch.long, device=device) - prefix_mask = generate_prefix_mask(delay, P, future_len + 1) + # Build t tensor [N] + t = torch.full((N,), t_value, device=device) # Normalize observation data data_normalized = model_args.observation_normalizer( @@ -118,10 +113,7 @@ def _build_model_inputs( merged = {**data_normalized} merged["gt_trajectories"] = all_gt merged["sampled_trajectories"] = x_t - merged["diffusion_time"] = t_4d - merged["prefix_mask"] = prefix_mask - if "delay" not in merged: - merged["delay"] = delay + merged["diffusion_time"] = t return merged, all_gt diff --git a/rlvr/grpo_loss.py b/rlvr/grpo_loss.py index acfaeba6d..ab592d533 100644 --- a/rlvr/grpo_loss.py +++ b/rlvr/grpo_loss.py @@ -39,37 +39,18 @@ def compute_trajectory_loss(model, data, trajectory, model_args, noise, t, device): - """V4-compatible trajectory loss using 4D diffusion timestep. + """Trajectory loss using scalar diffusion timestep. - Matches SFT training in decoder.py: includes prefix mask with random delay, - per-timestep t modulation, and clean prefix injection. + Matches SFT training in decoder.py: samples the denoiser at a scalar + diffusion time `t` and regresses the model output toward the GT trajectory. """ - import random as _random - from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear - from diffusion_planner.model.module.decoder import generate_prefix_mask B = data["ego_current_state"].shape[0] P = 1 + model_args.predicted_neighbor_num future_len = model_args.future_len eps = 1e-3 - # Expand t to [B, P, T+1, 1] - if t.dim() == 1: - t_4d = t.view(B, 1, 1, 1).expand(B, P, future_len + 1, 1).clone() - elif t.dim() == 4: - t_4d = t - else: - t_4d = t.view(B, 1, 1, 1).expand(B, P, future_len + 1, 1).clone() - - # Prefix mask with random delay — matches SFT (decoder.py line 95-100) - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (B,), device=device) - prefix_mask = generate_prefix_mask(delay, P, future_len + 1) # (B, P, T+1, 1) - mask_coeff = _random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t_4d * mask_coeff, torch.tensor(eps, device=device)) - t_4d = torch.where(prefix_mask, curr_mask_time, t_4d) - gt_trajectory = torch.as_tensor(trajectory, dtype=torch.float32, device=device) if gt_trajectory.dim() == 2: gt_trajectory = gt_trajectory.unsqueeze(0) # [T, 4] → [1, T, 4] @@ -91,12 +72,11 @@ def compute_trajectory_loss(model, data, trajectory, model_args, noise, t, devic all_gt = torch.cat([current_states[:, :, None, :], gt_future], dim=2) # [B, P, T+1, 4] - # Diffusion noise with prefix masking — matches SFT (decoder.py line 111-116) - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t_4d[..., 1:, :]) + # Diffusion noise — matches SFT (decoder.py) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * noise xT_full = torch.cat([all_gt[:, :, :1, :], xT], dim=2) # [B, P, T+1, 4] - # Prefix: replace noised steps with clean GT - xT_full = torch.where(prefix_mask, all_gt, xT_full) data_for_norm = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in data.items()} data_normalized = model_args.observation_normalizer(data_for_norm) @@ -104,10 +84,7 @@ def compute_trajectory_loss(model, data, trajectory, model_args, noise, t, devic merged_inputs = {**data_normalized} merged_inputs["gt_trajectories"] = all_gt merged_inputs["sampled_trajectories"] = xT_full - merged_inputs["diffusion_time"] = t_4d # [B, P, T+1, 1] - merged_inputs["prefix_mask"] = prefix_mask - if "delay" not in merged_inputs: - merged_inputs["delay"] = delay + merged_inputs["diffusion_time"] = t _, outputs = model(merged_inputs) @@ -276,10 +253,7 @@ def compute_batched_trajectory_losses( Returns: [N] tensor of per-trajectory MSE losses. """ - import random as _random - from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear - from diffusion_planner.model.module.decoder import generate_prefix_mask N = trajectories_tensor.shape[0] P = 1 + model_args.predicted_neighbor_num @@ -302,12 +276,8 @@ def compute_batched_trajectory_losses( else: batch_data[k] = v - # Expand t to [N, P, T+1, 1] — matches SFT (decoder.py line 90-92) - if t.dim() == 1: - t_N = t.expand(N) - t_4d = t_N.view(N, 1, 1, 1).expand(N, P, future_len + 1, 1).clone() - else: - t_4d = t.expand(N, *t.shape[1:]).contiguous() if t.shape[0] == 1 else t + # Expand t to [N] — matches SFT (decoder.py) + t_N = t.expand(N) # Expand noise to [N, P, T, 4] if noise.shape[0] == 1: @@ -315,16 +285,6 @@ def compute_batched_trajectory_losses( else: noise_N = noise - # Prefix mask with random delay — matches SFT (decoder.py line 95-100) - # Forces first `delay` steps to use clean GT, training the model to - # predict the trajectory conditioned on a clean prefix. - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (N,), device=device) - prefix_mask = generate_prefix_mask(delay, P, future_len + 1) # (N, P, T+1, 1) - mask_coeff = _random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t_4d * mask_coeff, torch.tensor(eps, device=device)) - t_4d = torch.where(prefix_mask, curr_mask_time, t_4d) - # Normalize trajectories: [N, T, 4] ego_mean = model_args.state_normalizer.mean[0].to(device) ego_std = model_args.state_normalizer.std[0].to(device) @@ -390,12 +350,11 @@ def compute_batched_trajectory_losses( full_neighbor_mask.unsqueeze(-1).expand_as(neighbor_slice), 0.0 ) - # Diffusion noise with prefix masking — matches SFT (decoder.py line 111-116) - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t_4d[..., 1:, :]) + # Diffusion noise — matches SFT (decoder.py) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t_N) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * noise_N xT_full = torch.cat([all_gt[:, :, :1, :], xT], dim=2) - # Prefix: replace noised steps with clean GT for delay steps - xT_full = torch.where(prefix_mask, all_gt, xT_full) # Normalize observation data data_for_norm = { @@ -406,10 +365,7 @@ def compute_batched_trajectory_losses( merged = {**data_normalized} merged["gt_trajectories"] = all_gt merged["sampled_trajectories"] = xT_full - merged["diffusion_time"] = t_4d - merged["prefix_mask"] = prefix_mask - if "delay" not in merged: - merged["delay"] = delay + merged["diffusion_time"] = t_N _, outputs = model(merged) @@ -572,10 +528,7 @@ def _compute_neighbor_reg_loss( When B>1 (batched trainers expand per-scene data), uses only the first element since all B entries come from the same scene. """ - import random as _random - from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear - from diffusion_planner.model.module.decoder import generate_prefix_mask B = data["ego_current_state"].shape[0] if B > 1: @@ -656,28 +609,17 @@ def _compute_neighbor_reg_loss( total_reg = torch.tensor(0.0, device=device) for _ in range(K): t = torch.rand(1, device=device) * (1 - eps) + eps - t_4d = t.view(1, 1, 1, 1).expand(1, P, future_len + 1, 1).clone() - - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (1,), device=device) - prefix_mask = generate_prefix_mask(delay, P, future_len + 1) - mask_coeff = _random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t_4d * mask_coeff, torch.tensor(eps, device=device)) - t_4d = torch.where(prefix_mask, curr_mask_time, t_4d) z = torch.randn(1, P, future_len, 4, device=device) - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t_4d[..., 1:, :]) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * z xT_full = torch.cat([all_gt[:, :, :1, :], xT], dim=2) - xT_full = torch.where(prefix_mask, all_gt, xT_full) merged = {**data_normalized} merged["gt_trajectories"] = all_gt merged["sampled_trajectories"] = xT_full - merged["diffusion_time"] = t_4d - merged["prefix_mask"] = prefix_mask - if "delay" not in merged: - merged["delay"] = delay + merged["diffusion_time"] = t # LoRA forward (with grad) _, lora_out = policy_model(merged) diff --git a/rlvr/grpo_sft_trainer.py b/rlvr/grpo_sft_trainer.py index 10a3878e3..51b4c8396 100644 --- a/rlvr/grpo_sft_trainer.py +++ b/rlvr/grpo_sft_trainer.py @@ -23,7 +23,6 @@ import torch import torch.nn.functional as F from diffusion_planner.model.diffusion_utils.sde import VPSDE_linear -from diffusion_planner.model.module.decoder import generate_prefix_mask from scipy.signal import savgol_filter from torch import nn from tqdm import tqdm @@ -283,32 +282,20 @@ def _compute_sft_diffusion_loss( for _ in range(K): # Sample random timestep t = torch.rand(B, device=device) * (1 - eps) + eps - t_4d = t.view(B, 1, 1, 1).expand(B, P, future_len + 1, 1).clone() - - # Prefix mask with random delay - max_delay = 5 - delay = torch.randint(0, max_delay + 1, (B,), device=device) - prefix_mask = generate_prefix_mask(delay, P, future_len + 1) - mask_coeff = _random.uniform(0.0, 1.0) - curr_mask_time = torch.maximum(t_4d * mask_coeff, torch.tensor(eps, device=device)) - t_4d = torch.where(prefix_mask, curr_mask_time, t_4d) # Noise and diffusion z = torch.randn(B, P, future_len, 4, device=device) - mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t_4d[..., 1:, :]) + mean, std = VPSDE_linear().marginal_prob(all_gt[..., 1:, :], t) + std = std.view(-1, *([1] * (len(all_gt[..., 1:, :].shape) - 1))) xT = mean + std * z xT_full = torch.cat([all_gt[:, :, :1, :], xT], dim=2) - xT_full = torch.where(prefix_mask, all_gt, xT_full) merged_inputs = { k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in data_normalized.items() } merged_inputs["gt_trajectories"] = all_gt merged_inputs["sampled_trajectories"] = xT_full - merged_inputs["diffusion_time"] = t_4d - merged_inputs["prefix_mask"] = prefix_mask - if "delay" not in merged_inputs: - merged_inputs["delay"] = delay + merged_inputs["diffusion_time"] = t _, outputs = model(merged_inputs) From 75ef0fe4dfad90ec523e0032126de47bc2221921 Mon Sep 17 00:00:00 2001 From: Shintaro Sakoda Date: Mon, 27 Jul 2026 16:57:25 +0900 Subject: [PATCH 2/3] Fixed Diffusion-Planner/diffusion_planner/diffusion_planner/utils/onnx_export.py Signed-off-by: Shintaro Sakoda --- diffusion_planner/diffusion_planner/utils/onnx_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffusion_planner/diffusion_planner/utils/onnx_export.py b/diffusion_planner/diffusion_planner/utils/onnx_export.py index c28f3dd80..9acfd4dc3 100644 --- a/diffusion_planner/diffusion_planner/utils/onnx_export.py +++ b/diffusion_planner/diffusion_planner/utils/onnx_export.py @@ -186,7 +186,7 @@ def forward( agent_num = 1 + self.decoder._predicted_neighbor_num sampled_trajectories = sampled_trajectories.reshape( - batch_size, agent_num, 1 + self.decoder._future_len, 4 + batch_size, agent_num, (1 + self.decoder._future_len) * 4 ) model_output = self.decoder.dit( From cc2b4b7540c5103f25f0ec7b16beed637985079e Mon Sep 17 00:00:00 2001 From: Shintaro Sakoda Date: Mon, 27 Jul 2026 19:20:39 +0900 Subject: [PATCH 3/3] Added dummy processing Signed-off-by: Shintaro Sakoda --- .../diffusion_planner/utils/onnx_export.py | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/diffusion_planner/diffusion_planner/utils/onnx_export.py b/diffusion_planner/diffusion_planner/utils/onnx_export.py index 9acfd4dc3..b2e7b00b8 100644 --- a/diffusion_planner/diffusion_planner/utils/onnx_export.py +++ b/diffusion_planner/diffusion_planner/utils/onnx_export.py @@ -41,6 +41,7 @@ "goal_pose", "ego_shape", "turn_indicators", + "delay", ] ENCODER_INPUT_NAMES = [ @@ -167,6 +168,10 @@ class DecoderONNXWrapper(nn.Module): This wrapper intentionally does not call a sampler. An external denoising loop should update x_t and timesteps, then call this ONNX model once per model evaluation. + + ``diffusion_time`` is taken in the (B, P, 1 + T, 1) layout the deployed ROS node + (``autoware_diffusion_planner``) binds, where every element holds the same timestep, and is + reduced to the per-sample scalar the DiT expects. """ def __init__(self, model: Diffusion_Planner): @@ -191,7 +196,7 @@ def forward( model_output = self.decoder.dit( sampled_trajectories, - diffusion_time, + diffusion_time.reshape(batch_size, -1)[:, 0], encoding, neighbor_current_mask, ).reshape(batch_size, agent_num, 1 + self.decoder._future_len, 4) @@ -219,12 +224,28 @@ def forward(self, encoding: torch.Tensor, final_x0: torch.Tensor) -> torch.Tenso class FullONNXWrapper(nn.Module): - """Original all-in-one planner export.""" + """Original all-in-one planner export. + + Takes a ``delay`` input that this model does not use. The deployed ROS node + (``autoware_diffusion_planner``) always binds a ``delay`` tensor for real-time chunking, and + refuses to load an engine whose signature lacks it. Real-time chunking is not part of this + model, so ``delay`` is routed through a term that is identically zero for the non-negative + step counts the node sends (see :meth:`zero_from_delay`); an input that reaches no output at + all would be pruned out of the exported graph. + """ def __init__(self, model: Diffusion_Planner): super().__init__() self.model = model + def zero_from_delay(self, delay: torch.Tensor) -> torch.Tensor: + """Return a (B, 1, 1, 1) tensor of zeros that data-depends on ``delay``. + + ``torch.clamp(delay, max=0)`` is zero for every ``delay >= 0`` while staying opaque to the + constant folding in ONNX / TensorRT that would eliminate a plain multiplication by zero. + """ + return torch.clamp(delay, max=0.0)[:, :, None, None] + def forward( self, sampled_trajectories: torch.Tensor, @@ -243,9 +264,10 @@ def forward( goal_pose: torch.Tensor, ego_shape: torch.Tensor, turn_indicators: torch.Tensor, + delay: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = { - "sampled_trajectories": sampled_trajectories, + "sampled_trajectories": sampled_trajectories + self.zero_from_delay(delay), "ego_agent_past": ego_agent_past, "ego_current_state": ego_current_state, "neighbor_agents_past": neighbor_agents_past, @@ -300,6 +322,7 @@ def build_dummy_inputs() -> TensorDict: inputs["goal_pose"] = torch.randn(1, POSE_DIM, dtype=torch.float32) inputs["ego_shape"] = torch.tensor([[2.75, 4.34, 1.70]], dtype=torch.float32) inputs["turn_indicators"] = torch.randint(0, 3, (1, INPUT_T + 1), dtype=torch.float32) + inputs["delay"] = torch.zeros(1, 1, dtype=torch.float32) return inputs @@ -307,7 +330,7 @@ def build_decoder_inputs(inputs: TensorDict, encoding: torch.Tensor) -> TensorDi return { "encoding": encoding, "sampled_trajectories": inputs["sampled_trajectories"], - "diffusion_time": torch.ones(1, dtype=torch.float32), + "diffusion_time": torch.ones(1, MAX_NUM_AGENTS, OUTPUT_T + 1, 1, dtype=torch.float32), "neighbor_agents_past": inputs["neighbor_agents_past"], }