action_head 前向传播
- python 代码
python
def _predict(
self,
noisy_actions: torch.Tensor,
vl_embs: torch.Tensor,
state: torch.Tensor,
timesteps: torch.Tensor,
embodiment_id: torch.Tensor,
) -> torch.Tensor:
if self.test == 0:
print("---------------test-------------")
print(f"noisy_actions shape: {noisy_actions.shape}")
print(f"state shape: {state.shape}")
print(f"timesteps shape: {timesteps.shape}")
print(f"embodiment_id shape: {embodiment_id.shape}")
features = self._encode_action_tokens(noisy_actions, state, timesteps, embodiment_id)
position_ids = torch.arange(features.shape[1], dtype=torch.long, device=features.device)
features = features + self.position_embedding(position_ids).unsqueeze(0)
output = self.model(
hidden_states=self._build_sequence(features),
encoder_hidden_states=vl_embs,
timestep=timesteps,
)
actions = self.action_decoder(output, embodiment_id)[:, -self.action_horizon :]
if self.test == 0:
print(f"actions shape: {actions.shape}")
self.test = 1
return actions
@torch.no_grad()
def predict_action(
self,
vl_embs: torch.Tensor,
state: torch.Tensor,
embodiment_id: torch.Tensor,
) -> torch.Tensor:
actions = torch.zeros(
(vl_embs.shape[0], self.action_horizon, self.action_dim),
dtype=vl_embs.dtype,
device=vl_embs.device,
)
noise = torch.randn_like(actions)
for step in range(self.num_inference_timesteps, 0, -1):
time = step / float(self.num_inference_timesteps)
timestep = min(int(time * self.num_timestep_buckets), self.num_timestep_buckets - 1)
timesteps = torch.full(
(vl_embs.shape[0],),
timestep,
device=vl_embs.device,
dtype=torch.long,
)
noisy_actions = time * noise + (1 - time) * actions
actions = self._predict(noisy_actions, vl_embs, state, timesteps, embodiment_id)
return actions
- 对应维度说明
bash
---------------test-------------
noisy_actions shape: torch.Size([1, 30, 80])
state shape: torch.Size([1, 1, 80])
timesteps shape: torch.Size([1])
embodiment_id shape: torch.Size([1])
actions shape: torch.Size([1, 30, 80])
- 流程图
#mermaid-svg-r3VvG1HFf0IfSIIg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-r3VvG1HFf0IfSIIg .error-icon{fill:#552222;}#mermaid-svg-r3VvG1HFf0IfSIIg .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-r3VvG1HFf0IfSIIg .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-r3VvG1HFf0IfSIIg .marker{fill:#333333;stroke:#333333;}#mermaid-svg-r3VvG1HFf0IfSIIg .marker.cross{stroke:#333333;}#mermaid-svg-r3VvG1HFf0IfSIIg svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-r3VvG1HFf0IfSIIg p{margin:0;}#mermaid-svg-r3VvG1HFf0IfSIIg .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster-label text{fill:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster-label span{color:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster-label span p{background-color:transparent;}#mermaid-svg-r3VvG1HFf0IfSIIg .label text,#mermaid-svg-r3VvG1HFf0IfSIIg span{fill:#333;color:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg .node rect,#mermaid-svg-r3VvG1HFf0IfSIIg .node circle,#mermaid-svg-r3VvG1HFf0IfSIIg .node ellipse,#mermaid-svg-r3VvG1HFf0IfSIIg .node polygon,#mermaid-svg-r3VvG1HFf0IfSIIg .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-r3VvG1HFf0IfSIIg .rough-node .label text,#mermaid-svg-r3VvG1HFf0IfSIIg .node .label text,#mermaid-svg-r3VvG1HFf0IfSIIg .image-shape .label,#mermaid-svg-r3VvG1HFf0IfSIIg .icon-shape .label{text-anchor:middle;}#mermaid-svg-r3VvG1HFf0IfSIIg .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-r3VvG1HFf0IfSIIg .rough-node .label,#mermaid-svg-r3VvG1HFf0IfSIIg .node .label,#mermaid-svg-r3VvG1HFf0IfSIIg .image-shape .label,#mermaid-svg-r3VvG1HFf0IfSIIg .icon-shape .label{text-align:center;}#mermaid-svg-r3VvG1HFf0IfSIIg .node.clickable{cursor:pointer;}#mermaid-svg-r3VvG1HFf0IfSIIg .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-r3VvG1HFf0IfSIIg .arrowheadPath{fill:#333333;}#mermaid-svg-r3VvG1HFf0IfSIIg .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-r3VvG1HFf0IfSIIg .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-r3VvG1HFf0IfSIIg .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-r3VvG1HFf0IfSIIg .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-r3VvG1HFf0IfSIIg .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-r3VvG1HFf0IfSIIg .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster text{fill:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg .cluster span{color:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-r3VvG1HFf0IfSIIg .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-r3VvG1HFf0IfSIIg rect.text{fill:none;stroke-width:0;}#mermaid-svg-r3VvG1HFf0IfSIIg .icon-shape,#mermaid-svg-r3VvG1HFf0IfSIIg .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-r3VvG1HFf0IfSIIg .icon-shape p,#mermaid-svg-r3VvG1HFf0IfSIIg .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-r3VvG1HFf0IfSIIg .icon-shape .label rect,#mermaid-svg-r3VvG1HFf0IfSIIg .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-r3VvG1HFf0IfSIIg .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-r3VvG1HFf0IfSIIg .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-r3VvG1HFf0IfSIIg :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-r3VvG1HFf0IfSIIg .input>*{fill:#e1f5fe!important;stroke:#311b92!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .input span{fill:#e1f5fe!important;stroke:#311b92!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .process>*{fill:#fff3e0!important;stroke:#e65100!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .process span{fill:#fff3e0!important;stroke:#e65100!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .tensor>*{fill:#e8f5e9!important;stroke:#e65100!important;stroke-width:1px!important;stroke-dasharray:5 5!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .tensor span{fill:#e8f5e9!important;stroke:#e65100!important;stroke-width:1px!important;stroke-dasharray:5 5!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .model>*{fill:#fce4ec!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .model span{fill:#fce4ec!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .output>*{fill:#f1f8e9!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .output span{fill:#f1f8e9!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .loop>*{fill:#f3e5f5!important;stroke:#4a148c!important;stroke-width:2px!important;stroke-dasharray:4 4!important;}#mermaid-svg-r3VvG1HFf0IfSIIg .loop span{fill:#f3e5f5!important;stroke:#4a148c!important;stroke-width:2px!important;stroke-dasharray:4 4!important;} 3. 最终输出
去噪迭代循环 (step = num_inference_timesteps 倒数至 1)
- 外层推理 (predict_action)
- 单步去噪前向 (_predict)
hidden_states
encoder_hidden_states
timestep
更新下一次迭代的 actions
循环继续
循环结束
vl_embs
(来自VLM)
state
1, 1, 80
embodiment_id
1
初始化 actions 为全零
1, 30, 80
生成纯噪声 noise
1, 30, 80
开始单步迭代
计算 time 和 timestep
timesteps shape: 1
计算 noisy_actions
= time * noise + (1-time) * actions
shape: 1, 30, 80
self._encode_action_tokens()
features
- self.position_embedding()
加入位置信息
features (with pos)
self._build_sequence()
hidden_states
self.model()
output (transformer hidden states)
self.action_decoder()
:, -action_horizon :
actions
1, 30, 80
结束单步迭代
最终预测 actions
1, 30, 80
DiT 前向传播
- python 代码
python
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: Optional[torch.LongTensor] = None,
return_all_hidden_states: bool = False,
encoder_attention_mask: Optional[torch.Tensor] = None,
):
print("---------------forward-------------")
print(f"hidden_states shape: {hidden_states.shape}")
print(f"encoder_hidden_states shape: {encoder_hidden_states.shape}")
print(f"timestep shape: {timestep.shape}")
if encoder_attention_mask is not None:
print(f"encoder_attention_mask: {encoder_attention_mask}")
print(f"return_all_hidden_states: {return_all_hidden_states}")
time_embedding = self.timestep_encoder(timestep)
hidden_states = hidden_states.contiguous()
encoder_hidden_states = encoder_hidden_states.contiguous()
all_hidden_states = [hidden_states]
for index, block in enumerate(self.transformer_blocks):
self_attention = index % 2 == 1 and self.config.interleave_self_attention
hidden_states = block(
hidden_states,
encoder_hidden_states=None if self_attention else encoder_hidden_states,
encoder_attention_mask=None if self_attention else encoder_attention_mask,
temb=time_embedding,
)
all_hidden_states.append(hidden_states)
shift, scale = self.proj_out_1(F.silu(time_embedding)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
output = self.proj_out_2(hidden_states)
print(f"output shape: {output.shape}")
return (output, all_hidden_states) if return_all_hidden_states else output
- 对应维度说明
bash
---------------forward-------------
hidden_states shape: torch.Size([1, 62, 768])
encoder_hidden_states shape: torch.Size([1, 88, 1024])
timestep shape: torch.Size([1])
return_all_hidden_states: False
output shape: torch.Size([1, 62, 1024])
- 流程图
#mermaid-svg-G5sp7diBA07Crsnt{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-G5sp7diBA07Crsnt .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-G5sp7diBA07Crsnt .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-G5sp7diBA07Crsnt .error-icon{fill:#552222;}#mermaid-svg-G5sp7diBA07Crsnt .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-G5sp7diBA07Crsnt .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-G5sp7diBA07Crsnt .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-G5sp7diBA07Crsnt .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-G5sp7diBA07Crsnt .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-G5sp7diBA07Crsnt .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-G5sp7diBA07Crsnt .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-G5sp7diBA07Crsnt .marker{fill:#333333;stroke:#333333;}#mermaid-svg-G5sp7diBA07Crsnt .marker.cross{stroke:#333333;}#mermaid-svg-G5sp7diBA07Crsnt svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-G5sp7diBA07Crsnt p{margin:0;}#mermaid-svg-G5sp7diBA07Crsnt .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-G5sp7diBA07Crsnt .cluster-label text{fill:#333;}#mermaid-svg-G5sp7diBA07Crsnt .cluster-label span{color:#333;}#mermaid-svg-G5sp7diBA07Crsnt .cluster-label span p{background-color:transparent;}#mermaid-svg-G5sp7diBA07Crsnt .label text,#mermaid-svg-G5sp7diBA07Crsnt span{fill:#333;color:#333;}#mermaid-svg-G5sp7diBA07Crsnt .node rect,#mermaid-svg-G5sp7diBA07Crsnt .node circle,#mermaid-svg-G5sp7diBA07Crsnt .node ellipse,#mermaid-svg-G5sp7diBA07Crsnt .node polygon,#mermaid-svg-G5sp7diBA07Crsnt .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-G5sp7diBA07Crsnt .rough-node .label text,#mermaid-svg-G5sp7diBA07Crsnt .node .label text,#mermaid-svg-G5sp7diBA07Crsnt .image-shape .label,#mermaid-svg-G5sp7diBA07Crsnt .icon-shape .label{text-anchor:middle;}#mermaid-svg-G5sp7diBA07Crsnt .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-G5sp7diBA07Crsnt .rough-node .label,#mermaid-svg-G5sp7diBA07Crsnt .node .label,#mermaid-svg-G5sp7diBA07Crsnt .image-shape .label,#mermaid-svg-G5sp7diBA07Crsnt .icon-shape .label{text-align:center;}#mermaid-svg-G5sp7diBA07Crsnt .node.clickable{cursor:pointer;}#mermaid-svg-G5sp7diBA07Crsnt .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-G5sp7diBA07Crsnt .arrowheadPath{fill:#333333;}#mermaid-svg-G5sp7diBA07Crsnt .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-G5sp7diBA07Crsnt .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-G5sp7diBA07Crsnt .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-G5sp7diBA07Crsnt .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-G5sp7diBA07Crsnt .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-G5sp7diBA07Crsnt .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-G5sp7diBA07Crsnt .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-G5sp7diBA07Crsnt .cluster text{fill:#333;}#mermaid-svg-G5sp7diBA07Crsnt .cluster span{color:#333;}#mermaid-svg-G5sp7diBA07Crsnt div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-G5sp7diBA07Crsnt .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-G5sp7diBA07Crsnt rect.text{fill:none;stroke-width:0;}#mermaid-svg-G5sp7diBA07Crsnt .icon-shape,#mermaid-svg-G5sp7diBA07Crsnt .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-G5sp7diBA07Crsnt .icon-shape p,#mermaid-svg-G5sp7diBA07Crsnt .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-G5sp7diBA07Crsnt .icon-shape .label rect,#mermaid-svg-G5sp7diBA07Crsnt .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-G5sp7diBA07Crsnt .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-G5sp7diBA07Crsnt .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-G5sp7diBA07Crsnt :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;}#mermaid-svg-G5sp7diBA07Crsnt .input>*{fill:#e1f5fe!important;stroke:#311b92!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .input span{fill:#e1f5fe!important;stroke:#311b92!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .process>*{fill:#fff3e0!important;stroke:#e65100!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .process span{fill:#fff3e0!important;stroke:#e65100!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .tensor>*{fill:#e8f5e9!important;stroke:#e65100!important;stroke-width:1px!important;stroke-dasharray:5 5!important;}#mermaid-svg-G5sp7diBA07Crsnt .tensor span{fill:#e8f5e9!important;stroke:#e65100!important;stroke-width:1px!important;stroke-dasharray:5 5!important;}#mermaid-svg-G5sp7diBA07Crsnt .model>*{fill:#fce4ec!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .model span{fill:#fce4ec!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .output>*{fill:#f1f8e9!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .output span{fill:#f1f8e9!important;stroke:#1b5e20!important;stroke-width:2px!important;}#mermaid-svg-G5sp7diBA07Crsnt .loop>*{fill:#f3e5f5!important;stroke:#4a148c!important;stroke-width:2px!important;stroke-dasharray:4 4!important;}#mermaid-svg-G5sp7diBA07Crsnt .loop span{fill:#f3e5f5!important;stroke:#4a148c!important;stroke-width:2px!important;stroke-dasharray:4 4!important;} 5. 最终输出 (Output)
4. 输出处理 (FiLM 调制与线性投影)
3. Transformer 模块循环 (for block in self.transformer_blocks)
2. 准备与时间嵌入
- 输入 (Inputs)
是: 仅计算自注意力
(encoder_hidden_states=None)
否: 计算交叉注意力
(传入 encoder_hidden_states)
提供视觉语言特征
提供时间条件(temb)
更新
进入下一层
最终隐藏状态
hidden_states
1, 62, 768
encoder_hidden_states
1, 88, 1024
timestep
1
.contiguous() 内存连续化
hidden_states
1, 62, 768
encoder_hidden_states
1, 88, 1024
self.timestep_encoder()
time_embedding
循环开始
index % 2 == 1 且
interleave_self_attention?
block()
更新后的 hidden_states
1, 62, 768
循环结束
F.silu(time_embedding)
self.proj_out_1()
chunk(2, dim=1)
拆分为 shift 和 scale
shift
scale
self.norm_out()
* (1 + scale) + shift
hidden_states (调制后)
self.proj_out_2()
output
1, 62, 1024