3. MiniCPM-Robot 中action head 前向传播

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)

  1. 外层推理 (predict_action)
  2. 单步去噪前向 (_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. 准备与时间嵌入

  1. 输入 (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