ACT----斯坦福具身智能入门梳理

0.简介

ALOHA算得上是今年年初很惊艳的具身项目了，其项目网站: https://mobile-aloha.github.io/ 。而ACT模仿学习算法能够有效处理复合误差，通过动作分块和时间集成减少任务的有效范围，缓解模仿学习中的复合误差问题，提高了在精细操作任务中的性能。将策略训练为条件变分自编码器（CVAE），能够更好地对有噪声的人类演示数据进行建模，准确预测动作序列。实验证明在多个模拟和真实世界的精细操作任务中显著优于之前的模仿学习算法。我们主要的代码可以在https://github.com/tonyzhaozh/act或者https://github.com/MarkFzp/act-plus-plus看到并学习。

1. ACT原理

ACT作为斯坦福最新的Mobile ALOHA系统的最核心的算法：Action Chunking with Transformer。它为啥效果这么好，其实主要的就是Transformer生成式算法，在ACT中主要使用的是CVAE，下面我们就AE（自动编码器）、VAE（变分自动编码器）、CVAE（条件变分自编码器）来进行介绍，这些都是比较老的方法了，但是思想可以借鉴

1.1 AE（自动编码器）

如上图所示，自动编码器主要由两部分组成：编码器（Encoder）和解码器（Decoder）。编码器和解码器可以看作是两个函数，一个用于将高维输入（如图片）映射为低维编码（code），另一个用于将低维编码（code）映射为高维输出（如生成的图片）。这两个函数可以是任意形式，但在深度学习中，我们用神经网络去学习这两个函数。

由上面介绍可以看出，AE的Encoder是将图片映射成“数值编码”，Decoder是将“数值编码”映射成图片。这样存在的问题是，在训练过程中，随着不断降低输入图片与输出图片之间的误差，模型会过拟合，泛化性能不好。也就是说对于一个训练好的AE，输入某个图片，就只会将其编码为某个确定的code，输入某个确定的code就只会输出某个确定的图片

这时候我们转变思路，不将图片映射成“数值编码”，而将其映射成“分布”。这就是VAE的思想。

1.2 VAE（变分自动编码器）

vae和常规自编码器不同的地方在于它的encoder的output不是一个latent vector让decoder去做decode，而是从某个连续的分布里（常见的是高斯分布）采样得到一个随机数or随机向量，然后decode再去针对这个scaler做解码。

常规的自动编码器的任务中没有任何东西被训练来强制获得这样的规律(但是vae就会假设latent vector服从高斯分布)：自动编码器被训练成以尽可能少的损失进行编码和解码，压根就不care latent vector服从什么分布。那么自然,我们是不可能使用一个预定义的随机分布产生随机的input然后又期望decoder能够decode出有意义的东西的.

既然我们不知道latent vector服从什么分布，我们就直接人为对其进行约束满足某种预定义的分布，这个预定义的分布和我们产生随机数的分布保持一致，不就完美解决问题了吗？

整体架构，就是输入图像会解析成均值和标准差两个部分，然后通过正态分布来重参数化获得均值Z，VAE计算以下两方面之间的损失：

• 重构损失（Reconstruction Loss）：这一部分的损失计算的是输入数据与重构数据之间的差异。（输入和输出贴近）

• KL散度（Kullback-Leibler Divergence Loss）：这一部分的损失衡量的是学习到的潜在表示的分布与先验分布（通常假设为标准正态分布）之间的差异。KL散度是一种衡量两个概率分布相似度的指标，VAE通过最小化KL散度来确保学习到的潜在表示的分布尽可能接近先验分布。这有助于模型生成性能的提升，因为它约束了潜在空间的结构，使其更加规整，便于采样和推断。（符合正态分布，来提供随机性）

对于VAE来说，为什么使用重参数化？具体来说，在不使用重参数化的情况下，模型会直接从参数化的分布（例如，正态分布，由均值 μ 和方差 σ2 参数化）中采样，这使得梯度无法回传。

重参数化技巧通过引入一个不依赖于模型参数的外部噪声源（通常是标准正态分布中抽取的），并对这个噪声与模型的均值和方差进行变换，来生成符合目标分布的样本。这样，模型的随机输出就可以表示为模型参数的确定性函数和一个随机噪声的组合。便可以完成梯度回传。详细可以看看这篇文章《【扩散模型基础知识】Diffusion Model中的重参数化和VAE中的重参数化的区别》

代码中的 eps.mul(std).add_(mu) 实现了上述公式的计算，即首先将随机噪声 ϵ 与标准差 σ 相乘，然后将结果加上均值 μ。这样，得到的 z 既包含了模型学习到的分布的特征（通过 μ 和 σ），同时也引入了必要的随机性（通过 ϵ），允许模型通过采样生成多样化的数据。

# VAE class definition
# Encode the input --> reparameterize --> decode
class VAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(VAE, self).__init__()
        self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
        self.decoder = Decoder(latent_dim, hidden_dim, input_dim)
 
    def reparameterize(self, mu, logvar):
        # Reparameterization trick to sample from the distribution represented by the mean and log variance
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return eps.mul(std).add_(mu)
 
    def forward(self, x):
        mu, logvar = self.encoder(x.view(-1, input_dim))
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

1.3 CVAE（条件变分自编码器）

条件变分自编码器（CVAE）是VAE的一种扩展，它能够在生成过程中引入额外的信息（条件），以控制生成结果的某些属性。CVAE通过将条件信息与输入数据一起传递给编码器和解码器，学习条件信息和数据之间的关系。
整体结构和VAE差不多，区别是在将数据输入Encoder时把数据内容与其标签(label)合并(cat)一起输入，将编码(Z)输入Decoder时把编码内容与数据标签(label)合并(cat)一起输入。注意label并不参与Loss计算，CVAE的Loss和VAE的Loss计算方式相同（input和output的MSE Loss以及$N(\mu ,{\sigma }^2)$与$N(0,1)$的KL Loss）。

因此我们在生成数据时，可以先从正态分布采样，然后cat上你想生成的数据的label，一起送入Decoder，就能生成和label类似的数据。

class CVAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim, condition_dim):
        super(CVAE, self).__init__()
        self.fc1 = nn.Linear(input_dim + condition_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        self.fc2 = nn.Linear(latent_dim + condition_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, input_dim)
​
    def encode(self, x, c):
        h = torch.relu(self.fc1(torch.cat([x, c], dim=-1)))
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
​
    def decode(self, z, c):
        h = torch.relu(self.fc2(torch.cat([z, c], dim=-1)))
        x_reconstructed = torch.sigmoid(self.fc3(h))
        return x_reconstructed
​
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
​
    def forward(self, x, c):
        mu, logvar = self.encode(x, c)
        z = self.reparameterize(mu, logvar)
        return self.decode(z, c), mu, logvar
​
# 条件信息c是one-hot编码的数字标签
# 定义损失函数和训练过程的代码同VAE，只是输入多了条件信息

2. CVAE到ACT

相比于CVAE而言，ACT目标是机器人领域的工作，右侧主要就是将图像作为编码器，然后关节角度+动作信息作为解码器，从而生成动作序列。而左侧主要的就是一个编码器，其会用到右侧的解码器输出动作信息以及角度信息，来输出一个z，这个Z和我们CVAE里面的z的概念是一致的。其就类似自编码的操作

下面是整体训练的一些评价loss，其中CVAE这类方法是希望能够来让我的控制符合正态分布，并且具有一定的随机性，提升它的泛化性

下面就是一个动作分块和时间集成的操作，在时间集成中，我们对于一个时间步的动作，会有过k次预测，使用加权的方式对这$k$次预测做平均，权重的设计如下（$w_0$指的是最久的预测）：
$$w_i=\exp (-m\ast i)$$
这是一个以$i$为变量，单调递减的函数，预测发生的越久（即我离预测比较近的情况下，我预测的比较远的部分得分比较小），其权重越大，$m$控制融合新信息的速度，$m$越小，融合的越快。

下面就是整体的训练和测试流程

…详情请参照古月居