revert: real time chunking

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]

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)]

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