GR00T N1.7源码学习（三）：动作头内部模块、DiT结构与多机器人条件编码解析

GR00T N1.7源码学习（一）：工程入口、模型结构与动作生成流程解析-CSDN博客

GR00T N1.7源码学习（二）：训练数据、Processor与多机器人动作空间解析-CSDN博客

前两篇博客已经介绍了GR00T N1.7的整体调用链和训练数据处理流程。这一篇会记录动作头内部的实现，涉及的脚本有3个，

复制代码

gr00t/model/gr00t_n1d7/gr00t_n1d7.py
gr00t/model/modules/embodiment_conditioned_mlp.py
gr00t/model/modules/dit.py

其中gr00t_n1d7.py负责把动作头各个模块串起来，embodiment_conditioned_mlp.py实现多机器人条件相关的MLP和Action Encoder，dit.py实现真正的Transformer动作网络。

1、动作头把状态和动作转换成DiT可以处理的Token

Gr00tN1d7ActionHead内部没有直接把低维状态和动作丢给DiT，而是先把它们映射成Transformer Token。初始化代码中最关键的是这几个模块，

复制代码

class Gr00tN1d7ActionHead(nn.Module):
    """Action head component for flow matching diffusion policy."""

    supports_gradient_checkpointing = True

    def __init__(self, config: Gr00tN1d7Config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.input_embedding_dim = config.input_embedding_dim

        if config.use_alternate_vl_dit:
            self.model = AlternateVLDiT(
                **config.diffusion_model_cfg,
                cross_attention_dim=config.backbone_embedding_dim,
                attend_text_every_n_blocks=config.attend_text_every_n_blocks,
            )
        else:
            self.model = DiT(
                **config.diffusion_model_cfg,
                cross_attention_dim=config.backbone_embedding_dim,
            )

动作头主体是DiT或AlternateVLDiT，但在DiT前后还包了三个映射模块，

复制代码

self.state_encoder = CategorySpecificMLP(
    num_categories=config.max_num_embodiments,
    input_dim=config.max_state_dim * config.state_history_length,
    hidden_dim=self.hidden_size,
    output_dim=self.input_embedding_dim,
)

self.action_encoder = MultiEmbodimentActionEncoder(
    action_dim=self.action_dim,
    hidden_size=self.input_embedding_dim,
    num_embodiments=config.max_num_embodiments,
)

self.action_decoder = CategorySpecificMLP(
    num_categories=config.max_num_embodiments,
    input_dim=self.hidden_size,
    hidden_dim=self.hidden_size,
    output_dim=self.action_dim,
)

state_encoder负责把机器人状态历史压成一个状态Token；action_encoder负责把带噪动作、Flow时间步和机器人类别编码成动作Token；action_decoder负责把DiT输出重新映射回动作维度。这里真正值得注意的是，这几个模块都和embodiment_id相关。换句话说，GR00T N1.7可以共享一个DiT主体，但在靠近机器人状态和动作的输入输出两端，为不同机器人保留类别相关映射参数。

状态进入动作头前的形状大致是，

复制代码

state: [B, state_history_length, max_state_dim]

在前向传播中会被展平成一个Token，

复制代码

action_input.state = action_input.state.view(
    action_input.state.shape[0],
    1,
    -1,
)

因此送给state_encoder的形状变成，

复制代码

[B, 1, state_history_length * max_state_dim]

输出就是一个状态Token，

复制代码

state_features: [B, 1, input_embedding_dim]

动作序列则保持时间维度，每一个未来动作位置都会被编码成一个动作Token。后面状态Token和动作Token会在序列维拼接，组成DiT的主输入序列。

2、CategorySpecificMLP用embodiment_id选择机器人专属映射

多机器人条件编码的底层模块是CategorySpecificLinear，定义在，

复制代码

gr00t/model/modules/embodiment_conditioned_mlp.py

源码如下，

复制代码

class CategorySpecificLinear(nn.Module):
    """Linear layer with category-specific weights and biases for multi-embodiment support."""

    def __init__(self, num_categories, input_dim, hidden_dim):
        super().__init__()
        self.num_categories = num_categories
        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))

    def forward(self, x, cat_ids):
        selected_W = self.W[cat_ids]
        selected_b = self.b[cat_ids]
        return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)

普通nn.Linear只有一组权重，所有样本共用同一个W和b。这里的权重多了一维类别维度，

复制代码

W: [num_categories, input_dim, hidden_dim]
b: [num_categories, hidden_dim]

前向传播时，cat_ids就是embodiment_id，形状为 $B$ 。每个样本会根据自己的机器人类别取出对应权重，

复制代码

selected_W = self.W[cat_ids]
selected_b = self.b[cat_ids]

如果输入x形状是，

复制代码

x: [B, T, input_dim]

那么selected_W形状是，

复制代码

selected_W: [B, input_dim, hidden_dim]

最终通过torch.bmm得到，

复制代码

[B, T, hidden_dim]

这就是GR00T N1.7多机器人动作头的一个关键设计。不同机器人可以共享动作生成主干，但状态输入和动作输出这些靠近物理本体的映射层不完全共享，而是根据embodiment_id选择不同参数。

CategorySpecificMLP是在CategorySpecificLinear基础上封装的两层MLP，

复制代码

class CategorySpecificMLP(nn.Module):
    """Two-layer MLP with category-specific weights for multi-embodiment support."""

    def __init__(self, num_categories, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.num_categories = num_categories
        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)

    def forward(self, x, cat_ids):
        hidden = F.relu(self.layer1(x, cat_ids))
        return self.layer2(hidden, cat_ids)

在动作头里，它主要用在两个位置。第一个是state_encoder，把不同机器人状态映射到统一Token空间；第二个是action_decoder，把DiT输出从统一Token空间映射回不同机器人的动作空间，

复制代码

pred = self.action_decoder(
    model_output,
    embodiment_id,
)

pred_actions = pred[:, -actions.shape[1]:]

因为DiT输入序列最前面拼了一个状态Token，所以输出时只保留最后的动作Token部分。这里的pred_actions才是动作头最终用于监督的动作速度预测。

3、MultiEmbodimentActionEncoder把动作值、Flow时间步和动作位置合成动作Token

动作序列的编码模块是MultiEmbodimentActionEncoder。它比状态编码器更复杂，因为动作Token除了包含动作值，还需要知道当前Flow时间步，并且需要区分Action Chunk中的第几步。

源码如下，

复制代码

class MultiEmbodimentActionEncoder(nn.Module):
    """Action encoder with multi-embodiment support and sinusoidal positional encoding."""

    def __init__(self, action_dim, hidden_size, num_embodiments):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_embodiments = num_embodiments

        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)

这里有三层类别相关线性层，

复制代码

W1: action_dim -> hidden_size
W2: 2 * hidden_size -> hidden_size
W3: hidden_size -> hidden_size

前向传播逻辑如下，

复制代码

def forward(self, actions, timesteps, cat_ids):
    B, T, _ = actions.shape

    if timesteps.dim() == 1 and timesteps.shape[0] == B:
        timesteps = timesteps.unsqueeze(1).expand(-1, T)
    else:
        raise ValueError(
            "Expected `timesteps` to have shape (B,) so we can replicate across T."
        )

    a_emb = self.W1(actions, cat_ids)
    tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)

    x = torch.cat([a_emb, tau_emb], dim=-1)
    x = swish(self.W2(x, cat_ids))

    x = self.W3(x, cat_ids)
    return x

输入timesteps形状是 $B$ ，表示一个样本对应一个Flow时间步。代码会把它扩展到所有动作Token上，

复制代码

timesteps = timesteps.unsqueeze(1).expand(-1, T)

所以同一个动作块里的每个动作Token拿到的是同一个Flow时间条件。动作块整体处在同一个从噪声到真实动作的路径位置上，而不是动作块内每一步使用不同Flow时间。

时间步编码来自SinusoidalPositionalEncoding，

复制代码

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, embedding_dim):
        super().__init__()
        self.embedding_dim = embedding_dim

    def forward(self, timesteps):
        timesteps = timesteps.float()

        B, T = timesteps.shape
        device = timesteps.device

        half_dim = self.embedding_dim // 2
        exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (
            torch.log(torch.tensor(10000.0)) / half_dim
        )
        freqs = timesteps.unsqueeze(-1) * exponent.exp()

        sin = torch.sin(freqs)
        cos = torch.cos(freqs)
        enc = torch.cat([sin, cos], dim=-1)

        return enc

这里的sin/cos编码处理的是Flow时间步，而不是动作块内部位置。动作块内部位置由position_embedding处理，

复制代码

if config.add_pos_embed:
    self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
    nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)

动作编码完成后，如果配置开启位置Embedding，会加到每个动作Token上，

复制代码

if self.config.add_pos_embed:
    pos_ids = torch.arange(
        action_features.shape[1],
        dtype=torch.long,
        device=device,
    )

    pos_embs = self.position_embedding(pos_ids).unsqueeze(0)
    action_features = action_features + pos_embs

这里可以把两个时间概念分开理解，

复制代码

Flow时间步：当前带噪动作处在从噪声到真实动作路径的哪个位置
动作块位置：当前Token是未来动作序列中的第几步

MultiEmbodimentActionEncoder里的sin/cos编码处理前者，position_embedding处理后者。这样每个动作Token同时携带动作值、Flow时间步、动作块位置和机器人类别。

4、视觉语言特征先做轻量整理再作为DiT的条件输入

动作头除了处理状态和动作，也会对Backbone输出的视觉语言特征做一层整理。初始化时有两个模块，

复制代码

self.vlln = (
    nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
)

vl_self_attention_cfg = getattr(config, "vl_self_attention_cfg", None)
if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0:
    self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg)
else:
    self.vl_self_attention = nn.Identity()

前向传播中对应调用，

复制代码

def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
    backbone_features = backbone_output["backbone_features"]
    backbone_features = self.vlln(backbone_features)
    backbone_features = self.vl_self_attention(backbone_features)
    backbone_output["backbone_features"] = backbone_features
    return backbone_output

vlln是视觉语言特征上的LayerNorm，使Backbone输出更适合进入动作头的交叉注意力。vl_self_attention则可以对视觉语言Token再做几层自注意力，相当于在送入DiT之前对VLM序列做一次轻量重整。如果配置里没有开启，这两个模块就退化为nn.Identity()，不改变原始Backbone特征。

SelfAttentionTransformer定义在dit.py里，它也是由BasicTransformerBlock堆叠而成，只不过不接收外部encoder_hidden_states，只对输入视觉语言序列自身做处理，

复制代码

def forward(
    self,
    hidden_states: torch.Tensor,
    return_all_hidden_states: bool = False,
):
    hidden_states = hidden_states.contiguous()
    all_hidden_states = [hidden_states]

    for idx, block in enumerate(self.transformer_blocks):
        hidden_states = block(hidden_states)
        all_hidden_states.append(hidden_states)

    if return_all_hidden_states:
        return hidden_states, all_hidden_states
    else:
        return hidden_states

5、DiT用时间调制和交叉注意力更新动作Token

DiT定义在，

复制代码

gr00t/model/modules/dit.py

该DiT不是图像生成里那种patch化图像Token的DiT，而是面向机器人动作Token序列的Transformer。初始化时会创建时间步编码器，

复制代码

class DiT(ModelMixin, ConfigMixin):
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 8,
        attention_head_dim: int = 64,
        output_dim: int = 26,
        num_layers: int = 12,
        dropout: float = 0.1,
        norm_type: str = "ada_norm",
        interleave_self_attention=False,
        cross_attention_dim: Optional[int] = None,
        ...
    ):
        super().__init__()

        self.attention_head_dim = attention_head_dim
        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim

        self.timestep_encoder = TimestepEncoder(
            embedding_dim=self.inner_dim, compute_dtype=self.compute_dtype
        )

TimestepEncoder会把离散时间步变成一个全局时间条件，

复制代码

class TimestepEncoder(nn.Module):
    def __init__(self, embedding_dim, compute_dtype=torch.float32):
        super().__init__()
        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)

    def forward(self, timesteps):
        dtype = next(self.parameters()).dtype
        timesteps_proj = self.time_proj(timesteps).to(dtype)
        timesteps_emb = self.timestep_embedder(timesteps_proj)
        return timesteps_emb

这里和MultiEmbodimentActionEncoder中的时间编码不是重复错误，而是两个层级的条件注入，

复制代码

action_encoder：把时间步和动作值融合，形成动作Token
DiT内部：把时间步作为全局条件，调制Transformer Block和输出层

DiT默认使用ada_norm时，会通过AdaLayerNorm把时间条件写入归一化层，

复制代码

class AdaLayerNorm(nn.Module):
    def __init__(
        self,
        embedding_dim: int,
        norm_elementwise_affine: bool = False,
        norm_eps: float = 1e-5,
        chunk_dim: int = 0,
    ):
        super().__init__()
        output_dim = embedding_dim * 2
        self.silu = nn.SiLU()
        self.linear = nn.Linear(embedding_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)

    def forward(
        self,
        x: torch.Tensor,
        temb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        temb = self.linear(self.silu(temb))
        scale, shift = temb.chunk(2, dim=1)
        x = self.norm(x) * (1 + scale[:, None]) + shift[:, None]
        return x

temb会被拆成scale和shift，用于调制LayerNorm结果。这样每个Flow时间步都会对应一组不同的归一化缩放和平移，高噪声阶段和低噪声阶段可以走不同的网络行为。

DiT里的BasicTransformerBlock用同一个Attention接口处理自注意力和交叉注意力，

复制代码

self.attn1 = Attention(
    query_dim=dim,
    heads=num_attention_heads,
    dim_head=attention_head_dim,
    dropout=dropout,
    bias=attention_bias,
    cross_attention_dim=cross_attention_dim,
    upcast_attention=upcast_attention,
    out_bias=attention_out_bias,
)

前向传播时是否传入encoder_hidden_states，决定当前做自注意力还是交叉注意力，

复制代码

attn_output = self.attn1(
    norm_hidden_states,
    encoder_hidden_states=encoder_hidden_states,
    attention_mask=(
        encoder_attention_mask if encoder_hidden_states is not None else attention_mask
    ),
)
hidden_states = attn_output + hidden_states

当encoder_hidden_states=None时，就是动作Token之间的自注意力；当传入视觉语言特征时，就是动作Token读取视觉语言条件的交叉注意力。

DiT主循环中还有一个interleave_self_attention配置，

复制代码

for idx, block in enumerate(self.transformer_blocks):
    if idx % 2 == 1 and self.config.interleave_self_attention:
        hidden_states = block(
            hidden_states,
            attention_mask=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            temb=temb,
        )
    else:
        hidden_states = block(
            hidden_states,
            attention_mask=None,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=None,
            temb=temb,
        )

如果interleave_self_attention=False，每一层都可以通过交叉注意力读取视觉语言特征；如果为True，则奇数层只做动作Token内部自注意力，偶数层读取视觉语言条件。这样动作Token可以在内部交互和读取多模态条件之间交替更新。

DiT输出端还会再次使用时间条件做调制，

复制代码

conditioning = temb
shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]

return self.proj_out_2(hidden_states)

也就是说，时间条件不只进入动作Token编码阶段，也进入DiT Block和DiT输出端。这个设计和扩散/流匹配类动作模型的特点是一致的：不同噪声时间位置对应不同生成行为。

6、AlternateVLDiT提供图像Token和非图像Token交替读取路径

如果配置中开启use_alternate_vl_dit，动作头会使用AlternateVLDiT，

复制代码

if config.use_alternate_vl_dit:
    self.model = AlternateVLDiT(
        **config.diffusion_model_cfg,
        cross_attention_dim=config.backbone_embedding_dim,
        attend_text_every_n_blocks=config.attend_text_every_n_blocks,
    )

AlternateVLDiT继承自DiT，主要改的是视觉语言条件的读取方式。普通DiT在cross attention时直接读取完整的encoder_hidden_states，而AlternateVLDiT会根据image_mask把Backbone输出Token拆成图像Token和非图像Token，

复制代码

image_attention_mask = image_mask & backbone_attention_mask
non_image_attention_mask = (~image_mask) & backbone_attention_mask

然后在不同cross attention层中交替选择当前读取哪类Token，

复制代码

for idx, block in enumerate(self.transformer_blocks):
    if idx % 2 == 1:
        hidden_states = block(
            hidden_states,
            attention_mask=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            temb=temb,
        )
    else:
        if idx % (2 * self.attend_text_every_n_blocks) == 0:
            curr_encoder_attention_mask = non_image_attention_mask
        else:
            curr_encoder_attention_mask = image_attention_mask

        hidden_states = block(
            hidden_states,
            attention_mask=None,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=curr_encoder_attention_mask,
            temb=temb,
        )

奇数层做动作Token内部自注意力，偶数层做交叉注意力；偶数层里又根据attend_text_every_n_blocks决定读取非图像Token还是图像Token。这样做相当于给动作生成加入一个更强的结构先验：有些层更关注语言任务描述，有些层更关注视觉场景，中间再穿插动作Token之间的内部交互。

需要注意的是，AlternateVLDiT不是默认一定启用的路径。普通配置会走DiT，只有config.use_alternate_vl_dit=True时才会进入这条分支。因此它更像N1.7动作头里预留的一种视觉语言条件注入方式，而不是理解主流程必须掌握的部分。