CL-DiffPhyCon 代码细节

相比于 DiffPhyCon 的代码，

• trainer.py 完全没变化。
• diffusion.py 里多定义了 asyn_t_seq 函数，两个开关参数 is_init_model 和 asynchronous 及其逻辑，以及 frames 参数。
• train.py 多了两个 args 参数定义 is_model_w 和 asynch_inference_mode，调用 ddpm 时加上这俩参数，其他没变化。
• inference.py 多了 args 参数定义 asynch_inference_mode 和 infer_interval，逻辑变化自然最大。

CL-DiffPhyCon 代码细节

在 CL-DiffPhyCon 里，做的是一个「滑动窗口」策略，窗口宽度是 $H$（代码里 self.frames），但 环境总共要跑的物理步数是 $T\ge H$。

• H：控制序列长度，即 控制窗口长度，是代码中的 self.frames。
• τ：当前闭环控制步，每次闭环循环的步数。
• u_τ：在 τ 时刻要执行的控制信号，u₀ 是代码中的 u_pred[:,:,0]。

区别「窗口长度」$H$ 和「总步数」$T$：

• 想象大雾中你要走一条总长 $T$ 步的小路，每次能看到前 $H$ 步。
• 第一次你从起点看见 $[0\dots H-1]$ 步，但第 0 步最清晰，第 $H-1$ 步最模糊，你走第 0 步；
• 再往前滑一步（窗口变成 $[1\dots H]$），但第 1 步最清晰，第 $H$ 步最模糊，你走第 1 步…
• 当走到第 $T-1$ 步的时候，窗口里或许还能看到 $[T-1\dots T-1+H-1]$ 这么一条“超出终点”的视野，但你永远只会走那第 $T-1$ 步，然后到达终点就停。剩下的在终点以外的那些格子（对应 pad 或多余预测）不会被走到。

左侧 Initialization 部分：

• 用 标准扩散流程（比如 1000 步去噪）从纯噪声一步步去噪，得到 完整控制信号序列：

  if self.is_init_model:
      img = torch.randn(shape, device=device)
      ...
      pred_img, x_start, pred_noise = self.p_sample(...)
  

• 最终得到 u_pred = [u₀, u₁, ..., u_{H-1}]，控制信号可以依次执行。

• 执行第一个控制动作 u₀ = u_pred[:,:,0]，送给环境：

  action = u_pred[:, :, 0]
  

闭环阶段 Closed-loop Control（从 τ ≥ 1 开始）：当前已执行 τ−1 步，下一步是 u_τ，

• 将之前去噪出来的 u_pred 去掉前面的 u_{τ-1}

  u_pred = u_pred[:, :, 1:]  # 弹出上一轮执行的控制
  

• 补两帧 随机噪声，用于 后续去噪预测 u_{τ+H-1} 和 维持输入长度：

  pad_img = torch.randn(...), noise_final = torch.randn(...)
  img = torch.cat([u_pred, noise_final, pad_img], dim=2)
  

• 把这个新的 img（长度又变成了 self.frames+1）作为输入送进扩散采样流程。

异步去噪 self.asyn_t_seq(...)，关键机制是 不同时间步控制帧用不同的噪声级别。输出 u_τ 即 u_pred[:,:,0]，也就是本轮控制动作。

这正是 CL-DiffPhyCon 的关键创新点：不必等完整序列生成完，就能边采样边控制，形成自然的闭环反馈控制流程。

• 完整序列一次性去噪：如果 模型预测有误，后面的动作会一直“背锅”（误差累积），而且要等完整去噪跑完才能开始执行。
• 每个物理步异步去噪一小段并马上输出：每步可立即执行，不用等完整序列；环境反馈马上进来，纠正模型偏差；采样只在新帧上跑最少的迭代，效率更高。

1D Burgers 方程数据集类定义

# data_1d.py
class DiffusionDataset(Dataset):
    def __init__(
        self, 
        fname,   # 存储数据的 .pt 文件路径
        preprocess=get_burgers_preprocess('all'),   # 预处理函数 以及默认调用
        load_all=True  # 一次性用 torch.load 把整个文件读进内存
    ):
        '''
        Arguments:

        '''
        self.load_all = load_all
        if load_all:
            self.db = torch.load(fname)
            self.x = preprocess(self.db)  # 预处理
        else:
            raise NotImplementedError

    def __len__(self):
        if self.load_all:
            return self.x.size(0)
        else:
            raise NotImplementedError

    def __getitem__(self, idx):
        if self.load_all:
            return self.x[idx]
        else:
            raise NotImplementedError

    def get(self, idx): # 别名，方便调用
        return self.__getitem__(idx)
    
    def len(self):  # 别名，方便调用
        return self.__len__()

DiffusionDataset 类用来包装处理好的数据以供 DataLoader 读取。继承自 torch.utils.data.Dataset，需要实现以下两个函数，

• __len__()：返回样本数，self.x.size(0)。

• __getitem__(idx)：返回第 idx 条预处理后的数据 self.x[idx]。

def get_burgers_preprocess(
    rescaler=None, 
    stack_u_and_f=False,  # 是否把 u（状态）和 f（作用力）在新的维度上堆叠，用于 2D 卷积输入（batch×3×…）
    pad_for_2d_conv=False,  # 配合 stack_u_and_f=True 使用，在时间维度上补零，使序列长度对齐到固定大小(15 或 16)
    partially_observed_fill_zero_unobserved=None,  # 若不为 None，则对 u 做部分观测模拟，将中间一段设置为 0, 针对空间的
):
    if rescaler is None:
        raise NotImplementedError('Should specify rescaler. If no rescaler is not used, specify 1.')
    
    def preprocess(db):
        '''We are only returning f and u for now, in the shape of 
        (u0, u1, ..., f0, f1, ...)
        '''
        u = db['u']   # 物理状态张量，shape 是 (batch, nt, nx) 或类似
        f = db['f']   # 作用力张量，shape 是 (batch, nt, nx)
        f = f[:,:15]  # 只取前 15 帧（时间步） # nt-1=15
        
        if stack_u_and_f:  # 若 stack_u_and_f=True，则必须 pad_for_2d_conv=True，先在时间轴（第 −2 维）补零
            assert pad_for_2d_conv
            nt = f.size(-2)
            f = nn.functional.pad(f, (0, 0, 0, 16 - nt), 'constant', 0)  
            u = nn.functional.pad(u, (0, 0, 0, 15 - nt), 'constant', 0) 
            u_target = u  # 把 (u, f, u_target) 三个通道堆到一起，得到 (batch, 3, nt_padded, nx)，方便 2D 卷积
            data = torch.stack((u, f, u_target), dim=1) 
        
        else:  # 若 stack_u_and_f=False，直接在通道维度上拼接
            assert not pad_for_2d_conv
            data = torch.cat((u, f), dim=1)  # shape = (batch, 16+15, nx)
    
        data = data / rescaler  # 将所有数值缩放到合适范围
        return data

    return preprocess

返回一个 preprocess 函数，用来对加载出来的“数据库” db（一个 dict，包含键 u 和 f）做一系列 预处理，最终得到送入模型的张量。

准备两种不同格式的「网络输入」——一种是 一维序列拼接（stack_u_and_f=False），另一种是 为了做 2D 卷积而在通道上堆叠（stack_u_and_f=True）最终 preprocess 返回形状为

• 堆叠：(batch, 3, nt_pad, nx)，堆叠 u、f 和 u_target，这里 u_target 是直接复制的 u 作为目标轨迹状态。
• 拼接：(batch, 16+15, nx)，在空间上对 16 帧状态的值 $(batch,16,n_x)$ 和 15 帧的控制力 $f_1,\dots ,f_{15}$，$(batch,15,n_x)$ 进行拼接。

train 时加载的数据 shape

训练时，使用的是

# train.py
RESCALER = 10.

def get_dataset(train_data_path):
    return DiffusionDataset(
        train_data_path, 
        preprocess=get_burgers_preprocess( 
            rescaler=RESCALER, 
            stack_u_and_f=True, 
            pad_for_2d_conv=True, 
            partially_observed_fill_zero_unobserved = args.partially_observed, 
        )
    )

if __name__ == "__main__":
	dataset = get_dataset(args.train_data_path)
    run_2d_Unet(dataset, args)

data = data / rescaler 这一行会对整个张量做逐元素除法──不论 data 的维度是 (batch, 3, frames+1, nx)（即 stack_u_and_f=True 且做了 pad 之后的 2D 卷积模式）还是 (batch, 16+15, nx)（stack_u_and_f=False 的情况），它都会把张量里所有的数值都除以同一个标量 rescaler。

RESCALER = 10.

def load_burgers_dataset(dataset):
    tmp_dataset = DiffusionDataset(
        dataset, # dataset of f varying in both space and time
        preprocess=get_burgers_preprocess(
            rescaler=RESCALER, 
            stack_u_and_f=True,    # (batch, channel=3, Nt, Nx)   channel: u + f + u_target
            pad_for_2d_conv=True, 
            partially_observed_fill_zero_unobserved = None, # does not matter since only for loading models
        )
    )
    return tmp_dataset

def load_2dconv_model(args):
	dataset = load_burgers_dataset(args.dataset)
    ddpm = get_2d_ddpm(args)
    trainer = Trainer(ddpm, dataset, ...)
    trainer.load(args.checkpoint__model_w)
    return ddpm

inference 时加载的数据 shape

def get_target(target_i, f=False, device=0, dataset='free_u_f_1e5', **dataset_kwargs):
    # repeating in the first dimension if one target is shared with multiple f
    test_dataset = DiffusionDataset(
        dataset, # dataset of f varying in both space and time
        preprocess=get_burgers_preprocess(
            rescaler=1., 
            stack_u_and_f=False,    # 拼接  (batch, 16+15, Nx)
            pad_for_2d_conv=False, 
            **dataset_kwargs, 
        )
    )

    if not f: # return only u
        ret = test_dataset.get(target_i).cuda(device)[..., :16, :]   #  u: (batch, 16, Nx)
    else: # return only f
        ret = test_dataset.get(target_i).cuda(device)[..., 16:, :]   #  f: (batch, 15, Nx)

    if len(ret.size()) == 2: # if target_i is int
        ret = ret.unsqueeze(0)

    return ret

inference 在线滚动机制

def evaluate(
    model_i, model_f, args, rep=1, …, conv2d=True,
):
    batch_size = 50 // rep
    for i in range(rep):
        target_idx = range(i * batch_size, (i+1) * batch_size)
        if conv2d:
            ddpm_mse, J_diffused, J_actual, energy = diffuse_2dconv(
                args,
                custom_metric=…,        # 评估函数，把预测的 f_diffused 计算出的 u_solved 和 u_ground‑truth 对比
                model_i=model_i,        # 初始化模型 checkpoint
                model_f=model_f,        # 反馈模型 checkpoint
                seed=i,
                nablaJ=…, J_scheduler=…, w_scheduler=…,
                clip_denoised=True,
                guidance_u0=True,
                batch_size=batch_size,
                u_init=…,               # 初始流场（真实）作为边界条件
                u_final=None,
                u_pred=None,            # ← 这里暂时传 None
            )
            # … collect metrics …
    return averages…

• model_i/model_f：分别是 初始化模型（Open‑loop 阶段用）和 反馈模型（Closed‑loop 阶段用）；

• u_init：真实环境中观测到的初态 $u_0$ （部分观测时带零填充）；

• u_final=None, u_pred=None：第一次调用 diffuse_2dconv 时，这俩是 None；真正的 u_pred 会在在线推理里由初始化模型产生。

  • u_final：作为 u_target，可以 从另外一条轨迹的状态拿来作为 target（代码实现里是从 db['u'] 里拿轨迹状态），也可以 全是 0 全是 1（根据你的控制目标来设置）。

  • u_pred：是当前从扩散模型中 采样出的未来控制序列 还没去物理演化（solve）的 [u_τ, u_{τ+1}, ..., u_{τ+H-1}]。u_pred 就是下图中反对角线上的模型采样出来的 x，即下面 anti_diagonal 函数实现的功能。
    

def anti_diagonal(pred):  
    return torch.cat([
        pred[len(pred) - 1 - t][:, t].unsqueeze(1)  # len(pred) == T, 取出扩散迭代第 k = T-1−t 步的输出块
        for t in range(len(pred))   # 对于每个预测时刻 t，
    ], dim=1)  # 在 dim=1 维度上拼接

anti_diagonal 函数把模型在整个扩散（denoising）过程里针对“不同预测时刻”给出的中间输出，按“反对角线”方式抽取出来，拼成一个随扩散迭代推进的“预测轨迹”。

采样流程中，.sample() 会返回一个张量 x（这里传入 anti_diagonal 的就是这个 x），对于扩散第 k 步 (k=0,1,…,T-1)，模型会对“预测 0…T-1 共 T 个未来时间点”都给出一份中间估计，存储在 pred[k]。反对角线方向 k + t = T-1 给出了从一开始全噪声，到最终全去噪，每一步都对应一个“预测时刻 t = T-1−k” 的估计结果。

异步

「窗口左移一格（弹出队头）＋尾部再插入一帧新预测」的 在线滚动机制：

def diffuse_2dconv(args, custom_metric, model_i, model_f, seed=0, ret_ls=False, **kwargs):
	def anti_diagonal(pred):  
        return torch.cat([pred[len(pred) - 1 - t][:,t].unsqueeze(1) for t in range(len(pred))], dim=1)
   
    u_from_x = lambda x: x[:, 0, :16, :]
    u0_from_x = lambda x: x[:, 0, 0, :]
    f_from_x = lambda x: x[:, 1, :15, :]

	db = torch.load(args.test_target)
    target_step = db['u'][-50:].cuda()*0.1  # (num_sequences, Nt, Nx) -> (50, Nt, Nx)
	
	total_time_steps = 80  # total_time_steps 是整个推理过程要走的步数
	
	# 加载数据 & 模型 ...
	ddpm_init = load_2dconv_model(model_i, args)  # is_init_model=True
	if args.asynch_inference_mode:
		# —— Open‐loop 初始化 —— 
		args.is_init_model = False
		ddpm_feedback = load_2dconv_model(model_f, args)
		
		kwargs['u_final'] = target_step[:, :15] # 第一轮我们用首 15 步的真值做“末态”条件，即 0,1,2,...,14
		x = ddpm_init.sample(**kwargs)   # x.shape = (B, C, frames+1, X)
		kwargs['u_pred'] = anti_diagonal(x) …   # ← 这里就是初始化模型给出的粗轨迹
		
		# —— Closed‑loop 在线反馈 —— 
	    for t in range(total_time_steps):
	    	kwargs['u_final'] = target_step[:, t:t+15]  # 真值滚动窗口
        	x = ddpm_feedback.sample(**kwargs) * RESCALER
        	u_controlled = burgers_numeric_solve_free(kwargs['u_init']*RESCALER, f_from_x(x)[:,[0]], visc=0.01, T=0.1, dt=1e-4, num_t=1)
        	# 计算多种 metric...  eg.
        	ddpm_mse = mse_deviation(x[:,0,1], u_controlled[:,-1], partially_observed=args.partially_observed).cpu()
        	diffused_mse = (u_from_x(x)[:, 1, :] - kwargs['u_final'][:,0]*RESCALER).square().mean(-1) # MSE
			
			# **关键：滚动更新 for next step**
        	kwargs['u_init'] = u_controlled[:,-1]/RESCALER    # next u0
            kwargs['u_pred'] = x[:,:,1:-1]/RESCALER     # next 粗预测轨迹

对于采样出的 x = ddpm_feedback.sample，给初始状态 kwargs['u_init'] 施加的控制是 f_from_x(x)[:,[0]] 只是 15 帧控制中的第一帧作为控制，求解得到控制后的状态即 u_controlled。

关于 ddpm_mse 的计算，x[:,0,1] 和 u_controlled[:,-1] 都是 (batch, Nx) 的 shape，所以才能进行逐点 MSE 计算。

• 从 ddpm_feedback 里采样得到 x 的 shape 是 (batch, channel=3, Nt, Nx)，同训练数据 shape，当写 x[:, 0, 1] 时，实际上是这样做的索引：
  • : ——保留整个 batch 维度
  • 0 ——固定通道为第 0 通道
  • 1 ——固定时间步为第 1 帧
  • 空下一维——自动对剩下那一维（也就是空间维度 Nx）全部取出

所以 x[:, 0, 1] 等价于 x[:, 0, 1, :]，它会返回一个 shape 为 (batch, Nx) 的张量──时间维被切掉了，只剩下空间维。

• 而 u_controlled 由 burgers_numeric_solve_free 生成，假设 burgers_numeric_solve_free 在 num_t=1 时返回 shape (batch, num_t+1, Nx)（含初始时刻），那么 u_controlled[:, -1] 就是取它最后一个时间点的数据，也是一个 (batch, Nx) 的张量。

for t in range(0, total_time_steps):
    kwargs['u_final'] = target_step[:, t:t+15]  # 每次滑动 window
    # 再用新的 15 帧条件做下一次采样／更新

	...
	kwargs['u_init'] = u_controlled[:,-1]/RESCALER    # next u0
    kwargs['u_pred'] = x[:,:,1:-1]/RESCALER     # next 粗预测轨迹

用 frames=15 来切出一个 15 帧的子序列，在整个 Nt 长度上做滑动，分别做多次有条件扩散。这样既保留了全序列信息，也让每次扩散仅专注于固定长度的局部时间窗口。

理解为一个 滑动窗口大小为 $H=frames$ 帧，在 $T$ 步的噪声强度上开始滑动，每次最左侧的一帧是最清晰的，可以直接作为输出与环境交互，然后窗口往右移动一帧，再进行一次去噪，这样最左侧的本来是倒数第二步清晰现在变成最清晰的，又可以输出与环境交互了…

• 初始规划（全序列）给出：
  $$u\_pred=[u_0,u_1,\dots ,u_{H-1}][f_0,f_1,\dots ,f_{H-2},0](\text{长度}=H)$$
• 执行了 $f_0$ 后，u_controlled=solved(u0, f0) 这个真实求解后的 u_solved 作为下一时刻的 u_init。
• 弹出队头：
  $$u\_pred\leftarrow [u_1,u_2,\dots ,u_{H-1}][f_1,\dots ,f_{H-2},0](\text{长度}=H-1)$$

此时 u_pred 里没有 $u_0$ 了；只剩下未来要做的 $[u_1\dots u_{H-1}]$。

同步

在 同步（else）分支 里，整个循环大致就是用给定的初始状态 u_init 和模型预测的第 t_ 步的控制量，求 t_+1 这一步的状态，然后把它当成下一步的初始状态；同时定期（每 infer_interval 步）重新采样扩散噪声并更新目标序列。

else:  
    for t in range(0,total_time_steps):
        print('online step:', t)

        t_ = t%args.infer_interval
            
        if t_ == 0:
            print('make new noise')
            kwargs['u_final'] = target_step[:,t:t+15]
            x = ddpm_init.sample(**kwargs) * RESCALER
                    
        x_gt = burgers_numeric_solve_free(x[:,0,t_], f_from_x(x)[:,[t_]], visc=0.01, T=0.1, dt=1e-4, num_t=1)
        u_controlled = burgers_numeric_solve_free(kwargs['u_init']*RESCALER, f_from_x(x)[:,[t_]], visc=0.01, T=0.1, dt=1e-4, num_t=1)
        
        ddpm_mse = mse_deviation(x[:,0,t_+1], x_gt[:,-1], partially_observed=args.partially_observed).cpu()
        
        diffused_mse = (u_from_x(x)[:, t_+1, :] - kwargs['u_final'][:,t_]*RESCALER).square().mean(-1) # MSE

用取模 t % infer_interval 来打上周期，保证 每隔 infer_interval 步做一次完整的扩散重采样，其它步数只用已有的结果。

• infer_interval 控制的是 多少步之后重新采样一次扩散噪声并更新 u_final。

• 当 t_ == 0 时，表示 刚好到了一个新的区间起点，代码才会 x = ddpm_init.sample 重新生成一段新的扩散轨迹 x（包含未来 15 步的控制预测）。

• 在区间内的其他步（t_ = 1,2,…,infer_interval-1），则沿用同一段 x，只是把第 t_ 步的 slice 拿出来做推进和对比，避免每一步都重新调用采样，节省开销。

u_controlled = burgers_numeric_solve_free(
    kwargs['u_init'] * RESCALER,
    f_from_x(x)[:, [t_]],
    visc=0.01, T=0.1, dt=1e-4, num_t=1
)

• u_init 作为当前状态：上一轮迭代输出的状态，代表“当前时刻”的数值解初值。
• f_from_x(x)[:, [t_]]：从扩散模型生成的张量 x 中提取第 t_ 步的“控制场”。
• 程序会把这一步算出来的最新状态赋回 u_init，即 kwargs['u_init'] = u_controlled[:, -1] / RESCALER，用于下一轮循环。

# 真正的“地面真值”解，也是单步积分，但用的是 x[:,0,t_] 作为初值
x_gt = burgers_numeric_solve_free(
    x[:, 0, t_],
    f_from_x(x)[:, [t_]],
    visc=0.01, T=0.1, dt=1e-4, num_t=1
)
# x_gt 的 shape 是 (batch, num_t+1=2, Nx)，
# 所以 x_gt[:, -1] 就是最后一步（也就是 t+1）的真值状态

异步模式和同步模式在计算 ddpm_mse 上有一点区别，

• 异步模式下，ddpm_mse 始终衡量的是“DDPM 在第 1 步（t=0→1）上的输出”与“受控数值解第 1 步的输出”之间的偏差。
• 同步模式下，ddpm_mse 衡量的是“DDPM 在 t_→t_+1 步”与“用纯数值方法，从 DDPM 的 t_ 步当作初值算出的地面真值”之间的偏差。

异步更侧重一次性生成的首个步质量，同步则侧重整个周期内多步的逐步精度。

扩散

# diffusion_1d.py
class GaussianDiffusion(nn.Module):
    def __init__(
        self,
        model,
        seq_length,
        timesteps = 900,
        frames=15,  # 用来告诉模型“整个时序里，要控制／条件化的那段轨迹一共有 15 帧”
        sampling_timesteps = None,
		...
	):
	super().__init__()
	self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training

	self.num_timesteps = int(timesteps)  # T = num_timesteps 
	self.frames = frames   				 # H = frames 
	self.gap_timesteps = self.num_timesteps // self.frames    # T/H = gap_timesteps 
	
	self.is_init_model = is_init_model
	self.asynchronous = asynchronous

这里的 frames=15 是指 除了初始时刻之外，后面一共要处理（或预测／条件化）的离散时刻数目 是 15。下面推理中从 ddpm 里 sample 得到的 x 的 shape 为 (batch, channel, Nt, Nx)，

• 而 u_from_x = lambda x: x[:, 0, :16, :] 是因为状态（c=0）有 初始帧（t=0） 再加上后面 15 帧状态。
• f_from_x = lambda x: x[:, 1, :15, :] 这个通道（c=1）上承载的，恰恰就是 未来的 15 帧 控制序列。

在 CL-DiffPhyCon 里，为了把「完整序列一次性去噪」改成「每个物理步异步去噪一小段并马上输出」，在 diffusion_1d.py 里给原来的 DiffPhyCon 加了 两个开关参数，并在采样逻辑里做了两套流程切换。

假设 num_timesteps = 1000：扩散过程被切成 1000 步噪声等级；frames = 4：想一次性输出 4 帧轨迹。

1. 训练（前向 q_process + 异步）
   • 随机抽一个基准 $t_{\mathrm{base}}\in [0,250)$。
   • 生成偏移 [0,250,500,750]，得到 $[t_{\mathrm{base}}+0,t_{\mathrm{base}}+250,t_{\mathrm{base}}+500,t_{\mathrm{base}}+750]$。
   • 对同一条真实 4 帧轨迹的每帧，分别在这 4 个 $t$ 上加噪，训练模型学会在不同噪声强度下回到干净状态。
2. 采样（反向 p_process + 异步截帧）
   • 初始化时 result = [ img₀ ]，img₀ 是 t≈∞（最强噪声）的随机图像，对应帧 0。
   • 从 $t=999$（最强噪声）开始，往下去噪到 $t=250$，共运行 750 步逆扩散。
   • 每当 $t$ 是 750、500、250 这三个 整除点时，就把当前去噪后的图像存下，连同“最初的纯噪声”一起，凑够 4 帧：
     $$result=\left[im{g}_{0}im{g}_{1}im{g}_{2}im{g}_3\right]=\left[\underset{\text{帧0：纯噪声}}{\underbrace{t=1000}},t=750,t=500,t=250\right]$$
   • img₁ 即帧 1，就是经过 1000→750 级别去噪 后的样本。返回的 result 是一个长度为 4 的列表，每项又是一个形状为 (B, frames, C) 的张量，相当于 一次采样直接生成了 4 帧 —— 从最乱到最清晰，模拟了一个轨迹的演化。

训练（forward + p_losses）

def asyn_t_seq(self, t):  # 异步噪声时间步 (B,) → (B, frames)
    t = t.unsqueeze(1).expand(-1, self.frames)  # (B,) → (B, frames)，frames: 你要预测的帧数
    t_offset = torch.arange(0, self.num_timesteps, step=self.num_timesteps // self.frames, device=t.device).long()
	# → (frames,)  因为范围是 [0, Nt)，但步长为 Nt/H，那么得到的就是 （frames,)，范围分别为 [0, Nt/H), (Nt/H, 2Nt/H)...
    t = t + t_offset   # 广播相加 → (B, frames)
    return t

def forward(self, img, *args, **kwargs):
    b, c, n, device, seq_length, = *img.shape, img.device, self.seq_length  # img = (B, C, N)
    assert n == seq_length, f'seq length must be {seq_length}'

	# diffusion timestep
    if self.asynchronous: # 先对每个样本随机抽一个 (B,)，再 (B, frames) 每一行对应这一样本不同帧的各自时间步
        t = torch.randint(0, self.gap_timesteps, (b,), device=device).long() # (B,) 个[0, gap_timesteps)范围内
        t = self.asyn_t_seq(t)     # (B,) → (B, frames)
    else:  # 为每个样本随机选择一个噪声时间步, 意味着所有通道／所有帧都是在同一个噪声强度下加噪、预测去噪
        t = torch.randint(0, self.num_timesteps, (b,), device=device).long() # (B,) 个 [0, Nt) 范围内
        
    img = self.normalize(img)
    return self.p_losses(img, t, *args, **kwargs)

• num_timesteps 表示扩散过程中前向／逆向的离散步数，也就是从完全干净到几乎全噪声离散成多少时间步，是扩散模型内部定义的噪声强度分辨率（比如 1000 步噪声等级）；
• frames 表示要生成或训练的物理序列长度，你想对多少帧做预测（比如 4 帧）。

在异步方案里，两者联动：用 gap_timesteps = num_timesteps // frames 作为每帧间隔，把整个噪声尺度均匀分给 frames 帧。

# t_base.unsqueeze(1) → shape (B,1)
# 扩展成 (B,4)，每列都是同一个基准 t
t_matrix = [[ 34,  34,  34,  34],
            [198, 198, 198, 198],
            [ 47,  47,  47,  47],
            [221, 221, 221, 221], ...]       # shape (B,4)

# 然后加上 t_offset → 得到真正的每帧扩散步数
t_async = t_matrix + [0,250,500,750] → 
    [[  34,  284,  534,  784],
     [ 198,  448,  698,  948],
     [  47,  297,  547,  797],
     [ 221,  471,  721,  971], ...]         # shape (B,4)

这样，每个 batch 样本的 4 帧就分别在不同的噪声进程阶段：第 1 帧是在 t≈[0,250) 的弱噪声阶段；第 2 帧是在 t≈[250,500) 的中等噪声阶段；第 3、4 帧则更靠后，噪声更重。

最终 p_losses 会接收一个形状为 同步: t of shape (B,) 或 异步: t of shape (B, frames) 的噪声时间步。

def p_losses(self, x_start, t, noise = None):
    noise = default(noise, lambda: torch.randn_like(x_start)) # 若外部没传入noise，就用标准正态随机噪声 (B,C,N)

    if self.asynchronous:  # t 的 shape=(B, frames)，多拼一个最后一帧的时间步
        t = torch.cat((t,t[:,[-1]]), dim=-1)   # (B, frames) → (B, frames+1)

    x = self.q_sample(x_start = x_start, t = t, noise = noise)
    
    if self.asynchronous:
        t = t[:,0]  # 把 t 取回 shape=(B,) 的基准噪声步（用来计算损失权重）
    
    x_self_cond = None
        if self.self_condition and random() < 0.5:
            with torch.no_grad():
                x_self_cond = self.model_predictions(x, t).pred_x_start
                x_self_cond.detach_()
   
    # is_condition_u0，conditioned_on_residual ...
    if self.is_model_w:  # training p(w|u0, uT)
        x[..., 0, 1:15, :] = 0 # when training p(w | u0, uT), unet does not see u_[1...T-1]
    
    model_out = self.model(x, t, x_self_cond, residual=residual)
	loss = F.mse_loss(model_out, target, reduction = 'none')

如果开启自条件机制，那么调用 model_predictions() 预测 ${\hat{x}}_0$ 作为 x_self_cond。其中，

def model_predictions(self, x, t, x_self_cond = None, residual=None, clip_x_start = False, rederive_pred_noise = False, **kwargs):
	...
	if not self.is_init_model:
	    init_t = t[:,0]  # eg. t=[200, 400, ...], shape=(b,)
	    model_output = self.model(x, init_t, x_self_cond, residual=residual) # 每个样本都用单值 t 预测
	else:  # shape=(b,frames)
	    model_output = self.model(x, t, x_self_cond, residual=residual) # 看全序列每帧的噪声步差异去预测
	...

• is_init_model=False（增量闭环去噪模式）
  • 同步模式：t 原本就是 (B,)，t[:,0] 仍是 (B,)，没变化。
  • 异步模式：t 是 (B, frames)，t[:,0] 就把它压回成一个 (B,)。也就是说，无论你异步地给了每帧不同的噪声步，真正送给模型去预测的却只是各样本的第 1 帧噪声步。
• is_init_model=True（全序列一次性去噪模式）
  • 不做压缩，直接把 (B,frames) 的 t 丢进去，让模型可以看到每帧的噪声步差异。

推理（sample）

def sample(self, batch_size=16, clip_denoised=True, **kwargs):
	if self.is_condition_uT:
        assert 'is_condition_uT' not in kwargs, 'specify this value in the model. not during sampling.'
        assert 'u_final' in kwargs and kwargs['u_final'] is not None
    
    seq_length, channels = self.seq_length, self.channels
    sample_size = (batch_size, channels, seq_length)
    return self.p_sample_loop(sample_size, clip_denoised=clip_denoised, **kwargs)

sample 出来的张量 (batch_size, channels, seq_length) 在维度排列上，和训练时样本的维度 (batch, channels, time_steps) 是一致的。

这需要在创建模型的时候就会把它的 seq_length 参数设成和训练数据里时间维相同的长度，

def get_2d_ddpm(args):
	sim_time_stamps, sim_space_grids = 16, 128 
	ddpm = GaussianDiffusion(u_net, seq_length=(sim_time_stamps, sim_space_grids), ...)
	return ddpm

从推理的入口 sample() 函数开始，选择调用 p_sample_loop() 或者 ddim_sample() 进行逆向扩散采样。is_init_model 在 p_sample_loop() 里的体现是：当 is_init_model=True，整条序列一次性走完所有去噪步；否则 只对新接进来的那一帧噪声做迭代。

is_init_model 决定「全序列一次性去噪」 or「增量闭环去噪」

• is_init_model=True：做 第一次的 open-loop 规划，模型要 从纯噪声里跑出一个完整的 $T$ 步控制序列，等同于传统一次性扩散采样。
• is_init_model=False：进入闭环阶段，每个物理时刻 只增量「去噪几步 → 输出当前控制 → 环境执行 → 再去噪」的循环。

# GaussianDiffusion
def p_sample_loop(self, shape, **kwargs):
    ...
    if self.is_init_model:
        # 1) 初始全序列规划，img 一开始就是一个形状为 [batch, channels, T] 的纯标准高斯噪声
        img = torch.randn(shape, device=device) 
        # 然后模型跑完所有去噪迭代，img → 变成一个“干净”的动作序列
    else:
        # 2) 闭环增量阶段，从上一步输出的 u_pred 接着去噪
        img = kwargs['u_pred']   # 取出外部给的轨迹预测  (B, C, H=frames)
        if img.shape[2] == self.frames:   # 这是第一次从 open-loop（全序列规划）切换到闭环增量
            img = self.normalize(img)  
            pad_img = torch.randn(([img.shape[0],img.shape[1],1,img.shape[3]]), device=device) # 拼初始噪声帧
            img = torch.cat([img[:,:,:], pad_img], dim=2)   # [B, C, H+1]
        elif img.shape[2] < self.frames:   # 这是第二步及以后的真正闭环增量
            # 先拼一个最后真值噪声
            noise_final = torch.randn(([img.shape[0],img.shape[1],1,img.shape[3]]), device=device)
            img = self.normalize(img)
            # 多 pad 一帧随机噪声，仅仅为了对齐输入长度
            pad_img = torch.randn(([img.shape[0],img.shape[1],1,img.shape[3]]), device=device) 
            img = torch.cat([img[:,:,:], noise_final, pad_img], dim=2)  # [B, C, H+1]

    if self.is_init_model:
        if self.asynchronous:  # Open‐loop 初始化 + 异步模式，总共从 t = T‑1, T‑2, …, (T/H) 这段去噪
            denoising_steps = self.num_timesteps - self.gap_timesteps   # T − (T/H)
            lower_step, upper_step = self.gap_timesteps-1, self.num_timesteps # [T/H -1, T]
        else:  # Open‐loop 初始化 + 同步模式， 从 0 到 T 这段去噪
            denoising_steps = self.num_timesteps  # T
            lower_step, upper_step = 0, self.num_timesteps  # [0, T]
    else:
    	# Closed‐loop 增量阶段，无论 async 与否，都只跑最浅的那段，从 0 到 T/H
        denoising_steps = self.num_timesteps  # T
        lower_step, upper_step = 0, self.gap_timesteps  # [0, T/H]
    print("denosing_steps, lower_step, upper_step: ", denoising_steps, lower_step, upper_step)   
    result = [img.permute(0,2,1,3)]    # [B, 3, T] -> [B, T, 3] 
    
    # sample loop
    for t in tqdm(reversed(range(lower_step, upper_step)), total = denoising_steps):
    	if self.is_condition_uT: 
    	...
    	self_cond = x_start if self.self_condition else None
        img_curr, x_start, pred_noise = self.p_sample(img, t, self_cond, residual=residual, **kwargs)
		...
		
		if self.is_init_model:
            if self.asynchronous:
                if t % (self.gap_timesteps) == 0 and t >= self.gap_timesteps:                        
                    result.append(img.permute(0,2,1,3))  # [B, 3, T] -> [B, T, 3]

	if self.is_init_model and self.asynchronous:
        assert self.frames == len(result)
        return result
    else:
        return self.unnormalize(img) 

拼初始噪声帧：扩散模型的逆向过程必须从一个高噪声的状态开始。这里给 img 的时间轴末尾再拼上一帧纯高斯噪声 pad_img（或先是真值加噪 noise_final 再是全噪声），这样合成的 img 就包含了：前面 H 帧（初始预测 u_pred） 和 最后一帧（纯噪声，作为扩散去噪的起点）。

在 Open‑loop + 异步模式 中，整条 $H$ 帧（在代码里 img = randn(shape)）都被当作一个整体的纯噪声初始图送入逆扩散。接着，它会从 $t=T-1$ 去噪到 $t=\frac{T}{H}-1$ 这段，每一步都对全部 H 帧去噪，只是在那些特定的噪声分界 $t=k\cdot \frac{T}{H}$ 处把当前的去噪结果“拍”下来，组成最终的 $H$ 帧序列。

• 整条序列一次性走完所有去噪步：对整个 $H$-帧张量都做了 $T-\frac{T}{H}$ 步的逆扩散。
• 即使是 result 中的最后一帧，对应 $t=\frac{T}{H}$，即最清晰的一帧，也不会跑到 $t=0$，所以“完全干净”并不在这个 loop 里出现。虽然已经比纯噪声清了很多，但仍残留一点噪声。

这里，只有在“初始化模型＋异步模式” 下，采样过程会把每次到达关键时刻的中间产物都放到一个 result 列表里，并最终返回这个列表；其它模式下，就只返回最后一次迭代得到的单一张量 img。

为什么要 img.permute(0,2,1)？

• img.permute(0,2,1) 后，拿到的 result 列表就是若干个形状 [B, T, 3] 的条目，每一项对应一次“关键时间点”的完整 [时间, 通道] 数据。
• 假设一共 frames 步要 append，那么最后拿到的 result = [ tensor([B, T, 3]), tensor([B, T, 3]), … ] # 列表长度 == self.frames，
• 这就和 inference.py中的 kwargs['u_pred'] = anti_diagonal(x).permute(0,2,1) 对应上了，沿着“从第0帧的末端到第T帧的开头”的那条反对角线，把每一帧的预测都抽出来，重新组织成一个 [batch, frames, channels] 的张量，再 .permute(0,2,1) 回去，即 [batch, channels, frames]，就得到了一个和训练时一致、可以直接作为 u_pred 条件输入的新张量。

# GaussianDiffusion
def p_sample(self, x, t: int, x_self_cond=None, residual=None, **kwargs):
    b, *_, device = *x.shape, x.device
    batched_times = torch.full((b,), t, device = x.device, dtype = torch.long) # 生成 [t, t, …], shape (B,)
        
    if not self.is_init_model:
        batched_times = self.asyn_t_seq(batched_times)  
        batched_times = torch.cat((batched_times,batched_times[:,[-1]]), dim=-1)
        
    # 传给 p_mean_variance 的就是一个形状 [batch, H+1] 的时序向量
    model_mean, _, model_log_variance, x_start, pred_noise = self.p_mean_variance(
    	x = x, t = batched_times, x_self_cond = x_self_cond, residual=residual, **kwargs
    )
    noise = torch.randn_like(x) if t > 0 else 0.
    pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
    return pred_img, x_start, pred_noise

• 标准一次性去噪（is_init_model=True），batched_times 是一个长度为 batch 的向量，值全是同一个 $t$，模型会对输入的整条 $H+1$ 帧 按同样的噪声级别 t 去噪 —— 这就是在 open-loop 阶段看到的「一次性全序列去噪」。

• 异步闭环去噪（is_init_model=False），传给 p_mean_variance 的就是一个形状 [batch, H+1] 的时序向量，

  batched_times = self.asyn_t_seq(batched_times)
  batched_times = torch.cat((batched_times, batched_times[:, [-1]]), dim=-1) 
  

  • asyn_t_seq 会把它映射成一个长度-$H$ 的不一样的噪声等级序列 $[t_{\tau },t_{\tau +1},\dots ,t_{\tau +H-1}]$。
  • 再在末尾多复制一位（为了给那帧 pad 噪声也塞一个时间标签），得到 $[t_{\tau },\dots ,t_{\tau +H-1},t_{\tau +H-1}]$。

p_mean_variance 在对输入的 $H+1$ 帧做去噪时，就能按照 每帧不同的噪声强度 来计算 model mean/variance —— 正对应图里 同一条队列里，不同帧跑到不同的噪声强度 的设计。

def p_mean_variance(self, x, t, x_self_cond = None, residual=None, **kwargs):
    preds = self.model_predictions(x, t, x_self_cond, residual=residual, **kwargs)
    x_start = preds.pred_x_start

    if kwargs['clip_denoised']:
        x_start.clamp_(-1., 1.)

    model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)
    return model_mean, posterior_variance, posterior_log_variance, x_start, preds.pred_noise

总结：is_init_model 就是这个切换开关 —— 它决定了 p_sample 是走常规深度去噪，还是走 CL-DiffPhyCon 的「每帧不一样噪声强度」去噪。

• Open-loop（is_init_model=True）：轨迹中所有帧都在同一噪声等级下去噪，只要准备一个标量 t，UNet 接到的时序向量更扁平 (B,)。
• Closed-loop（is_init_model=False）：需要给前后帧分配不同噪声强度，让靠前的帧去得更干净（小噪声强度、深度去噪），靠后的帧保留更多噪声（大噪声强度、浅度去噪）以便快速重新规划，UNet 得到的时间输入为 (B,frames+1)。

asynchronous 决定「在一次性去噪里要不要按 gap 分段拿中间结果」

# GaussianDiffusion
def p_sample_loop(self, shape, **kwargs):
	...
	x_start = None
    for t in tqdm(reversed(range(lower_step, upper_step)), desc = 'sampling loop time step', total = denoising_steps):
        for k in range(self.recurrence_k):
			... 
            img = img.detach() 
            
            if self.is_init_model:  # 做 open-loop 采样
                if self.asynchronous:
                	# 当 t 刚好是 gap_timesteps 的倍数，就把当前 img（也就是当前去噪深度下的控制序列）存下来
                    if t % (self.gap_timesteps) == 0 and t >= self.gap_timesteps:                        
                        result.append(img.permute(0,2,1,3))
			...

model = Diffusion(..., is_init_model=True, asynchronous=True) 就意味着 一边跑 1000 步去噪，一边每隔 gap_timesteps（比如 50）就得到一次控制序列快照。

asynchronous：

• True：在 Open-Loop 的那一大段去噪里，每隔 gap_timesteps 个扩散步，就把“半去噪”的 img（形状 [batch,H, …]）提取出来放进 result。
• False：Open-Loop 只在最后（全部去噪完）才输出一次 img，不做中间截点。

两个开关精细地管理了「什么时候去噪多深」、「要不要多次输出中间结果」和「闭环时只做浅度迭代」的行为。

• Open-loop + 异步：只跑后 $H-1$ 段的去噪，每 $T/H$ 步截一次点，得到多帧快照。
• Open-loop + 同步：跑完全部 $T$ 步，一次性输出。
• Closed-loop：只跑最浅的 $T/H$ 步，迅速输出当前控制——这正是 CL-DiffPhyCon 每帧增量闭环所需的微步迭代。