本篇文章将向大家介绍支持异构机器人跨模态数据的CrossFormer模型如何实现通用策略模型的模仿学习算法。
在CrossFormer的实现过程中，主要将token区分为task tokens，observations tokens以及readout tokens。依此来实现多模态输入的分词和输出语义的提取。
其中task tokens具体实现时被安排在整个token输入序列的前面（wrapper类为PrefixGroup，the beginning of the token sequence），task的描述规范根据任务不同有不同的定义，如基于抓取动作的任务描述规范一般基于自然语言，而基于导航任务的任务描述规范一般为目标图像（goal image）。基于语言的任务指令描述（language instructions）一般通过FiLM层实现和图像信息的融合，可能是借鉴了RT-1的实现思路，具体的实现可以参考ResNet26FILM类[1]。
observations tokens的wrapper类为TimestepGroup，论文中提到，输入序列中的观测历史长度为k，即为k个时间点的观测数据的集合。
readout的tokens的wrapper类也为TimestepGroup，append到all_timestep_groups的list中，放在observeations的token后面。readout token和BERT的[CLS]特殊token类似，用于提取整个输入的token序列对应的总体语义信息，作为具体的任务head模型组件的输入。但实现时和BERT专门定义的词元[CLS]不同，在Cross Former中以positonal_embedding的形式实现，不同任务readout的tokens个数不同，在生成时将从positonal_embedding中按设定的个数截取即可）。BlockTransformer类的assemble_input_tokens函数实现了将prefix tokens和timestamps tokens进行concatenate形成整体输入。split_output_tokens在transformer架构的backbone的输出张量上做类似于assemble_input_tokens的逆操作逻辑后读取readout对应的张量输出后将其作为输入对接到具体的任务head。
问题1：该开源代码实现中语言指令（language instructions）和图像融合的FiLM代码具体的体现在什么地方？
回答：FiLM的实现可能借鉴了RT-1中的模型中的视觉和语言融合实现的思路，其中具体的实现解释参考引文[3]的问题1及解答。
问题2：不同的观测tokens类型（如图像和本体观测数据）如何实现的融合，是通过一个映射层可学习参数实现不同modality的观测token的融合处理的吗？
回答：get_model_config函数里有定义，针对不同的观测类型有不同的tokenizer，如图像（和自然指令融合）的tokenizer为ImageTokenizer，双臂机器人，四组机器人等的本体感知观测数据的tokenizer为LowdimObsTokenizer。不同类型的tokenizer的词汇表独立，通过线性可学习参数层实现映射到统一的token_embedding_size维度，具体的解释参考引文[3]的问题2及解答。
问题3：readouts的token具体有哪些，如何定义的，和BERT模型的[CLS]类似吗？
回答：在该开源系统的实现中，readouts在实现上以生成位置嵌入的形式存在并concatenate到observations的tokens后面。具体的实现可以参考引文[3]的问题3的回答。
问题4：异构的机器人数据集如何在一起训练，不同horizone窗口的数据对应的tokens输入序列的长度大小不同，是通过padding和mask相结合进行统一处理的吗？以及有什么标记来将transformer的output的embeddings输出给对应的任务head吗？
回答：和一般的LLM模型的批处理方法类似，通过padding和mask实现了不同长度序列的数据的统一训练，可以在generate_attention_mask函数中看一下具体的实现。同时在transformer结构输出时可以根据分组的名称完成split（函数split_output_tokens），和输入时的tokens的组装（assemble_input_tokens）可以看成是互逆的过程，在不同的head进行前向的时候根据不同的readout对应的group的名称的不同，就可以取对应的transformer bacbone输出的output对应的输出（可以类别取BERT的[CLS]位置对应输出的张量）作为对应head的输入进行计算。具体的过程可以参考引文[3]的问题4的回答。
References
- [1]、Scaling Cross-Embodied Learning: One Policy for Manipulation, Navigation, Locomotion and Aviation：2408.11812
- [2]、CrossFormer代码实现：rail-berkeley/crossformer
- [3]、模型的关键实现细节分析：crossformer_model_impl.md
Leave a Reply