pi0 : Our First Generalist Policy (Part. 2)

type

status

date

slug

summary

Part 2: Code

本身是不开源的，但有大佬根据论文复现了一个（实在牛x）

modules.py

从简单的开始（雾）

SinusoidalPosEmb

实现正弦位置编码（ Sinusoidal Positional Embedding ），用于对时间步做 positional encoding （ encoding/embedding 应该一个意思），是从 transformer 那边来的，公式如下

其中，是序列中的位置索引，在这里是时间步；是位置编码的维度（通常是模型的隐藏层维度）；是维度索引（从到）。

但一般不直接用这个公式，而是简单变换一下，先计算一组频率因子

其中，是一般设定为 10000.0 的超参数。再用频率因子与时间步相乘，最后分别应用正弦和余弦函数。

torch.arange() 在这里生成一个从 0 到 self.half_dim - 1 的整数序列。

可以看出， SinusoidalPosEmb 将时间步原始序列顺序信息映射到了模型隐藏层维度。

与 transformer 的原始公式不同，这里是直接把 sin() 和 cos() 直接拼起来了，没有奇偶交错。

ActionEncoder

基本就是对公式的代码实现。

代码中将时间步的嵌入列为了可选项，默认为不嵌入时间步。

self.linear_1 ， self.linear_2 和 self.linear_3 分别对应了，，。

unsqueeze(1) 在 1 位置增加一个维度， expand(-1, action.size(1), -1) 保持第一个维度和第三个维度不变，将第二个维度大小扩展为与 action.size(1) 一样（注：expand 只能扩展大小为 1 的维度）。

GaussianFourierFeatureTransform

这个类实现了一个高斯傅立叶特征变换（Gaussian Fourier Feature Transform），用于将低维输入映射到高维空间，以便神经网络更好地学习高频函数，基于论文 “Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains”

Fourier Features Let Networks Learn High Frequency Functions in...

We show that passing input points through a simple Fourier feature mapping enables a multilayer perceptron (MLP) to learn high-frequency functions in low-dimensional problem domains. These results...

https://arxiv.org/abs/2006.10739

论文中的原公式如下

不过代码中也没管 sin() 和 cos() 交错的问题。

self.b 是一个形状为 [input_dim, embed_dim] 的随机矩阵，元素从标准正态分布中采样，并乘以 scale ，用于将输入数据映射到高维空间。

torch.matmul(...) 计算矩阵乘法。

AdaptiveRMSNorm

这个模块执行自适应 RMS 归一化（Adaptive RMSNorm）操作。

可以把这个模块分成两部分：

RMS 归一化

通过均方根（Root Mean Square, RMS）对输入数据做归一化，从而使数据的分布更加稳定。

稳定训练：通过归一化输入数据，RMS 归一化可以缓解梯度爆炸或梯度消失问题，使训练过程更加稳定。

加速收敛：归一化后的数据分布更加集中，有助于加速模型的收敛。

无需学习参数：RMS 归一化是一种无参数的操作，不需要额外的可学习参数。

条件自适应

条件自适应是一种根据条件信息动态调整模型参数或输出的方法。条件信息可以是额外的输入（如类别标签、时间步长、上下文信息等），模型会根据这些条件信息生成不同的参数或调整输出。~~（deepseek 说的，不知道靠不靠谱）~~

根据条件信息 cond 生成缩放因子 gamma 和偏移因子 beta ，做线性变换

nn.Sequential() 是 PyTorch 中的一个容器，用于按顺序组合多个子模块，输入数据会依次通过每个字模块，最终输出结果。

“条件自适应”部分，self.to_gamma 让条件 cond 通过线性变换从 dim_cond 映射到 dim ，再通过 Sigmoid 约束到。self.to_beta 让条件 cond 通过线性变换映射到 dim ，但因为本身就是一个可学习的偏移因子，所以不需要 bias ，也不需要约束。

“RMS 归一化”部分，直接翻译公式。

rearrange 函数来自einops 库，用于对张量的维度进行重新排列或扩展。

AdaptiveLayerscale

这个模块个人感觉和上一个差不多？

根据输入的条件 cond 对输入特征 x 做线性变换，相当于把条件信息嵌入了特征中。

adaln_zero_gamma_linear 是一个线性层，将条件 cond 从 dim_cond 映射到 dim 。

nn.init.zeros_ 将线性层的权重 Weight 初始化为。

nn.init.constant_ 将线性层的偏置 bias 初始化为 adaln_zero_bias_init_value 。

为什么 bias 初始化为 -2.0

我也不知道qwq

好像是个经验值，相当于 gamma 初始化为 -2.0，经过 sigmoid() 后大概是 0.1192，是一个比较合适的值，有利于训练的稳定。

为什么不像 AdaptiveRMSNorm 一样用 cond 算一个 bias 加到特征 x 上。

还是不知道🤷

mixture.py

可以分为三个部分：

MixtureDecoderLyaer 实现模型的一层，一层中包括自注意力机制和前馈网络 MLP；

MixtureAttention 实现自注意力机制，支持多头注意力、旋转位置编码和 LoRA/QLoRA；

Mixture 组织各个 MixtureDecoderLayer 层和最后的归一化层。

MixtureDecoderLayer

用 config 初始化自注意力模块和 GemmaMLP 。

GemmaMLP 是一种多层感知机（MLP），目的与普通 MLP 一样：对输出做非线性变换。但实现方式应该是基于门控 MLP（gMLP）。这个还挺有意思的，有空可以仔细看看。

GemmaMLP

关于 gMLP 的原论文：

Pay Attention to MLPs

Transformers have become one of the most important architectural innovations in deep learning and have enabled many breakthroughs over the past few years. Here we propose a simple network...

https://arxiv.org/abs/2105.08050

GemmaMLP 代码实现：

一个普通的 MLP 可以用数学公式这样描述：

而 GemmaMLP 可以这样描述：

根据 config 可以选择直接用 GemmaRMSNorm ，或者用 modules.py 里写的 AdaptiveRMSNorm 和 AdaptiveLayerscale 。

getattr(self, norm_name) 动态获取 self 对象中名为 norm_name 的方法。norm_name 是方法名称的字符串，因此可以根据传入的名称来调用不同的归一化方法。

~~感觉不太重要~~

也是模式的设定，感觉也不太重要。

MixtureAttention

这个模块实现一个多头注意力机制。

浏览了一下代码，代码应该不难，但好像对注意力机制还是理解的不到位。决定岔出去彻底理解一下 attention。

🗽pi0 : Our First Generalist Policy (Part. 2)