1、GR00T N1.7由视觉语言骨干和动作模型组成GR00T N1.7是一套面向机器人控制的Vision-Language-Action模型。模型接收相机图像、语言指令和机器人当前状态输出一段连续的机器人动作序列。GR00T N1.7内部包含两条比较明显的处理链路。第一条链路负责处理图像和语言得到视觉语言特征第二条链路负责将机器人状态、带噪动作以及视觉语言特征组合起来生成连续动作。模型主类定义在gr00t/model/gr00t_n1d7/gr00t_n1d7.py其中最主要的两个类是class Gr00tN1d7ActionHead(nn.Module) Action head component for flow matching diffusion policy. class Gr00tN1d7(PreTrainedModel) Gr00tN1d7: VLA model with Cosmos-Reason2-2B (Qwen3-VL) backbone.Gr00tN1d7负责组织完整模型内部包含视觉语言Backbone和动作头Gr00tN1d7ActionHead负责根据视觉语言特征、机器人状态和带噪动作预测动作速度。完整的数据流如下图像、语言 - Cosmos-Reason2-2B基于Qwen3-VL - 视觉语言特征视觉语言特征 机器人状态 带噪动作 - DiT - 动作速度场 - Euler积分 - 连续动作序列这里需要先区分两个经常混在一起的概念。GR00T N1.7的动作网络使用DiT也就是Diffusion Transformer但动作的训练目标采用Flow Matching。DiT描述的是网络结构Flow Matching描述的是训练和采样方法两者并不冲突。2、N1.7微调入口完成模型与数据配置N1.7微调入口位于gr00t/experiment/launch_finetune.py命令行参数由tyro根据FinetuneConfig自动生成ft_config tyro.cli(FinetuneConfig, description__doc__)FinetuneConfig定义在gr00t/configs/finetune_config.py该配置类中包含模型路径、数据集路径、机器人类型、训练哪些模块、学习率、Batch Size、保存频率等参数。例如dataclass class FinetuneConfig: base_model_path: str dataset_path: str embodiment_tag: str modality_config_path: str | None None tune_llm: bool False tune_visual: bool False tune_projector: bool True tune_diffusion_model: bool True global_batch_size: int 64 learning_rate: float 1e-4 gradient_accumulation_steps: int 1 output_dir: str ./outputs从默认值可以看出常规微调不会更新完整的语言模型和视觉编码器而是训练动作头中的状态编码器、动作编码器、动作解码器以及DiT主体。需要注意的是配置项虽然叫tune_projector但从后面的源码可以看到它控制的并不是单独一个传统意义上的多模态投影层而是动作头前后的多组特征映射模块。tune_llm False tune_visual False tune_projector True tune_diffusion_model True这种设置比较符合机器人数据集的实际情况。机器人示范数据的规模通常远小于视觉语言模型的预训练数据如果直接更新完整VLM不仅显存开销很高也容易破坏原有的视觉语言能力。进入主函数后代码先解析embodiment_tagfrom gr00t.data.embodiment_tags import EmbodimentTag ft_config.embodiment_tag EmbodimentTag.resolve(ft_config.embodiment_tag) embodiment_tag ft_config.embodiment_tag.valueembodiment_tag可以理解为机器人本体类型标识。不同机器人可能有完全不同的状态维度和动作维度例如单臂机械臂、双臂机器人、人形机器人以及仿真环境中的机械臂其关节数量和动作定义都不同。GR00T通过该标识选择对应的数据配置以及类别相关的动作编码器。数据集路径支持传入多个目录dataset_paths [ path for path in ft_config.dataset_path.split(os.pathsep) if path ]在Linux系统中os.pathsep通常是冒号因此可以通过下面的方式传入多个数据集--dataset-path /data/task_a:/data/task_b随后代码加载默认配置并将用户传入的数据集注册到配置中config get_default_config().load_dict( { data: { download_cache: False, datasets: [ { dataset_paths: dataset_paths, mix_ratio: 1.0, embodiment_tag: embodiment_tag, } ], } } )接下来会覆盖N1.7使用的模型参数config.model.tune_llm ft_config.tune_llm config.model.tune_visual ft_config.tune_visual config.model.tune_projector ft_config.tune_projector config.model.tune_diffusion_model ft_config.tune_diffusion_model config.model.load_bf16 False config.model.reproject_vision False config.model.model_name nvidia/Cosmos-Reason2-2B config.model.backbone_trainable_params_fp32 True config.model.use_relative_action True这里有几个配置值得单独说明。model_name被固定为nvidia/Cosmos-Reason2-2B说明N1.7默认使用Cosmos-Reason2-2B作为视觉语言骨干use_relative_actionTrue表示训练时默认使用相对动作backbone_trainable_params_fp32True表示VLM中可训练的参数会保留FP32精度避免低精度训练导致数值不稳定。最后将训练参数写入总配置并调用run(config)因此launch_finetune.py本身并不包含训练循环它的主要工作是读取命令行参数、构造配置然后将真正的训练交给gr00t/experiment/experiment.py。3、Gr00tN1d7主类组织完整的VLA模型Gr00tN1d7继承自Hugging Face的PreTrainedModelclass Gr00tN1d7(PreTrainedModel): config_class Gr00tN1d7Config supports_gradient_checkpointing True继承PreTrainedModel以后模型可以使用Hugging Face提供的权重保存和加载接口例如AutoModel.from_pretrained(...) model.save_pretrained(...)初始化函数中首先构造Backbonebackbone_cls get_backbone_cls(config) self.backbone backbone_cls( model_nameconfig.model_name, tune_llmconfig.tune_llm, tune_visualconfig.tune_visual, select_layerconfig.select_layer, reproject_visionconfig.reproject_vision, use_flash_attentionconfig.use_flash_attention, load_bf16config.load_bf16, tune_top_llm_layersconfig.tune_top_llm_layers, trainable_params_fp32config.backbone_trainable_params_fp32, transformers_loading_kwargstransformers_loading_kwargs, )get_backbone_cls(config)会根据配置返回对应的视觉语言模型封装类。N1.7配置中使用的是Cosmos-Reason2-2B其底层结构属于Qwen3-VL系列因此后续动作模型拿到的不是原始图片而是Qwen3-VL编码后的序列特征。动作头的初始化比较直接self.action_head Gr00tN1d7ActionHead(config)此外还会创建数据Collatorfrom .processing_gr00t_n1d7 import Gr00tN1d7DataCollator self.collator Gr00tN1d7DataCollator( model_nameconfig.model_name, model_typeconfig.backbone_model_type, transformers_loading_kwargstransformers_loading_kwargs, )Collator负责把一个Batch中的图像和语言整理成VLM能够接受的输入格式包括input_ids、attention_mask、图像Tensor以及视觉相关的网格信息等。模型的forward函数没有额外的复杂逻辑def forward(self, inputs: dict) - BatchFeature: backbone_inputs, action_inputs self.prepare_input(inputs) backbone_outputs self.backbone(backbone_inputs) action_outputs self.action_head(backbone_outputs, action_inputs) return action_outputs这段代码将整个训练过程分成三步prepare_input - backbone - action_headprepare_input会将同一个输入字典拆成两部分。图像、文本等数据交给Backbone状态、动作、动作Mask和embodiment_id交给动作头backbone_inputs self.backbone.prepare_input(inputs) action_inputs self.action_head.prepare_input(inputs)随后使用统一函数把Tensor移动到模型所在设备def to_device_with_dtype(x): if torch.is_floating_point(x): return x.to(self.device, dtypeself.dtype) else: return x.to(self.device)浮点Tensor不仅会移动到GPU还会转换为模型当前使用的数据类型整数类型的Token ID和Mask只移动设备不进行浮点转换。4、动作头由状态编码器、动作编码器和DiT组成动作头类名为Gr00tN1d7ActionHeadclass Gr00tN1d7ActionHead(nn.Module): Action head component for flow matching diffusion policy.注释中已经直接写出了flow matching diffusion policy。这个表述容易让人误以为代码同时实现了两套方法实际上这里只实现一套动作生成流程使用DiT预测Flow Matching速度场。 动作头首先根据配置创建DiTif config.use_alternate_vl_dit: self.model AlternateVLDiT( **config.diffusion_model_cfg, cross_attention_dimconfig.backbone_embedding_dim, attend_text_every_n_blocksconfig.attend_text_every_n_blocks, ) else: self.model DiT( **config.diffusion_model_cfg, cross_attention_dimconfig.backbone_embedding_dim, )正常配置下使用DiT开启use_alternate_vl_dit后使用AlternateVLDiT。两种结构都会将状态和动作特征作为主序列并通过交叉注意力读取视觉语言特征。下面几个成员负责状态和动作的维度映射self.state_encoder CategorySpecificMLP( num_categoriesconfig.max_num_embodiments, input_dimconfig.max_state_dim * config.state_history_length, hidden_dimself.hidden_size, output_dimself.input_embedding_dim, ) self.action_encoder MultiEmbodimentActionEncoder( action_dimself.action_dim, hidden_sizeself.input_embedding_dim, num_embodimentsconfig.max_num_embodiments, ) self.action_decoder CategorySpecificMLP( num_categoriesconfig.max_num_embodiments, input_dimself.hidden_size, hidden_dimself.hidden_size, output_dimself.action_dim, )state_encoder将机器人状态映射到DiT的输入维度action_encoder将连续动作和时间步编码为动作Tokenaction_decoder将DiT输出还原为连续动作速度。这里使用了CategorySpecificMLP和MultiEmbodimentActionEncoder原因是不同机器人本体的数据分布不同。即使两台机器人都使用7维动作其中每一维的物理含义也可能不同。通过embodiment_id选择类别相关参数可以在共享DiT主体的同时为不同机器人保留各自的输入输出映射。动作头还会记录以下参数self.action_dim config.max_action_dim self.action_horizon config.action_horizon self.num_inference_timesteps config.num_inference_timestepsaction_dim是补齐后的最大动作维度action_horizon表示模型一次预测多少步动作num_inference_timesteps表示推理时执行多少次速度积分。5、动作头同时控制参数梯度和模块运行模式动作头通过set_trainable_parameters控制不同模块是否参与训练def set_trainable_parameters( self, tune_projector: bool, tune_diffusion_model: bool, tune_vlln: bool, ): for p in self.parameters(): p.requires_grad True if not tune_projector: self.state_encoder.requires_grad_(False) self.action_encoder.requires_grad_(False) self.action_decoder.requires_grad_(False) if self.config.add_pos_embed: self.position_embedding.requires_grad_(False) if not tune_diffusion_model: self.model.requires_grad_(False) if not tune_vlln: self.vlln.requires_grad_(False) self.vl_self_attention.requires_grad_(False)tune_projector控制状态编码器、动作编码器、动作解码器以及位置编码tune_diffusion_model控制DiT主体tune_vlln控制视觉语言特征进入动作头之前的LayerNorm和自注意力层。代码中还实现了set_frozen_modules_to_eval_modedef set_frozen_modules_to_eval_mode(self): if self.training: if not self.tune_projector: self.state_encoder.eval() self.action_encoder.eval() self.action_decoder.eval() if not self.tune_diffusion_model: self.model.eval() if not self.tune_vlln: self.vlln.eval() self.vl_self_attention.eval()仅设置requires_gradFalse并不会自动将模块切换到评估模式。Hugging Face Trainer在训练期间会调用model.train()被冻结模块中的Dropout等训练态行为仍可能生效因此源码又在forward开头调用set_frozen_modules_to_eval_mode()将这些模块重新切换到eval()状态。6、Flow Matching在噪声和真实动作之间构造训练轨迹动作头forward函数的输入包括视觉语言特征和动作相关数据def forward( self, backbone_output: BatchFeature, action_input: BatchFeature, ) - BatchFeature:Backbone输出中主要包含backbone_features: [B, seq_len, backbone_embedding_dim] backbone_attention_mask:[B, seq_len]动作输入中主要包含state: [B, state_history, state_dim] action: [B, action_horizon, action_dim] embodiment_id: [B] action_mask: [B, action_horizon, action_dim]首先对Backbone特征做额外处理backbone_output self.process_backbone_output(backbone_output) vl_embeds backbone_output.backbone_featuresprocess_backbone_output中包含LayerNorm和可选的视觉语言自注意力层def process_backbone_output(self, backbone_output): 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接着处理机器人状态。状态历史会先展平assert action_input.state.shape[1] self.config.state_history_length action_input.state action_input.state.view( action_input.state.shape[0], 1, -1, )假设状态历史长度为4每一帧状态维度为32展平后得到128维向量再由state_encoder映射为一个状态Tokenstate_features self.state_encoder( action_input.state, embodiment_id, )训练阶段还会按概率丢弃完整状态特征if self.training and self.state_dropout_prob 0: do_dropout ( torch.rand( state_features.shape[0], devicestate_features.device, ) self.state_dropout_prob ) do_dropout do_dropout[:, None, None].to( dtypestate_features.dtype ) state_features state_features * (1 - do_dropout)这里不是对状态向量中的单个元素做Dropout而是对一个样本的整个状态Token置零。这样训练出来的模型不会完全依赖机器人状态在状态传感器存在噪声或缺失时仍能利用图像和语言生成动作。Flow Matching目标的构造集中在下面几行actions action_input.action noise torch.randn( actions.shape, deviceactions.device, dtypeactions.dtype, ) t self.sample_time( actions.shape[0], deviceactions.device, dtypeactions.dtype, ) t t[:, None, None] noisy_trajectory (1 - t) * noise t * actions velocity actions - noise模型输入noisy_trajectory、时间步以及条件信息输出该位置上的速度。训练目标不是预测高斯噪声也不是直接恢复干净动作而是预测从噪声指向真实动作的速度方向。时间t并不是均匀采样而是由Beta分布产生self.beta_dist Beta( config.noise_beta_alpha, config.noise_beta_beta, ) def sample_time(self, batch_size, device, dtype): sample self.beta_dist.sample([batch_size]).to( device, dtypedtype, ) sample (1 - sample) * self.config.noise_s return sampleBeta分布可以控制训练样本更多地落在轨迹的哪个区域。不同于简单的均匀采样这种方式可以调整模型对高噪声阶段和低噪声阶段的学习比例。连续时间t随后被离散化t_discretized ( t[:, 0, 0] * self.num_timestep_buckets ).long()离散时间步会传给动作编码器和DiTaction_features self.action_encoder( noisy_trajectory, t_discretized, embodiment_id, )如果启用了位置编码还会给动作序列中的每一个时间位置添加独立Embeddingif self.config.add_pos_embed: pos_ids torch.arange( action_features.shape[1], dtypetorch.long, devicedevice, ) pos_embs self.position_embedding(pos_ids).unsqueeze(0) action_features action_features pos_embs这样DiT可以区分动作块中的第0步、第1步和后续时间步。7、DiT根据多模态条件预测动作速度并计算损失状态特征和动作特征会在序列维度拼接sa_embs torch.cat( (state_features, action_features), dim1, )如果状态被编码为1个Token动作块长度为16那么拼接后的序列长度就是17[state_token, action_0, action_1, ..., action_15]视觉语言特征不会直接拼接进这个序列而是作为交叉注意力的encoder_hidden_states输入DiTmodel_output, _ self.model( hidden_statessa_embs, encoder_hidden_statesvl_embeds, encoder_attention_maskvl_attn_mask, timestept_discretized, return_all_hidden_statesTrue, )这种结构可以看成条件生成模型。主序列是机器人状态和动作视觉语言特征提供环境和任务条件。每一层DiT都可以根据当前图像和语言指令更新动作Token。DiT输出经过动作解码器pred self.action_decoder( model_output, embodiment_id, )由于输出中还包含状态Token对应的位置因此只保留最后的动作部分pred_actions pred[:, -actions.shape[1]:]动作损失采用逐元素MSEaction_loss F.mse_loss( pred_actions, velocity, reductionnone, )随后乘以action_mask。动作Mask的作用是忽略补零维度以及无效时间步。GR00T需要支持不同机器人一些机器人可能只有7维动作另一些机器人可能有20维动作模型内部会统一补齐到max_action_dim但补齐部分不应参与损失计算。源码随后将逐元素MSE乘以action_mask再根据有效动作元素的数量计算平均损失action_loss action_loss * action_mask loss action_loss.sum() / ( action_mask.sum() 1e-6 )分母使用有效元素数量而不是直接对整个Tensor求平均可以避免不同动作维度和不同Padding长度对损失尺度造成影响。8、推理阶段从高斯噪声逐步生成动作序列训练阶段学习速度场推理阶段则沿着该速度场进行积分。动作生成入口为Gr00tN1d7.get_action(...)主模型先处理输入并计算视觉语言特征backbone_inputs, action_inputs self.prepare_input(inputs) backbone_outputs self.backbone(backbone_inputs) action_outputs self.action_head.get_action( backbone_outputs, action_inputs, options, )动作头首先创建一段高斯噪声actions torch.randn( size( batch_size, self.config.action_horizon, self.action_dim, ), dtypevl_embeds.dtype, devicedevice, )然后根据推理步数计算步长dt 1.0 / self.num_inference_timesteps如果num_inference_timesteps4每次积分步长就是0.25。每轮迭代都会对当前动作轨迹重新编码for t in range(self.num_inference_timesteps): t_cont t / float(self.num_inference_timesteps) t_discretized int( t_cont * self.num_timestep_buckets ) timesteps_tensor torch.full( size(batch_size,), fill_valuet_discretized, devicedevice, ) action_features self.action_encoder( actions, timesteps_tensor, embodiment_id, )随后将状态Token和动作Token拼接交给DiTsa_embs torch.cat( (state_features, action_features), dim1, ) model_output self.model( hidden_statessa_embs, encoder_hidden_statesvl_embeds, timesteptimesteps_tensor, )动作解码器输出当前轨迹上的速度pred self.action_decoder( model_output, embodiment_id, ) pred_velocity pred[:, -self.action_horizon:]最后使用Euler方法更新动作actions actions dt * pred_velocity * vel_strength
GR00T N1.7源码学习（一）：工程入口、模型结构与动作生成流程解析
1、GR00T N1.7由视觉语言骨干和动作模型组成GR00T N1.7是一套面向机器人控制的Vision-Language-Action模型。模型接收相机图像、语言指令和机器人当前状态输出一段连续的机器人动作序列。GR00T N1.7内部包含两条比较明显的处理链路。第一条链路负责处理图像和语言得到视觉语言特征第二条链路负责将机器人状态、带噪动作以及视觉语言特征组合起来生成连续动作。模型主类定义在gr00t/model/gr00t_n1d7/gr00t_n1d7.py其中最主要的两个类是class Gr00tN1d7ActionHead(nn.Module) Action head component for flow matching diffusion policy. class Gr00tN1d7(PreTrainedModel) Gr00tN1d7: VLA model with Cosmos-Reason2-2B (Qwen3-VL) backbone.Gr00tN1d7负责组织完整模型内部包含视觉语言Backbone和动作头Gr00tN1d7ActionHead负责根据视觉语言特征、机器人状态和带噪动作预测动作速度。完整的数据流如下图像、语言 - Cosmos-Reason2-2B基于Qwen3-VL - 视觉语言特征视觉语言特征 机器人状态 带噪动作 - DiT - 动作速度场 - Euler积分 - 连续动作序列这里需要先区分两个经常混在一起的概念。GR00T N1.7的动作网络使用DiT也就是Diffusion Transformer但动作的训练目标采用Flow Matching。DiT描述的是网络结构Flow Matching描述的是训练和采样方法两者并不冲突。2、N1.7微调入口完成模型与数据配置N1.7微调入口位于gr00t/experiment/launch_finetune.py命令行参数由tyro根据FinetuneConfig自动生成ft_config tyro.cli(FinetuneConfig, description__doc__)FinetuneConfig定义在gr00t/configs/finetune_config.py该配置类中包含模型路径、数据集路径、机器人类型、训练哪些模块、学习率、Batch Size、保存频率等参数。例如dataclass class FinetuneConfig: base_model_path: str dataset_path: str embodiment_tag: str modality_config_path: str | None None tune_llm: bool False tune_visual: bool False tune_projector: bool True tune_diffusion_model: bool True global_batch_size: int 64 learning_rate: float 1e-4 gradient_accumulation_steps: int 1 output_dir: str ./outputs从默认值可以看出常规微调不会更新完整的语言模型和视觉编码器而是训练动作头中的状态编码器、动作编码器、动作解码器以及DiT主体。需要注意的是配置项虽然叫tune_projector但从后面的源码可以看到它控制的并不是单独一个传统意义上的多模态投影层而是动作头前后的多组特征映射模块。tune_llm False tune_visual False tune_projector True tune_diffusion_model True这种设置比较符合机器人数据集的实际情况。机器人示范数据的规模通常远小于视觉语言模型的预训练数据如果直接更新完整VLM不仅显存开销很高也容易破坏原有的视觉语言能力。进入主函数后代码先解析embodiment_tagfrom gr00t.data.embodiment_tags import EmbodimentTag ft_config.embodiment_tag EmbodimentTag.resolve(ft_config.embodiment_tag) embodiment_tag ft_config.embodiment_tag.valueembodiment_tag可以理解为机器人本体类型标识。不同机器人可能有完全不同的状态维度和动作维度例如单臂机械臂、双臂机器人、人形机器人以及仿真环境中的机械臂其关节数量和动作定义都不同。GR00T通过该标识选择对应的数据配置以及类别相关的动作编码器。数据集路径支持传入多个目录dataset_paths [ path for path in ft_config.dataset_path.split(os.pathsep) if path ]在Linux系统中os.pathsep通常是冒号因此可以通过下面的方式传入多个数据集--dataset-path /data/task_a:/data/task_b随后代码加载默认配置并将用户传入的数据集注册到配置中config get_default_config().load_dict( { data: { download_cache: False, datasets: [ { dataset_paths: dataset_paths, mix_ratio: 1.0, embodiment_tag: embodiment_tag, } ], } } )接下来会覆盖N1.7使用的模型参数config.model.tune_llm ft_config.tune_llm config.model.tune_visual ft_config.tune_visual config.model.tune_projector ft_config.tune_projector config.model.tune_diffusion_model ft_config.tune_diffusion_model config.model.load_bf16 False config.model.reproject_vision False config.model.model_name nvidia/Cosmos-Reason2-2B config.model.backbone_trainable_params_fp32 True config.model.use_relative_action True这里有几个配置值得单独说明。model_name被固定为nvidia/Cosmos-Reason2-2B说明N1.7默认使用Cosmos-Reason2-2B作为视觉语言骨干use_relative_actionTrue表示训练时默认使用相对动作backbone_trainable_params_fp32True表示VLM中可训练的参数会保留FP32精度避免低精度训练导致数值不稳定。最后将训练参数写入总配置并调用run(config)因此launch_finetune.py本身并不包含训练循环它的主要工作是读取命令行参数、构造配置然后将真正的训练交给gr00t/experiment/experiment.py。3、Gr00tN1d7主类组织完整的VLA模型Gr00tN1d7继承自Hugging Face的PreTrainedModelclass Gr00tN1d7(PreTrainedModel): config_class Gr00tN1d7Config supports_gradient_checkpointing True继承PreTrainedModel以后模型可以使用Hugging Face提供的权重保存和加载接口例如AutoModel.from_pretrained(...) model.save_pretrained(...)初始化函数中首先构造Backbonebackbone_cls get_backbone_cls(config) self.backbone backbone_cls( model_nameconfig.model_name, tune_llmconfig.tune_llm, tune_visualconfig.tune_visual, select_layerconfig.select_layer, reproject_visionconfig.reproject_vision, use_flash_attentionconfig.use_flash_attention, load_bf16config.load_bf16, tune_top_llm_layersconfig.tune_top_llm_layers, trainable_params_fp32config.backbone_trainable_params_fp32, transformers_loading_kwargstransformers_loading_kwargs, )get_backbone_cls(config)会根据配置返回对应的视觉语言模型封装类。N1.7配置中使用的是Cosmos-Reason2-2B其底层结构属于Qwen3-VL系列因此后续动作模型拿到的不是原始图片而是Qwen3-VL编码后的序列特征。动作头的初始化比较直接self.action_head Gr00tN1d7ActionHead(config)此外还会创建数据Collatorfrom .processing_gr00t_n1d7 import Gr00tN1d7DataCollator self.collator Gr00tN1d7DataCollator( model_nameconfig.model_name, model_typeconfig.backbone_model_type, transformers_loading_kwargstransformers_loading_kwargs, )Collator负责把一个Batch中的图像和语言整理成VLM能够接受的输入格式包括input_ids、attention_mask、图像Tensor以及视觉相关的网格信息等。模型的forward函数没有额外的复杂逻辑def forward(self, inputs: dict) - BatchFeature: backbone_inputs, action_inputs self.prepare_input(inputs) backbone_outputs self.backbone(backbone_inputs) action_outputs self.action_head(backbone_outputs, action_inputs) return action_outputs这段代码将整个训练过程分成三步prepare_input - backbone - action_headprepare_input会将同一个输入字典拆成两部分。图像、文本等数据交给Backbone状态、动作、动作Mask和embodiment_id交给动作头backbone_inputs self.backbone.prepare_input(inputs) action_inputs self.action_head.prepare_input(inputs)随后使用统一函数把Tensor移动到模型所在设备def to_device_with_dtype(x): if torch.is_floating_point(x): return x.to(self.device, dtypeself.dtype) else: return x.to(self.device)浮点Tensor不仅会移动到GPU还会转换为模型当前使用的数据类型整数类型的Token ID和Mask只移动设备不进行浮点转换。4、动作头由状态编码器、动作编码器和DiT组成动作头类名为Gr00tN1d7ActionHeadclass Gr00tN1d7ActionHead(nn.Module): Action head component for flow matching diffusion policy.注释中已经直接写出了flow matching diffusion policy。这个表述容易让人误以为代码同时实现了两套方法实际上这里只实现一套动作生成流程使用DiT预测Flow Matching速度场。 动作头首先根据配置创建DiTif config.use_alternate_vl_dit: self.model AlternateVLDiT( **config.diffusion_model_cfg, cross_attention_dimconfig.backbone_embedding_dim, attend_text_every_n_blocksconfig.attend_text_every_n_blocks, ) else: self.model DiT( **config.diffusion_model_cfg, cross_attention_dimconfig.backbone_embedding_dim, )正常配置下使用DiT开启use_alternate_vl_dit后使用AlternateVLDiT。两种结构都会将状态和动作特征作为主序列并通过交叉注意力读取视觉语言特征。下面几个成员负责状态和动作的维度映射self.state_encoder CategorySpecificMLP( num_categoriesconfig.max_num_embodiments, input_dimconfig.max_state_dim * config.state_history_length, hidden_dimself.hidden_size, output_dimself.input_embedding_dim, ) self.action_encoder MultiEmbodimentActionEncoder( action_dimself.action_dim, hidden_sizeself.input_embedding_dim, num_embodimentsconfig.max_num_embodiments, ) self.action_decoder CategorySpecificMLP( num_categoriesconfig.max_num_embodiments, input_dimself.hidden_size, hidden_dimself.hidden_size, output_dimself.action_dim, )state_encoder将机器人状态映射到DiT的输入维度action_encoder将连续动作和时间步编码为动作Tokenaction_decoder将DiT输出还原为连续动作速度。这里使用了CategorySpecificMLP和MultiEmbodimentActionEncoder原因是不同机器人本体的数据分布不同。即使两台机器人都使用7维动作其中每一维的物理含义也可能不同。通过embodiment_id选择类别相关参数可以在共享DiT主体的同时为不同机器人保留各自的输入输出映射。动作头还会记录以下参数self.action_dim config.max_action_dim self.action_horizon config.action_horizon self.num_inference_timesteps config.num_inference_timestepsaction_dim是补齐后的最大动作维度action_horizon表示模型一次预测多少步动作num_inference_timesteps表示推理时执行多少次速度积分。5、动作头同时控制参数梯度和模块运行模式动作头通过set_trainable_parameters控制不同模块是否参与训练def set_trainable_parameters( self, tune_projector: bool, tune_diffusion_model: bool, tune_vlln: bool, ): for p in self.parameters(): p.requires_grad True if not tune_projector: self.state_encoder.requires_grad_(False) self.action_encoder.requires_grad_(False) self.action_decoder.requires_grad_(False) if self.config.add_pos_embed: self.position_embedding.requires_grad_(False) if not tune_diffusion_model: self.model.requires_grad_(False) if not tune_vlln: self.vlln.requires_grad_(False) self.vl_self_attention.requires_grad_(False)tune_projector控制状态编码器、动作编码器、动作解码器以及位置编码tune_diffusion_model控制DiT主体tune_vlln控制视觉语言特征进入动作头之前的LayerNorm和自注意力层。代码中还实现了set_frozen_modules_to_eval_modedef set_frozen_modules_to_eval_mode(self): if self.training: if not self.tune_projector: self.state_encoder.eval() self.action_encoder.eval() self.action_decoder.eval() if not self.tune_diffusion_model: self.model.eval() if not self.tune_vlln: self.vlln.eval() self.vl_self_attention.eval()仅设置requires_gradFalse并不会自动将模块切换到评估模式。Hugging Face Trainer在训练期间会调用model.train()被冻结模块中的Dropout等训练态行为仍可能生效因此源码又在forward开头调用set_frozen_modules_to_eval_mode()将这些模块重新切换到eval()状态。6、Flow Matching在噪声和真实动作之间构造训练轨迹动作头forward函数的输入包括视觉语言特征和动作相关数据def forward( self, backbone_output: BatchFeature, action_input: BatchFeature, ) - BatchFeature:Backbone输出中主要包含backbone_features: [B, seq_len, backbone_embedding_dim] backbone_attention_mask:[B, seq_len]动作输入中主要包含state: [B, state_history, state_dim] action: [B, action_horizon, action_dim] embodiment_id: [B] action_mask: [B, action_horizon, action_dim]首先对Backbone特征做额外处理backbone_output self.process_backbone_output(backbone_output) vl_embeds backbone_output.backbone_featuresprocess_backbone_output中包含LayerNorm和可选的视觉语言自注意力层def process_backbone_output(self, backbone_output): 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接着处理机器人状态。状态历史会先展平assert action_input.state.shape[1] self.config.state_history_length action_input.state action_input.state.view( action_input.state.shape[0], 1, -1, )假设状态历史长度为4每一帧状态维度为32展平后得到128维向量再由state_encoder映射为一个状态Tokenstate_features self.state_encoder( action_input.state, embodiment_id, )训练阶段还会按概率丢弃完整状态特征if self.training and self.state_dropout_prob 0: do_dropout ( torch.rand( state_features.shape[0], devicestate_features.device, ) self.state_dropout_prob ) do_dropout do_dropout[:, None, None].to( dtypestate_features.dtype ) state_features state_features * (1 - do_dropout)这里不是对状态向量中的单个元素做Dropout而是对一个样本的整个状态Token置零。这样训练出来的模型不会完全依赖机器人状态在状态传感器存在噪声或缺失时仍能利用图像和语言生成动作。Flow Matching目标的构造集中在下面几行actions action_input.action noise torch.randn( actions.shape, deviceactions.device, dtypeactions.dtype, ) t self.sample_time( actions.shape[0], deviceactions.device, dtypeactions.dtype, ) t t[:, None, None] noisy_trajectory (1 - t) * noise t * actions velocity actions - noise模型输入noisy_trajectory、时间步以及条件信息输出该位置上的速度。训练目标不是预测高斯噪声也不是直接恢复干净动作而是预测从噪声指向真实动作的速度方向。时间t并不是均匀采样而是由Beta分布产生self.beta_dist Beta( config.noise_beta_alpha, config.noise_beta_beta, ) def sample_time(self, batch_size, device, dtype): sample self.beta_dist.sample([batch_size]).to( device, dtypedtype, ) sample (1 - sample) * self.config.noise_s return sampleBeta分布可以控制训练样本更多地落在轨迹的哪个区域。不同于简单的均匀采样这种方式可以调整模型对高噪声阶段和低噪声阶段的学习比例。连续时间t随后被离散化t_discretized ( t[:, 0, 0] * self.num_timestep_buckets ).long()离散时间步会传给动作编码器和DiTaction_features self.action_encoder( noisy_trajectory, t_discretized, embodiment_id, )如果启用了位置编码还会给动作序列中的每一个时间位置添加独立Embeddingif self.config.add_pos_embed: pos_ids torch.arange( action_features.shape[1], dtypetorch.long, devicedevice, ) pos_embs self.position_embedding(pos_ids).unsqueeze(0) action_features action_features pos_embs这样DiT可以区分动作块中的第0步、第1步和后续时间步。7、DiT根据多模态条件预测动作速度并计算损失状态特征和动作特征会在序列维度拼接sa_embs torch.cat( (state_features, action_features), dim1, )如果状态被编码为1个Token动作块长度为16那么拼接后的序列长度就是17[state_token, action_0, action_1, ..., action_15]视觉语言特征不会直接拼接进这个序列而是作为交叉注意力的encoder_hidden_states输入DiTmodel_output, _ self.model( hidden_statessa_embs, encoder_hidden_statesvl_embeds, encoder_attention_maskvl_attn_mask, timestept_discretized, return_all_hidden_statesTrue, )这种结构可以看成条件生成模型。主序列是机器人状态和动作视觉语言特征提供环境和任务条件。每一层DiT都可以根据当前图像和语言指令更新动作Token。DiT输出经过动作解码器pred self.action_decoder( model_output, embodiment_id, )由于输出中还包含状态Token对应的位置因此只保留最后的动作部分pred_actions pred[:, -actions.shape[1]:]动作损失采用逐元素MSEaction_loss F.mse_loss( pred_actions, velocity, reductionnone, )随后乘以action_mask。动作Mask的作用是忽略补零维度以及无效时间步。GR00T需要支持不同机器人一些机器人可能只有7维动作另一些机器人可能有20维动作模型内部会统一补齐到max_action_dim但补齐部分不应参与损失计算。源码随后将逐元素MSE乘以action_mask再根据有效动作元素的数量计算平均损失action_loss action_loss * action_mask loss action_loss.sum() / ( action_mask.sum() 1e-6 )分母使用有效元素数量而不是直接对整个Tensor求平均可以避免不同动作维度和不同Padding长度对损失尺度造成影响。8、推理阶段从高斯噪声逐步生成动作序列训练阶段学习速度场推理阶段则沿着该速度场进行积分。动作生成入口为Gr00tN1d7.get_action(...)主模型先处理输入并计算视觉语言特征backbone_inputs, action_inputs self.prepare_input(inputs) backbone_outputs self.backbone(backbone_inputs) action_outputs self.action_head.get_action( backbone_outputs, action_inputs, options, )动作头首先创建一段高斯噪声actions torch.randn( size( batch_size, self.config.action_horizon, self.action_dim, ), dtypevl_embeds.dtype, devicedevice, )然后根据推理步数计算步长dt 1.0 / self.num_inference_timesteps如果num_inference_timesteps4每次积分步长就是0.25。每轮迭代都会对当前动作轨迹重新编码for t in range(self.num_inference_timesteps): t_cont t / float(self.num_inference_timesteps) t_discretized int( t_cont * self.num_timestep_buckets ) timesteps_tensor torch.full( size(batch_size,), fill_valuet_discretized, devicedevice, ) action_features self.action_encoder( actions, timesteps_tensor, embodiment_id, )随后将状态Token和动作Token拼接交给DiTsa_embs torch.cat( (state_features, action_features), dim1, ) model_output self.model( hidden_statessa_embs, encoder_hidden_statesvl_embeds, timesteptimesteps_tensor, )动作解码器输出当前轨迹上的速度pred self.action_decoder( model_output, embodiment_id, ) pred_velocity pred[:, -self.action_horizon:]最后使用Euler方法更新动作actions actions dt * pred_velocity * vel_strength
