基于当前仓库源码整理,目标是把 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 说明。

  1. 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。

  1. 训练数据从哪里来
    训练数据不是人工一帧一帧标注出来的,而是来自专家策略的 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 拼错。

  1. 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] -> [B
2, 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) # [B
To, N, C]
global_feat = max(per_point_feat, dim=1) # [BTo, C]
point_feat = projection(global_feat) # [B
To, 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: [B
2, 64]
final_feat: [B*2, 128]
然后恢复时间信息:

[B2, 128] -> [B, 2128] -> [B, 256]
这个 [B, 256] 就是 global_cond。

  1. 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 内部的时间结构。

  1. 条件是怎么注入 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 本身。

  1. 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] -> [B
2, 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

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐