DP3 核心原理详解:从训练数据到模型、Loss 和推理
基于当前仓库源码整理,目标是把 3D Diffusion Policy,也就是 DP3,讲成一条完整、容易理解的链路:
专家策略生成 demonstration
-> 保存成 zarr 轨迹数据
-> Dataset 从轨迹中切时间窗口
-> 模型读取最近几帧观测
-> 条件扩散模型学习生成动作序列
-> 推理时从随机噪声逐步去噪得到未来动作
对应核心源码:
3D-Diffusion-Policy/train.py
3D-Diffusion-Policy/diffusion_policy_3d/policy/dp3.py
3D-Diffusion-Policy/diffusion_policy_3d/model/vision/pointnet_extractor.py
3D-Diffusion-Policy/diffusion_policy_3d/model/diffusion/conditional_unet1d.py
3D-Diffusion-Policy/diffusion_policy_3d/dataset/*.py
3D-Diffusion-Policy/diffusion_policy_3d/common/sampler.py
3D-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 checkpoint
DexArt -> PPO expert checkpoint
MetaWorld -> scripted policy
RealDex -> 真实机器人采集数据转换
生成后的数据路径类似:
3D-Diffusion-Policy/data/adroit_hammer_expert.zarr
3D-Diffusion-Policy/data/metaworld_basketball_expert.zarr
3D-Diffusion-Policy/data/dexart_laptop_expert.zarr
3. zarr 训练数据格式
zarr 数据可以理解成一个磁盘上的数组数据库。项目里每个 demonstration 数据集大致长这样:
adroit_hammer_expert.zarr/
data/
img
point_cloud
depth
state
action
meta/
episode_ends
核心字段:
data/point_cloud 每一步的点云观测
data/state 每一步的机器人状态
data/action 每一步专家动作
meta/episode_ends 每条 episode 的结束位置
注意:所有 episode 会被拼成一条长时间轴。
假设有三条成功轨迹:
episode 0: 100 steps
episode 1: 120 steps
episode 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_buffer=self.replay_buffer,
sequence_length=horizon,
pad_before=pad_before,
pad_after=pad_after,
episode_mask=train_mask
)
默认配置:
horizon: 16
n_obs_steps: 2
n_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 = 16
pad_before = 1
pad_after = 7
也就是说,每个训练样本是从一条专家轨迹中切出来的长度为 16 的时间窗口。
单个样本大致是:
obs.point_cloud: [16, N, 3]
obs.agent_pos: [16, D]
action: [16, A]
经过 DataLoader(batch_size=B) 后:
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 步 action:
obs: obs[s], obs[s+1], …, obs[s+15]
action: act[s], act[s+1], …, act[s+15]
但模型实际不会使用全部 16 帧 obs。默认只使用前 2 帧:
obs[s], obs[s+1]
如果把 obs[s+1] 看成当前时刻 t,那么:
模型看到:
obs[t-1], obs[t]
训练 loss 覆盖:
act[t-1], act[t], …, act[t+14]
推理实际执行:
act[t], act[t+1], …, act[t+7]
所以更精确地说:
训练时模型学习复原一个完整 horizon 的动作轨迹。
真正用于控制的是从当前观测时刻开始的一段动作。
代码里执行动作的切片是:
start = n_obs_steps - 1
end = start + n_action_steps
action = action_pred[:, start:end]
默认:
start = 1
end = 9
所以推理时执行:
action_pred[:, 1:9]
6. DP3 模型整体结构
DP3 模型可以拆成两部分:
DP3Encoder
负责把观测编码成条件向量
ConditionalUnet1D
负责对动作序列做条件扩散去噪
整体数据流:
point_cloud + agent_pos
|
v
DP3Encoder
|
v
global_cond
|
v
noisy action trajectory + timestep + global_cond
|
v
ConditionalUnet1D
|
v
predicted clean action trajectory
默认情况下:
obs_as_global_cond: true
condition_type: film
prediction_type: sample
这表示:
观测被编码成一个全局条件 global_cond。
U-Net 通过 FiLM 使用这个条件。
模型预测的是干净动作序列本身,而不是噪声。
7. 观测编码器 DP3Encoder
代码位置:
3D-Diffusion-Policy/diffusion_policy_3d/model/vision/pointnet_extractor.py
DP3Encoder 输入:
point_cloud: [B, To, N, 3]
agent_pos: [B, To, D]
其中:
To = n_obs_steps
默认 To=2。
在进入 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, dim=1) # [BTo, C]
point_feat = projection(global_feat) # [BTo, 64]
return point_feat
默认:
point_feat: [B*2, 64]
7.2 机器人状态分支
agent_pos 走 MLP:
agent_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
|
v
down blocks
|
v
middle blocks
|
v
up blocks + skip connections
|
v
output 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, bias
feature -> scale * feature + bias
伪代码:
def conditional_residual_block(x, cond):
h = conv_block_1(x)
scale, bias = cond_encoder(cond).chunk(2, dim=channel)
h = scale * h + bias
h = conv_block_2(h)
return h + residual(x)
直观理解:
同一个去噪网络,在不同场景、不同机器人状态下,会通过 scale 和 bias 改变自己的行为。
10. Diffusion 训练目标是什么
普通行为克隆可能直接学:
obs -> action
DP3 学的是:
obs + noisy_action_trajectory + noise_level -> clean_action_trajectory
也就是说,训练时模型不是从零直接输出动作,而是学习一个去噪任务。
把专家动作序列记作:
x0 = expert action trajectory
例如:
x0: [B, 16, 26]
随机采样一个扩散时间步 t:
t: [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 输入 batch
batch[“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_steps=2:
point_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]
再 reshape:
global_cond = nobs_features.reshape(batch_size, -1)
得到:
global_cond: [B, 256]
11.5 准备 clean trajectory
trajectory = nactions
也就是:
trajectory = x0 = normalized expert action trajectory
trajectory: [B, 16, A]
11.6 生成 condition mask
condition_mask = self.mask_generator(trajectory.shape)
默认 obs_as_global_cond=True 时,trajectory 里只有 action,没有把 obs 拼进去。LowdimMaskGenerator 基本不会把 action 维 mask 成已知条件。
所以这里可以近似理解成:
loss 会覆盖整个动作轨迹。
11.7 随机采样噪声和扩散时间步
noise = torch.randn(trajectory.shape, device=trajectory.device)
timesteps = torch.randint(
0,
self.noise_scheduler.config.num_train_timesteps,
(batch_size,),
device=trajectory.device
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐

所有评论(0)