基于当前仓库源码整理目标是把 3D Diffusion Policy也就是 DP3讲成一条完整、容易理解的链路专家策略生成 demonstration- 保存成 zarr 轨迹数据- Dataset 从轨迹中切时间窗口- 模型读取最近几帧观测- 条件扩散模型学习生成动作序列- 推理时从随机噪声逐步去噪得到未来动作对应核心源码3D-Diffusion-Policy/train.py3D-Diffusion-Policy/diffusion_policy_3d/policy/dp3.py3D-Diffusion-Policy/diffusion_policy_3d/model/vision/pointnet_extractor.py3D-Diffusion-Policy/diffusion_policy_3d/model/diffusion/conditional_unet1d.py3D-Diffusion-Policy/diffusion_policy_3d/dataset/*.py3D-Diffusion-Policy/diffusion_policy_3d/common/sampler.py3D-Diffusion-Policy/diffusion_policy_3d/common/replay_buffer.py默认例子主要用 dp3 adroit_hammer 说明。DP3 到底想解决什么问题机器人控制里最朴素的问题是机器人当前看到了什么机器人当前自己的关节/末端状态是什么下一步应该怎么动传统行为克隆可以训练一个模型obs[t] - action[t]但是机器人操作通常不是单步动作能解决的。比如开门、敲钉子、按按钮都需要一小段连续动作并且未来动作之间要协调。DP3 的核心想法是给定最近几帧 3D 观测和机器人状态生成一段动作轨迹而不是只生成一个动作。DP3 里的 3D 观测主要是点云point_cloud: [N, 3] 或 [N, 6]其中[N, 3] 表示每个点有 x, y, z[N, 6] 表示每个点有 x, y, z, r, g, b默认配置里use_pc_color: false所以模型只使用 xyz不使用 rgb。训练数据从哪里来训练数据不是人工一帧一帧标注出来的而是来自专家策略的 demonstration。例如bash scripts/gen_demonstration_adroit.sh hammer这个脚本会创建 Adroit hammer 仿真环境。加载已经训练好的专家策略。让专家策略在环境里执行动作。每一步保存观测和动作。只保留成功 episode。把数据保存成 zarr 格式。不同任务的专家来源不同Adroit - VRL3 expert checkpointDexArt - PPO expert checkpointMetaWorld - scripted policyRealDex - 真实机器人采集数据转换生成后的数据路径类似3D-Diffusion-Policy/data/adroit_hammer_expert.zarr3D-Diffusion-Policy/data/metaworld_basketball_expert.zarr3D-Diffusion-Policy/data/dexart_laptop_expert.zarr3. zarr 训练数据格式zarr 数据可以理解成一个磁盘上的数组数据库。项目里每个 demonstration 数据集大致长这样adroit_hammer_expert.zarr/data/imgpoint_clouddepthstateactionmeta/episode_ends核心字段data/point_cloud 每一步的点云观测data/state 每一步的机器人状态data/action 每一步专家动作meta/episode_ends 每条 episode 的结束位置注意所有 episode 会被拼成一条长时间轴。假设有三条成功轨迹episode 0: 100 stepsepisode 1: 120 stepsepisode 2: 80 steps那么 zarr 里不是保存成三个单独数组而是point_cloud.shape [300, N, 3]state.shape [300, D]action.shape [300, A]episode_ends [100, 220, 300]episode_ends 告诉代码episode 0 [0, 100)episode 1 [100, 220)episode 2 [220, 300)这样 SequenceSampler 在切训练样本时就不会把两个 episode 拼错。Dataset 怎么把轨迹变成训练样本以 AdroitDataset 为例训练时会读取ReplayBuffer.copy_from_path(zarr_path,keys[“state”, “action”, “point_cloud”, “img”])然后创建SequenceSampler(replay_bufferself.replay_buffer,sequence_lengthhorizon,pad_beforepad_before,pad_afterpad_after,episode_masktrain_mask)默认配置horizon: 16n_obs_steps: 2n_action_steps: 8任务配置里pad_before:KaTeX parse error: Expected }, got EOF at end of input: {eval:{n_obs_steps}-1’}pad_after:KaTeX parse error: Expected }, got EOF at end of input: {eval:{n_action_steps}-1’}所以默认horizon 16pad_before 1pad_after 7也就是说每个训练样本是从一条专家轨迹中切出来的长度为 16 的时间窗口。单个样本大致是obs.point_cloud: [16, N, 3]obs.agent_pos: [16, D]action: [16, A]经过 DataLoader(batch_sizeB) 后obs.point_cloud: [B, 16, N, 3]obs.agent_pos: [B, 16, D]action: [B, 16, A]以 adroit_hammer 为例配置是point_cloud:shape: [512, 3]agent_pos:shape: [24]action:shape: [26]所以 batch 形状是obs.point_cloud: [B, 16, 512, 3]obs.agent_pos: [B, 16, 24]action: [B, 16, 26]5. 一个容易误解的时间对齐问题Dataset 会给出 16 帧 obs 和 16 步 actionobs: obs[s], obs[s1], …, obs[s15]action: act[s], act[s1], …, act[s15]但模型实际不会使用全部 16 帧 obs。默认只使用前 2 帧obs[s], obs[s1]如果把 obs[s1] 看成当前时刻 t那么模型看到:obs[t-1], obs[t]训练 loss 覆盖:act[t-1], act[t], …, act[t14]推理实际执行:act[t], act[t1], …, act[t7]所以更精确地说训练时模型学习复原一个完整 horizon 的动作轨迹。真正用于控制的是从当前观测时刻开始的一段动作。代码里执行动作的切片是start n_obs_steps - 1end start n_action_stepsaction action_pred[:, start:end]默认start 1end 9所以推理时执行action_pred[:, 1:9]6. DP3 模型整体结构DP3 模型可以拆成两部分DP3Encoder负责把观测编码成条件向量ConditionalUnet1D负责对动作序列做条件扩散去噪整体数据流point_cloud agent_pos|vDP3Encoder|vglobal_cond|vnoisy action trajectory timestep global_cond|vConditionalUnet1D|vpredicted clean action trajectory默认情况下obs_as_global_cond: truecondition_type: filmprediction_type: sample这表示观测被编码成一个全局条件 global_cond。U-Net 通过 FiLM 使用这个条件。模型预测的是干净动作序列本身而不是噪声。7. 观测编码器 DP3Encoder代码位置3D-Diffusion-Policy/diffusion_policy_3d/model/vision/pointnet_extractor.pyDP3Encoder 输入point_cloud: [B, To, N, 3]agent_pos: [B, To, D]其中To n_obs_steps默认 To2。在进入 encoder 前代码会把时间维和 batch 维合并point_cloud: [B, 2, N, 3] - [B2, N, 3]agent_pos: [B, 2, D] - [B2, D]7.1 点云分支点云走一个简化 PointNet每个点 [x,y,z]- Linear- LayerNorm- ReLU- Linear- LayerNorm- ReLU- Linear- LayerNorm- ReLU- max pooling over points- final projection核心思想先对每个点独立提特征再对所有点做 max pooling得到整个点云的全局特征。伪代码def encode_point_cloud(points):# points: [BTo, N, 3]per_point_feat mlp(points) # [BTo, N, C]global_feat max(per_point_feat, dim1) # [BTo, C]point_feat projection(global_feat) # [BTo, 64]return point_feat默认point_feat: [B*2, 64]7.2 机器人状态分支agent_pos 走 MLPagent_pos: [B2, D] - state_feat: [B2, 64]伪代码def encode_state(agent_pos):return state_mlp(agent_pos)7.3 拼接成观测特征final_feat concat([point_feat, state_feat], dim-1)默认point_feat: [B2, 64]state_feat: [B2, 64]final_feat: [B*2, 128]然后恢复时间信息[B2, 128] - [B, 2128] - [B, 256]这个 [B, 256] 就是 global_cond。ConditionalUnet1D 动作扩散模型代码位置3D-Diffusion-Policy/diffusion_policy_3d/model/diffusion/conditional_unet1d.py它处理的是动作序列sample: [B, horizon, action_dim]以 adroit_hammer 为例sample: [B, 16, 26]内部会转换成 Conv1D 格式[B, 16, 26] - [B, 26, 16]其中26 是动作维度也就是 channel。16 是动作序列长度也就是时间维。U-Net 结构input noisy action|vdown blocks|vmiddle blocks|vup blocks skip connections|voutput action-shaped tensor默认通道变化26 - 512 - 1024 - 2048 - 1024 - 512 - 26这里的 U-Net 不是图像 U-Net而是 1D temporal U-Net。它沿动作序列时间维做卷积学习一个动作 chunk 内部的时间结构。条件是怎么注入 U-Net 的U-Net 每个 residual block 都接收一个条件向量condition timestep_embedding global_cond其中timestep_embedding: 当前扩散噪声等级的编码global_cond: 观测编码器输出的场景和机器人状态特征默认 condition_type: film。FiLM 的形式是condition - scale, biasfeature - scale * feature bias伪代码def conditional_residual_block(x, cond):h conv_block_1(x)scale, bias cond_encoder(cond).chunk(2, dimchannel) h scale * h bias h conv_block_2(h) return h residual(x)直观理解同一个去噪网络在不同场景、不同机器人状态下会通过 scale 和 bias 改变自己的行为。10. Diffusion 训练目标是什么普通行为克隆可能直接学obs - actionDP3 学的是obs noisy_action_trajectory noise_level - clean_action_trajectory也就是说训练时模型不是从零直接输出动作而是学习一个去噪任务。把专家动作序列记作x0 expert action trajectory例如x0: [B, 16, 26]随机采样一个扩散时间步 tt: [B]采样高斯噪声epsilon: [B, 16, 26]根据 scheduler 加噪xt add_noise(x0, epsilon, t)模型学习model(xt, t, obs_condition) - x0当前默认 prediction_type: sample所以模型预测的就是 x0 本身。compute_loss 逐行解释核心函数def compute_loss(self, batch):11.1 输入 batchbatch[“obs”][“point_cloud”]: [B, 16, N, 3]batch[“obs”][“agent_pos”]: [B, 16, D]batch[“action”]: [B, 16, A]11.2 归一化nobs self.normalizer.normalize(batch[“obs”])nactions self.normalizer[“action”].normalize(batch[“action”])归一化后nactions: [B, 16, A]归一化的原因点云坐标、关节状态、动作的数值范围不同。扩散模型对数值尺度敏感所以先把不同字段拉到合适范围。11.3 只取前 n_obs_steps 帧观测this_nobs dict_apply(nobs,lambda x: x[:, :self.n_obs_steps, …].reshape(-1, *x.shape[2:]))默认 n_obs_steps2point_cloud: [B, 16, N, 3] - [B, 2, N, 3] - [B2, N, 3]agent_pos: [B, 16, D] - [B, 2, D] - [B2, D]11.4 编码观测nobs_features self.obs_encoder(this_nobs)默认输出nobs_features: [B*2, 128]再 reshapeglobal_cond nobs_features.reshape(batch_size, -1)得到global_cond: [B, 256]11.5 准备 clean trajectorytrajectory nactions也就是trajectory x0 normalized expert action trajectorytrajectory: [B, 16, A]11.6 生成 condition maskcondition_mask self.mask_generator(trajectory.shape)默认 obs_as_global_condTrue 时trajectory 里只有 action没有把 obs 拼进去。LowdimMaskGenerator 基本不会把 action 维 mask 成已知条件。所以这里可以近似理解成loss 会覆盖整个动作轨迹。11.7 随机采样噪声和扩散时间步noise torch.randn(trajectory.shape, devicetrajectory.device)timesteps torch.randint(0,self.noise_scheduler.config.num_train_timesteps,(batch_size,),devicetrajectory.device