模仿学习算法——ATC / Diffusion Policy

模仿学习算法——ACT / Diffusion Policy

最近又开始接触具身智能相关的一些工作，本文主要介绍两篇较为基础的用来做模仿学习的算法，后续有空再进一步学习分享RT系列和Pi0等文章

1. ACT

ACT算法全称为Action Chunking with Transformers，发表于《Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware》，该论文主要内容为如何使用低成本的硬件和模仿学习算法来实的双手操作，我们主要对算法部分进行剖析

1.1 网络结构及特点

ACT算法的网络结构如下图所示，使用的采用的是生成式模型VAE（变分自动编码器）的方式训练模型，由左侧的Encoder部分和右侧的Decoder部分两部分构成：

这里我们通过类比标准VAE结构来理解ACT算法的原理：
在VAE中，

1. Encoder负责提取隐变量$z$，通常输入是数据样本（如图像），然后编码为隐变量$z$；
2. Decoder仅使用隐变量$z$来生成输出（如图像重建）；

在ACT里，

1. Encoder处理的是Action Sequence，并且额外加入了机器人关节状态（Joints）和CLS 标记（Classification Token, CLS），机器人关节状态的作用是是提供机器人当前的运动状，增强动作的稳定性，而CLS 标记的作用是作为Token在Self-Attention层中聚合整个输入序列的信息来构建隐变量$z$。
2. Decoder除了使用隐变量$z$来生成预测的Action Sequence，额外输入还有当前图像特征和机器人关节状态（Joints），图像信息的作用是提供任务的实时环境信息和增强任务的泛化能力，关节状态（Joints）的作用则是提供机器人当前的运动状态增强时序稳定性。

模型训练和推理的具体过程如下图所示：
训练过程：

训练过程中指的注意的点是：

1. 编码器$q_{\varphi }\left(z\mid a_{t:t+k},{\bar{o}}_t\right)$输出的隐变量$z$是一个满足高斯分布的参数，即$z\sim N\left(\mu ,{\sigma }^2\right)$，其中$\mu$和$\sigma$分别是隐变量$z$的均值和方差，直接从$N\left(\mu ,{\sigma }^2\right)$采样$z$是不可微的，因此采用重参数技巧，$$z=\mu +\sigma \cdot ϵ,ϵ\sim \mathcal{N}(0,I)$$这里$ϵ$是从标注正态分布$N\left(\mu ,{\sigma }^2\right)$采样的噪声项。采样的隐变量$z$被输入编码器${\pi }_{\theta }\left(a_{t:t+k}\mid o_t,z\right)$，用于预测$k$个时间步的动作。
2. 损失函数由两部分构成，一部分是动作序列的L1 Loss，另一部分是KL散度损失，其中KL散度损失定义如下：$$L_{\mathrm{reg}}=D_{\mathrm{KL}}\left(q_{\varphi }\left(z\mid a_{t:t+k},{\bar{o}}_t\right)\Vert \mathcal{N}(0,I)\right)$$其主要作用是在训练时$q_{\varphi }(z)$不会过度拟合演示数据，生成的$z$会形成一个平滑、连续的空间，保持足够的通用性，使得解码器能够适应不同风格的演示。如果没有KL散度损失，那么可能会导致这会导致推理时，输入图片略有不同，解码器可能会输出完全不同的动作，导致不稳定的机器人行为。

推理过程：

推理过程中值得注意的点是：

1. 在推理阶段，我们将隐变量$z$设为$0$，目的是使解码器输出一个最符合训练数据的平均策略，确保机器人执行稳定、可预测的动作，而不是带有风格变化的随机动作。
2. 我们可以通过通过隐变量$z$控制多模态输出，但是ACT 通过 隐变量$z$控制不同的演示风格（比如快/慢执行任务、不同力度），但它不会在同一个输入图像输入下预测多种可能的未来轨迹，因此其多模态性是有限制的。

这里我们可以来讨论下隐变量$z$的作用？如果去掉隐变量$z$，Decoder 直接输入图像特征和 Joints，是否可以输出预测的 Action？

1. 从直观的理解上，隐变量$z$能学习到的是全局结构信息，主要是因为它在训练过程中被强制去表示数据的高层特征，并且它的设计使其成为一个信息瓶颈，迫使模型从输入数据中提取全局模式，而不仅仅是局部依赖关系。假设你想写一篇小说：
   没有隐变量：你只是根据前面的句子，机械地接着写下一句话。这样可能导致文章缺乏风格和全局连贯性。
   有隐变量：你先决定小说的整体风格（比如是浪漫爱情、科幻还是侦探小说）。然后，每写一句话的时候，都会受到这个高层次设定的引导，使得整篇文章更加连贯。
2. 从数据学上隐变量$z$是一个 低维的潜在表示，它的作用类似于对输入数据$x$进行信息压缩，使得：$$p(z\mid x)=q_{\varphi }(z\mid x)$$即隐变量$z$是输入数据$x$的一个高层次抽象表示。Decoder 根据$z$生成输出：$$p(x\mid z)$$由于$z$的维度比原始数据低，它只能保留最关键信息，而不是所有细节。因此$z$会学习到全局模式，而不是局部信息。
3. 我们完全是可以去隐变量$z$的，如果 Transformer Encoder 足够强大，它可以直接学习输入特征（图像 + 关节信息）到动作预测之间的关系，而不需要额外的隐变量$z$，类似于 Seq2Seq 模型，但是可能导致丢失潜在的全局信息，使得的模型容易过拟合数训练数据，无法生成更丰富的动作变化。

1.2 代码解析

ACT的代码是比较简洁的，readme.txt中有说明各个脚本的作用：

1. imitate_episodes.py，训练和评估ACT；
2. policy.py，对不同Policy的封装；
3. detr，ACT的模型定义修改自DETR；
4. sim_env.py，通过关节空间控制的Mujoco仿真环境；
5. ee_sim_env.py，通过末端执行器空间控制的Mujoco仿真环境；
6. scripted_policy.py，模拟环境的脚本化策略，通过该脚本可以与Mujoco仿真环境进行交互
7. constants.py，跨文件共享的常量；
8. utils.py，数据加载和辅助函数，DataLoader定义在这个文件中；
9. visualize_episodes.py，保存.hdf5数据集中的视频；

这里我们主要对数据加载、模型定义和仿真环境三部分代码进行解析：
DataSet定义如下：

class EpisodicDataset(torch.utils.data.Dataset):
    """情节式数据集类
    
    用于加载和处理机器人操作的情节数据。每个情节包含状态、动作和图像数据。
    
    Args:
        episode_ids: 情节ID列表
        dataset_dir: 数据集目录路径
        camera_names: 相机名称列表
        norm_stats: 归一化统计信息字典
    """
    def __init__(self, episode_ids, dataset_dir, camera_names, norm_stats):
        super(EpisodicDataset).__init__()
        self.episode_ids = episode_ids
        self.dataset_dir = dataset_dir
        self.camera_names = camera_names
        self.norm_stats = norm_stats
        self.is_sim = None
        self.__getitem__(0) # 初始化self.is_sim

    def __len__(self):
        """返回数据集中的情节数量"""
        return len(self.episode_ids)

    def __getitem__(self, index):
        """获取指定索引的数据样本
        
        Args:
            index: 样本索引
            
        Returns:
            tuple: (图像数据, 关节位置数据, 动作数据, padding掩码)
        """
        sample_full_episode = False # 硬编码,是否采样完整情节

        # 加载情节数据
        episode_id = self.episode_ids[index]
        dataset_path = os.path.join(self.dataset_dir, f'episode_{episode_id}.hdf5')
        with h5py.File(dataset_path, 'r') as root:
            is_sim = root.attrs['sim']  # 是否为仿真数据
            original_action_shape = root['/action'].shape
            episode_len = original_action_shape[0]
            
            # 确定起始时间步
            if sample_full_episode:
                start_ts = 0
            else:
                start_ts = np.random.choice(episode_len)
                
            # 获取起始时间步的观测数据
            qpos = root['/observations/qpos'][start_ts]  # 关节位置
            qvel = root['/observations/qvel'][start_ts]  # 关节速度
            image_dict = dict()
            for cam_name in self.camera_names:
                image_dict[cam_name] = root[f'/observations/images/{cam_name}'][start_ts]
                
            # 获取从起始时间步开始的所有动作
            if is_sim:
                action = root['/action'][start_ts:]
                action_len = episode_len - start_ts
            else:
                # 真实机器人数据需要特殊处理以对齐时间步
                action = root['/action'][max(0, start_ts - 1):]
                action_len = episode_len - max(0, start_ts - 1)

        self.is_sim = is_sim
        
        # 对动作序列进行padding
        padded_action = np.zeros(original_action_shape, dtype=np.float32)
        padded_action[:action_len] = action
        is_pad = np.zeros(episode_len)
        is_pad[action_len:] = 1

        # 堆叠不同相机的图像
        all_cam_images = []
        for cam_name in self.camera_names:
            all_cam_images.append(image_dict[cam_name])
        all_cam_images = np.stack(all_cam_images, axis=0)

        # 转换为PyTorch张量
        image_data = torch.from_numpy(all_cam_images)
        qpos_data = torch.from_numpy(qpos).float()
        action_data = torch.from_numpy(padded_action).float()
        is_pad = torch.from_numpy(is_pad).bool()

        # 调整图像通道顺序为channel-first
        image_data = torch.einsum('k h w c -> k c h w', image_data)

        # 归一化数据
        image_data = image_data / 255.0  # 图像归一化到[0,1]
        action_data = (action_data - self.norm_stats["action_mean"]) / self.norm_stats["action_std"]  # 动作标准化
        qpos_data = (qpos_data - self.norm_stats["qpos_mean"]) / self.norm_stats["qpos_std"]  # 关节位置标准化

        return image_data, qpos_data, action_data, is_pad

网络定义如下：

class DETRVAE(nn.Module):
    """ 这是一个基于DETR的变分自编码器(VAE)模型 """
    def __init__(self, backbones, transformer, encoder, state_dim, num_queries, camera_names):
        """ 初始化模型
        参数:
            backbones: 用于特征提取的骨干网络,见backbone.py
            transformer: transformer架构的模型,见transformer.py
            encoder: 编码器模型
            state_dim: 机器人状态的维度
            num_queries: 查询的数量,即检测槽的数量。这是DETR在单张图像中可以检测的最大对象数
            camera_names: 相机名称列表
        """
        super().__init__()
        self.num_queries = num_queries
        self.camera_names = camera_names
        self.transformer = transformer
        self.encoder = encoder
        hidden_dim = transformer.d_model
        # 动作预测头和填充预测头
        self.action_head = nn.Linear(hidden_dim, state_dim)
        self.is_pad_head = nn.Linear(hidden_dim, 1)
        self.query_embed = nn.Embedding(num_queries, hidden_dim)

        # 处理视觉和机器人状态输入
        if backbones is not None:
            self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1)
            self.backbones = nn.ModuleList(backbones)
            self.input_proj_robot_state = nn.Linear(14, hidden_dim)
        else:
            # 如果没有backbone,只处理状态输入
            self.input_proj_robot_state = nn.Linear(14, hidden_dim)
            self.input_proj_env_state = nn.Linear(7, hidden_dim)
            self.pos = torch.nn.Embedding(2, hidden_dim)
            self.backbones = None

        # 编码器相关参数
        self.latent_dim = 32  # 潜在变量z的维度
        self.cls_embed = nn.Embedding(1, hidden_dim)  # CLS token的嵌入
        self.encoder_action_proj = nn.Linear(14, hidden_dim)  # 动作序列投影
        self.encoder_joint_proj = nn.Linear(14, hidden_dim)   # 关节状态投影
        self.latent_proj = nn.Linear(hidden_dim, self.latent_dim*2)  # 投影到潜在空间的均值和方差
        # 位置编码表:[CLS], qpos, action_sequence
        self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim))

        # 解码器相关参数
        self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim)  # 潜在变量投影
        self.additional_pos_embed = nn.Embedding(2, hidden_dim)  # 本体感受和潜在变量的位置编码

    def forward(self, qpos, image, env_state, actions=None, is_pad=None):
        """前向传播函数
        参数:
            qpos: [batch, qpos_dim] 机器人关节状态
            image: [batch, num_cam, channel, height, width] 相机图像
            env_state: 环境状态(未使用)
            actions: [batch, seq, action_dim] 动作序列,训练时使用
            is_pad: padding掩码
        """
        is_training = actions is not None  # 判断是训练还是推理
        bs, _ = qpos.shape

        ### 从动作序列获取潜在变量z
        if is_training:
            # 将动作序列投影到嵌入空间,并与CLS token拼接
            action_embed = self.encoder_action_proj(actions)  # (bs, seq, hidden_dim)
            qpos_embed = self.encoder_joint_proj(qpos)  # (bs, hidden_dim)
            qpos_embed = torch.unsqueeze(qpos_embed, axis=1)  # (bs, 1, hidden_dim)
            cls_embed = self.cls_embed.weight  # (1, hidden_dim)
            cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1)  # (bs, 1, hidden_dim)
            encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1)  # (bs, seq+1, hidden_dim)
            encoder_input = encoder_input.permute(1, 0, 2)  # (seq+1, bs, hidden_dim)
            
            # 处理padding掩码
            cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device)  # CLS和关节状态不需要mask
            is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1)  # (bs, seq+1)
            
            # 获取位置编码
            pos_embed = self.pos_table.clone().detach()
            pos_embed = pos_embed.permute(1, 0, 2)  # (seq+1, 1, hidden_dim)
            
            # 通过编码器获取潜在变量
            encoder_output = self.encoder(encoder_input, pos=pos_embed, src_key_padding_mask=is_pad)
            encoder_output = encoder_output[0]  # 只取CLS token的输出
            latent_info = self.latent_proj(encoder_output)
            mu = latent_info[:, :self.latent_dim]  # 均值
            logvar = latent_info[:, self.latent_dim:]  # 对数方差
            latent_sample = reparametrize(mu, logvar)  # 重参数化采样
            latent_input = self.latent_out_proj(latent_sample)
        else:
            # 推理时从标准正态分布采样
            mu = logvar = None
            latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32).to(qpos.device)
            latent_input = self.latent_out_proj(latent_sample)

        # 处理视觉和状态输入
        if self.backbones is not None:
            # 提取图像特征和位置编码
            all_cam_features = []
            all_cam_pos = []
            for cam_id, cam_name in enumerate(self.camera_names):
                features, pos = self.backbones[0](image[:, cam_id])  # 使用第一个backbone处理所有相机
                features = features[0]  # 取最后一层特征
                pos = pos[0]
                all_cam_features.append(self.input_proj(features))
                all_cam_pos.append(pos)
            # 处理本体感受特征
            proprio_input = self.input_proj_robot_state(qpos)
            # 将相机维度展开到宽度维度
            src = torch.cat(all_cam_features, axis=3)
            pos = torch.cat(all_cam_pos, axis=3)
            hs = self.transformer(src, None, self.query_embed.weight, pos, latent_input, proprio_input, self.additional_pos_embed.weight)[0]
        else:
            # 只处理状态输入
            qpos = self.input_proj_robot_state(qpos)
            env_state = self.input_proj_env_state(env_state)
            transformer_input = torch.cat([qpos, env_state], axis=1)  # seq length = 2
            hs = self.transformer(transformer_input, None, self.query_embed.weight, self.pos.weight)[0]
            
        # 预测动作和padding
        a_hat = self.action_head(hs)
        is_pad_hat = self.is_pad_head(hs)
        return a_hat, is_pad_hat, [mu, logvar]

仿真环境相关的代码主要集中在scripted_policy.py、sim_env.py和ee_sim_env.py三个文件中，其中在仿真环境中物体拾取和传递策略类代码（当然还有一些其他的策略类）如下：

class PickAndTransferPolicy(BasePolicy):
    """
    物体拾取和传递策略类
    继承自BasePolicy基类,实现了双臂机器人拾取物体并在手臂之间传递的轨迹生成
    """
    def generate_trajectory(self, ts_first):
        """
        生成双臂机器人的运动轨迹
        
        参数:
            ts_first: 初始时刻的时间步,包含机器人和环境的初始状态信息
        
        生成的轨迹包含:
        - 左臂轨迹: 等待 -> 移动到交接位置 -> 抓取物体 -> 移动到左侧
        - 右臂轨迹: 等待 -> 接近物体 -> 抓取 -> 移动到交接位置 -> 释放物体 -> 移动到右侧
        """
        # 获取初始机械臂末端执行器的位姿
        init_mocap_pose_right = ts_first.observation['mocap_pose_right']
        init_mocap_pose_left = ts_first.observation['mocap_pose_left']

        # 获取物体的位置和姿态信息
        box_info = np.array(ts_first.observation['env_state'])
        box_xyz = box_info[:3]  # 物体位置
        box_quat = box_info[3:] # 物体姿态四元数

        # 计算抓取时的夹爪姿态(右手)
        gripper_pick_quat = Quaternion(init_mocap_pose_right[3:])
        gripper_pick_quat = gripper_pick_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-60)

        # 计算交接时左臂的姿态
        meet_left_quat = Quaternion(axis=[1.0, 0.0, 0.0], degrees=90)

        # 定义交接位置
        meet_xyz = np.array([0, 0.5, 0.25])

        # 生成左臂轨迹点序列
        self.left_trajectory = [
            {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # 初始等待
            {"t": 100, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # 接近交接位置
            {"t": 260, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # 移动到交接位置
            {"t": 310, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 0}, # 闭合夹爪
            {"t": 360, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # 向左移动
            {"t": 400, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # 保持位置
        ]

        # 生成右臂轨迹点序列
        self.right_trajectory = [
            {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # 初始等待
            {"t": 90, "xyz": box_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat.elements, "gripper": 1}, # 接近物体
            {"t": 130, "xyz": box_xyz + np.array([0, 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 1}, # 下降
            {"t": 170, "xyz": box_xyz + np.array([0, 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 0}, # 抓取物体
            {"t": 200, "xyz": meet_xyz + np.array([0.05, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 0}, # 接近交接位置
            {"t": 220, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 0}, # 移动到交接位置
            {"t": 310, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 1}, # 释放物体
            {"t": 360, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # 向右移动
            {"t": 400, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # 保持位置
        ]

在仿真环境中调用的方式如下：

policy = PickAndTransferPolicy(inject_noise)
for step in range(episode_len):
    action = policy(ts)
    ts = env.step(action)
    episode.append(ts)
    if onscreen_render:
        plt_img.set_data(ts.observation['images']['angle'])
        plt.pause(0.02)
plt.close()

此外sim_env.py和ee_sim_env.py和中主要是控制方式不同

# sim_env.py
# Action space: [left_arm_qpos (6), left_gripper_positions (1), right_arm_qpos (6), right_gripper_positions (1)]
# 14维，直接控制关节角度

# ee_sim_env.py
# Action space: [left_arm_pose (7), left_gripper_positions (1), right_arm_pose (7), right_gripper_positions (1)]
# 16维，控制末端执行器的位置(3)和姿态(4)

1.3 实验结论

这里首先解释下真实数据采集方式：

1. 每个演示（episode）持续 8-14 秒，具体取决于任务的复杂度。
2. 控制频率为 50Hz，意味着每个演示包含 400-700 个时间步（time steps）。
3. 每个任务收集 50 次示范（demonstrations），但 Thread Velcro 任务 由于较高的复杂度，额外收集 100 次示范。
4. 每个任务的演示数据时间总计约 10-20 分钟（不包括重置和操作错误）。考虑到重置任务和操作员的失误，整个数据采集过程的墙上时间（wall-clock time）为 30-60 分钟。

在真实数据训练测试结果如下图所示（右侧两列为两个真实任务）：

ACT比所有SOTA方法成功率都要高，并且在真实数据中成功率几乎都达到90%以上。
另外还有几个比较有价值的消融实验：

图（a）的实验是说明是不同的 k 值（动作块大小）对任务成功率的影响，k=100 时，成功率 44%，显著提高。k=200-400 时，成功率稍微下降，表明过大的 k 可能导致对视觉反馈的依赖降低，影响精确度，结论是适当的 k 值（如 k=100）能减少误差累积，提高任务成功率。

图（b）和图（c）的实验是研究CVAE（Conditional Variational Autoencoder） 是否有助于学习人类演示数据中的多样性和随机性。在图（c）的人类数据上，With CVAE 成功率 35.3%，No CVAE 成功率 仅 2%。说明CVAE 在处理随机性较大的数据时至关重要，否则模型无法学习稳定策略。

图（d）则是研究降低控制频率（5Hz）是否会影响任务成功率和操作效率。结论是高频控制（50Hz）能显著提高任务完成效率，减少误差积累，提高机器人精细操作能力。

以上就是对ACT方法的总结，总体感觉是该方法简洁明了并且效果好，在10分钟的训练数据上就可以达到90%的成功率，后面值得一试。

2. Diffusion Policy

该算法发表于《Diffusion Policy: Visuomotor Policy Learning via Action Diffusion》，基于条件去噪扩散模型来生成机器人动作，其核心思想是不直接预测动作，而是学习动作分布的梯度场，其优势是能够自然表达多模态动作和高维动作序列生成，并且训练过程稳定

2.1 网络结构及特点

Diffusion Policy的模型结构如下图所示：

这里我们同样结合基础的Diffusion算法来理解Diffusion Policy，Diffusion算法模型全称为Denoising Diffusion Probabilistic Model (DDPM) ，DDPM是一种生成模型，其目标是通过逐步去噪来生成数据。算法分为前向扩散过程和逆向去噪过程：

1. 前向扩散过程
   将真实数据$x_0$逐步添加噪声，生成一系列噪声样本$x_k$，直到完全变成高斯噪声，其公式为：$$x_k=\sqrt{{\alpha }_k}{x}_{k-1}+\sqrt{1-{\alpha }_k}ϵ$$其中，$x_k$为第$k$步的样本，${\alpha }_k$为控制噪声强度的超参数，$ϵ\sim \mathcal{N}(0,I)$为标准高斯噪声，经过$k$步加噪后等效公式为：$$x_k=\sqrt{\overline{{\alpha }_k}}{x}_0+\sqrt{1-\overline{{\alpha }_k}}ϵ$$其中，$\overline{{\alpha }_k}={\prod }_{i=1}^k{\alpha }_i$表示累积的噪声参数。经过$K$步前向扩散后，$x_K$将近似服从一个标准高斯分布，即$$q\left(x_K\mid x_0\right)\approx \mathcal{N}(0,I)$$

2. 逆向去噪过程
   通过学习去噪网络，从纯噪声 $x_k$，逐步去噪，最终还原出真实数据 $x_0$。其公式为：$$x_{k-1}=x_k-\gamma \nabla E\left(x_k\right)+\sqrt{2\gamma }ϵ$$其中，$\nabla E\left(x_k\right)$代表能量函数的梯度，即去噪网络需要学习的核心部分，$\gamma$表示学习率，用于控制去噪幅度，等效公式为：$$x_{k-1}=\alpha \left(x_k-\gamma {ϵ}_{\theta }\left(x_k,k\right)\right)+N\left(0,{\sigma }^{2}I\right)$$其中，${ϵ}_{\theta }\left(x_k,k\right)$为去噪网络，用于预测噪声，$N\left(0,{\sigma }^{2}I\right)$是一个额外的高斯噪声项，用于确保去噪过程的随机性，去噪过程相当于“梯度下降”，每一步都在“修正”当前的动作，使其逐步变得合理。

3. 训练过程
   训练上述网络时，目标是最小化一下均方误差：$$L={\mathbb{E}}_{x_0,k,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }\left(x_k,k\right)\Vert }^2\right]$$其中，$ϵ$是真实噪声，${ϵ}_{\theta }\left(x_k,k\right)$是去噪网络的输出，训练的目标是让模型尽可能准确地预测噪声，从而在去噪过程中有效地恢复数据。

以上就介绍了Diffusion算法的基本概念，下面我们来详细看下Diffusion Policy的网络的整理流程：

1. 网络的输入是机器人在时间$t$时，历史行最近的$T_0$步观测数据$O_t$以及初始动作序列$A_t^K$，这些观测数据$O_t$包括图像信息和机器人状态信息，这里需要注意的是Diffusion Policy是时序网络，在论文中提到$T_0$在通常取$2-10$，也就是说每次预测时都会将当前帧和历史若干帧的观测都输入网络，这有助于网络判断物体运动区域、避免短时遮挡误判等问题。 初始动作序列$A_t^K$从高斯噪声中采样：$$A_t^K\sim \mathcal{N}(0,I)$$
2. 网络的输出是未来$T_{\alpha }$步的动作序列$A_t$，这里同样需要注意的是Diffusion Policy采用的是Receding Horizon Control，即一次性预测$T_{\alpha }$的时间步的步骤（通常为8-32步），但是只执行前$m$个动作（通常$m=1$或$m=2$），这样可以保证动作时不断调整的，即使有预测误差累积，也可以在下一次迭代中修正。
3. 网络本身提供了CNN-based和Transformer-based两类去噪网络，其中CNN-based模型适合低速、短时的任务，Transformer-based模型适合快速变化、长时的任务，

下面就两种CNN-based的模型和Transformer-based的模型分别进行展开
CNN-based模型如下图所示

1. 采用的是1D CNN进行动作预测；
2. 采用Feature-wise Linear Modulation (FiLM) 机制，将观测数据$O_t$经过MLP特征提取后，作为条件信号作用于动作序列特征的每一层通道，类似于Batch Normalization，但是允许网络基于输入数据进行调整。

Transformer-based模型如下图所示：

1. 采用的是Transformer进行动作预测；
2. 采用Cross Attention机制，将观测数据$O_t$经过MLP特征提取后和动作序列特征进行Cross Attention$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$这里Query来自于动作序列$A_t^K$，Key和Value来自于观测数据$O_t$
3. 每个时间步的动作$A_t^K$只能关注自己和之前的动作，为了防止未来信息泄漏因此需要特别注意因果掩码的设计。

2.2 代码解析

Diffusion Policy的代码我的主要看的是LeRobot中的源码，这里我们看下Diffusion Policy的推理过程，推理过程如下：

Diffusion Policy的推导过程如下，只有在self._queues[‘action’]为空的时候才进行一次模型推理：

@torch.no_grad
def select_action(self, batch: dict[str, Tensor]) -> Tensor:
     """        
     根据环境观察选择单个动作。
     
     此方法处理观察历史的缓存和由底层扩散模型生成的动作轨迹。工作原理如下：
       - 缓存`n_obs_steps`步的观察（对于前几步，观察会被复制`n_obs_steps`次以填充缓存）。
       - 扩散模型生成`horizon`步的动作。
       - 实际上只保留从当前步开始的`n_action_steps`步的动作用于执行。
     示意图如下：
         ----------------------------------------------------------------------------------------------
         (图例: o = n_obs_steps, h = horizon, a = n_action_steps)
         |时间步             | n-o+1 | n-o+2 | ..... | n     | ..... | n+a-1 | n+a   | ..... | n-o+h |
         |观察是否被使用     | 是    | 是    | 是    | 是    | 否    | 否    | 否    | 否    | 否    |
         |动作是否被生成     | 是    | 是    | 是    | 是    | 是    | 是    | 是    | 是    | 是    |
         |动作是否被使用     | 否    | 否    | 否    | 是    | 是    | 是    | 否    | 否    | 否    |
         ----------------------------------------------------------------------------------------------
     注意，这意味着我们需要：`n_action_steps <= horizon - n_obs_steps + 1`。另外，"horizon"可能不是描述该变量实际含义的最佳名称，
     因为这个周期实际上是从第一个观察开始测量的，而这个观察（如果`n_obs_steps` > 1）发生在过去。
     """
     # 归一化输入
     batch = self.normalize_inputs(batch)
     if self.config.image_features:
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         batch["observation.images"] = torch.stack(
             [batch[key] for key in self.config.image_features], dim=-4
         )
     # 注意：这必须在将图像堆叠到单个键之后发生
     self._queues = populate_queues(self._queues, batch)

     if len(self._queues["action"]) == 0:
         # 从队列中堆叠最新的n个观察
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         # 生成动作序列
         actions = self.diffusion.generate_actions(batch)

         # TODO(rcadene): make above methods return output dictionary?
         # 将归一化的动作转换回原始范围
         actions = self.unnormalize_outputs({"action": actions})["action"]

         # 将动作添加到队列中
         self._queues["action"].extend(actions.transpose(0, 1))

     # 从队列中取出下一个要执行的动作
     action = self._queues["action"].popleft()
     return action

模型推理的代码如下：

def generate_actions(self, batch: dict[str, Tensor]) -> Tensor:
        """       
        生成动作序列。此函数期望`batch`包含：
        {
            "observation.state": (B, n_obs_steps, state_dim) - 机器人状态观察
            
            "observation.images": (B, n_obs_steps, num_cameras, C, H, W) - 图像观察
                和/或
            "observation.environment_state": (B, environment_dim) - 环境状态观察
        }
        
        Args:
            batch: 包含观察数据的批次
            
        Returns:
            生成的动作序列
        """
        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
        assert n_obs_steps == self.config.n_obs_steps

        # 编码图像特征并将它们与状态向量一起连接
        global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)

        # 运行采样
        actions = self.conditional_sample(batch_size, global_cond=global_cond)

        # 提取`n_action_steps`步的动作（从当前观察开始）
        start = n_obs_steps - 1
        end = start + self.config.n_action_steps
        actions = actions[:, start:end]

        return actions

其中_prepare_global_conditioning中对图像特征编码，conditional_sample则执行条件生成的过程，代码如下

def conditional_sample(
        self, batch_size: int, global_cond: Tensor | None = None, generator: torch.Generator | None = None
    ) -> Tensor:
        """
        条件采样方法，生成动作序列
        
        Args:
            batch_size: 批次大小
            global_cond: 全局条件特征
            generator: 随机数生成器
            
        Returns:
            生成的动作序列
        """
        device = get_device_from_parameters(self)
        dtype = get_dtype_from_parameters(self)

        # 采样先验噪声
        sample = torch.randn(
            size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]),
            dtype=dtype,
            device=device,
            generator=generator,
        )

        # 设置噪声调度器的时间步
        self.noise_scheduler.set_timesteps(self.num_inference_steps)

        # 逐步去噪
        for t in self.noise_scheduler.timesteps:
            # 预测模型输出
            model_output = self.unet(
                sample,
                torch.full(sample.shape[:1], t, dtype=torch.long, device=sample.device),
                global_cond=global_cond,
            )
            # 计算前一个图像: x_t -> x_t-1
            sample = self.noise_scheduler.step(model_output, t, sample, generator=generator).prev_sample

        return sample

其中模型结构如下：

class DiffusionConditionalUnet1d(nn.Module):
    """A 1D convolutional UNet with FiLM modulation for conditioning.
    带有FiLM调制用于条件控制的一维卷积UNet。

    Note: this removes local conditioning as compared to the original diffusion policy code.
    注意：与原始扩散策略代码相比，这里移除了局部条件控制。
    """

    def __init__(self, config: DiffusionConfig, global_cond_dim: int):
        super().__init__()

        self.config = config

        # 扩散时间步的编码器
        self.diffusion_step_encoder = nn.Sequential(
            DiffusionSinusoidalPosEmb(config.diffusion_step_embed_dim),
            nn.Linear(config.diffusion_step_embed_dim, config.diffusion_step_embed_dim * 4),
            nn.Mish(),
            nn.Linear(config.diffusion_step_embed_dim * 4, config.diffusion_step_embed_dim),
        )

        # FiLM条件控制的维度
        cond_dim = config.diffusion_step_embed_dim + global_cond_dim

        # UNet编码器中每个下采样块的输入/输出通道。对于解码器，我们只需反转这些
        in_out = [(config.action_feature.shape[0], config.down_dims[0])] + list(
            zip(config.down_dims[:-1], config.down_dims[1:], strict=True)
        )

        # UNet编码器
        common_res_block_kwargs = {
            "cond_dim": cond_dim,
            "kernel_size": config.kernel_size,
            "n_groups": config.n_groups,
            "use_film_scale_modulation": config.use_film_scale_modulation,
        }
        self.down_modules = nn.ModuleList([])
        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (len(in_out) - 1)
            self.down_modules.append(
                nn.ModuleList(
                    [
                        # 第一个残差块，输入通道到输出通道
                        DiffusionConditionalResidualBlock1d(dim_in, dim_out, **common_res_block_kwargs),
                        # 第二个残差块，保持通道数不变
                        DiffusionConditionalResidualBlock1d(dim_out, dim_out, **common_res_block_kwargs),
                        # 如果不是最后一个块，则进行下采样
                        nn.Conv1d(dim_out, dim_out, 3, 2, 1) if not is_last else nn.Identity(),
                    ]
                )
            )

        # 自编码器中间的处理
        self.mid_modules = nn.ModuleList(
            [
                # 两个残差块，保持通道数不变
                DiffusionConditionalResidualBlock1d(
                    config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs
                ),
                DiffusionConditionalResidualBlock1d(
                    config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs
                ),
            ]
        )

        # UNet解码器
        self.up_modules = nn.ModuleList([])
        for ind, (dim_out, dim_in) in enumerate(reversed(in_out[1:])):
            is_last = ind >= (len(in_out) - 1)
            self.up_modules.append(
                nn.ModuleList(
                    [
                        # dim_in * 2，因为它还接收编码器的跳跃连接
                        DiffusionConditionalResidualBlock1d(dim_in * 2, dim_out, **common_res_block_kwargs),
                        # 第二个残差块，保持通道数不变
                        DiffusionConditionalResidualBlock1d(dim_out, dim_out, **common_res_block_kwargs),
                        # 如果不是最后一个块，则进行上采样
                        nn.ConvTranspose1d(dim_out, dim_out, 4, 2, 1) if not is_last else nn.Identity(),
                    ]
                )
            )

        # 最终卷积层
        self.final_conv = nn.Sequential(
            DiffusionConv1dBlock(config.down_dims[0], config.down_dims[0], kernel_size=config.kernel_size),
            nn.Conv1d(config.down_dims[0], config.action_feature.shape[0], 1),
        )

    def forward(self, x: Tensor, timestep: Tensor | int, global_cond=None) -> Tensor:
        """
        Args:
            x: (B, T, input_dim) tensor for input to the Unet.
               输入到UNet的形状为(B, T, input_dim)的张量
            timestep: (B,) tensor of (timestep_we_are_denoising_from - 1).
                      形状为(B,)的张量，表示(我们正在去噪的时间步 - 1)
            global_cond: (B, global_cond_dim)
                         形状为(B, global_cond_dim)的全局条件
            output: (B, T, input_dim)
                    形状为(B, T, input_dim)的输出
        Returns:
            (B, T, input_dim) diffusion model prediction.
            形状为(B, T, input_dim)的扩散模型预测
        """
        # 对于一维卷积，我们需要特征维度在前
        x = einops.rearrange(x, "b t d -> b d t")

        # 编码时间步
        timesteps_embed = self.diffusion_step_encoder(timestep)

        # 如果有全局条件特征，将其与时间步嵌入连接
        if global_cond is not None:
            global_feature = torch.cat([timesteps_embed, global_cond], axis=-1)
        else:
            global_feature = timesteps_embed

        # 运行编码器，跟踪跳跃特征以传递给解码器
        encoder_skip_features: list[Tensor] = []
        for resnet, resnet2, downsample in self.down_modules:
            x = resnet(x, global_feature)
            x = resnet2(x, global_feature)
            encoder_skip_features.append(x)
            x = downsample(x)

        # 处理中间层
        for mid_module in self.mid_modules:
            x = mid_module(x, global_feature)

        # 运行解码器，使用来自编码器的跳跃特征
        for resnet, resnet2, upsample in self.up_modules:
            # 连接跳跃连接
            x = torch.cat((x, encoder_skip_features.pop()), dim=1)
            x = resnet(x, global_feature)
            x = resnet2(x, global_feature)
            x = upsample(x)

        # 最终卷积处理
        x = self.final_conv(x)

        # 将特征维度移回原位置
        x = einops.rearrange(x, "b d t -> b t d")
        return x

2.3 实验结论

本文主要有这么几个实验有借鉴意义

1. 视觉编码器的训练策略
   论文中对视觉编码器采用了三种不同的训练方式，分别是从零开始训练、使用冻结的预训练模型、微调预训练模型，结论是微调预训练模型是最佳方案，采用比 Diffusion Policy 网络低 10 倍的学习率进行微调效果最佳，特别是 CLIP 预训练的 ViT-B/16，微调后仅用 50 轮训练就达到了 98% 成功率。此外，尽管 ResNet-18、ResNet-34 和 ViT-B/16 理论上模型能力差距较大，但它们在任务上的最终性能差距并不显著。
   

2. 动作归一化
   动作数据归一化至 [-1,1]，而不是使用零均值标准化，原始是DDPM在每次去噪迭代中都会将预测结果裁剪到$[-1,1]$，所以使用零均值标准化可能会导致某些动作区域不可访问，影响学习效果。

3. 旋转表示
   在速度控制（Velocity Control）环境中，使用 3D 轴角（Axis-Angle）表示法，原因是由于速度控制的指令通常接近 0，轴角表示的奇异性和不连续性通常不会带来问题。在位置控制（Position Control）环境中，使用 6D 旋转表示，原因是6D 旋转表示相比于四元数更稳定，能够避免 Gimbal Lock（万向锁）问题。

4. 图像增强
   训练时采用随机裁剪（Random Crop）增强，增强方式参考 Mandlekar et al. (2021)。推理时使用固定的中心裁剪（Static Center Crop），以保持输入图像的一致性。图像增强在训练时能提升模型的泛化能力，而推理时使用固定裁剪以确保稳定性。

5. 网络超参数选择
   CNN 版本超参数稳定，增大模型规模通常有助于提升性能，而 Transformer 版本需要针对不同任务微调。此外在 CNN 版本的 Diffusion Policy 中，使用 FiLM 机制进行条件输入效果更好，但是Push-T 任务使用 Imputation 代替 FiLM，效果更佳。

6. 观察时间窗口
   对于状态输入的 Diffusion Policy（State-Based），对观察时间窗口不敏感，模型性能变化不大。但是对于视觉输入的 Diffusion Policy（Vision-Based），CNN 版本的 Diffusion Policy 随着观察窗口的增加，性能下降。Transformer 版本未提及明显下降，可能对观察窗口的影响较小。实验发现，Observation Horizon = 2 对于大多数任务效果最
   

本篇博客只简单总结了ACT和Diffusion Policy两篇Imitation Learning算法的基石工作，后续还有许多衍生工作，比如

1. 3D Diffusion Policy（DP3）的改进主要是将3D点云通过MLP Encoder到网络的输入中，使得任务成功率有明显提升：
   

2. Improved 3D Diffusion Policy（iDP3）的3D DiFfusion Policy的改进主要是（1）是将世界坐标调整成以机器人为中心的坐标系（指的就是相机坐标系），这样做当环境发生变化时，对系统的影响会显著降低。（2）通过提升输入图像的分辨率来提高效果。