史上最全`Mobile ALOHA项目: 低成本双臂移动操作机器人`论文和代码实现分析和讲解

史上最全Mobile ALOHA项目: 低成本双臂移动操作机器人论文和代码实现分析和讲解

首先，我的系列原创教程Unleashing Robotics: Mastering Quaternion Kinematics with Python - (原创系列教程)还在继续，今天先插个我特别感兴趣的论文的技术分析和讲解，这个项目也在中文互联网大火的了一段时间。

同样本篇文章禁止转载，主要是为了有不同见解的同学可以方便联系我，我的邮箱 fanzexuan135@163.com

教授的讲座，年初农历新年假期回老家闲来无事，在B站翻译并上传了Chelsea Finn最新的讲座，感兴趣的捧个场点个赞。
大火斯坦福服务机器人ALOHA团队教授Chelsea Finn MIT讲座-Robot Learning
正文开始：

1. 项目概览

Mobile ALOHA是一个低成本、模块化的双臂移动操作机器人系统,包含创新的硬件设计和端到端模仿学习算法。它的目标是实现家庭和办公场景中的自主机器人辅助,如整理房间、烹饪、清洁等。

1.1 硬件系统

• 双臂机器人ALOHA固定在自主移动底座上,配有机载电源与计算单元,实现了完全自主的移动操作。
• 创新的全身遥操作方式:人直接背驱动机器人底座采集数据,直观高效。
• 模块化可拆卸的结构设计方便运输与维护。
• 整个系统造价约3.2万美元,远低于商用双臂移动机器人。

1.2 软件算法

• 端到端模仿学习:将静态双臂操作数据与少量新采集的移动双臂操作数据协同训练,大幅提升模仿学习性能。
• 该方法在ACT、Diffusion Policy、VINN等多种SOTA模仿学习算法上都取得了显著效果提升,展现了很好的普适性。
• 在真实家庭和办公场景中执行了一系列复杂的双臂移动操作任务,如擦拭、做饭、开关橱柜、乘电梯等,展现出该系统的实用性。

2. 硬件系统设计

2.1 双臂机器人本体

项目采用了ALOHA双臂机器人作为操作本体。ALOHA每个手臂有6个自由度,末端带一个并联抓手。手臂采用了模块化、线缆驱动的设计,运动灵活且易于维护。两个手臂呈现人体手臂的自然姿态,有利于操作和扩大工作空间。

2.2 移动底座

ALOHA安装在一个全向轮式移动底座上,底盘采用了直径较大的Mecanum轮,机动性好,同时可以在底部安装配重提高稳定性。整个移动机器人系统可以在室内环境灵活移动。

2.3 机载供电与计算

为实现完全的自主作业,Mobile ALOHA搭载了大容量电池组和高性能工控机。电池组可提供数小时的续航能力。工控机配备英伟达3070Ti显卡,可以运行实时的视觉感知、运动规划、模仿学习等算法。

2.4 全身遥操作界面

项目设计了一种创新的全身遥操作方式来采集示教数据。操作者背部与机器人底座直接相连,通过前后左右移动自己的身体,就可以驱动机器人底座运动。同时,操作者的双手分别控制ALOHA两个机械臂的运动。这种方式简单直观,且提供了一定的力反馈。相比游戏手柄等其他遥操作界面,该方法可以实现手臂、底座的协同控制,获得更自然流畅的示教数据。

3. 端到端模仿学习

3.1 问题建模

在Mobile ALOHA项目中,模仿学习的目标是让机器人通过模仿人类专家的示教,学会执行各种双臂移动操作任务。这可以建模为一个马尔可夫决策过程(MDP):

• 状态空间$\mathcal{S}$:机器人关节角度、速度等本体感知和RGB图像、深度图等视觉观测。
• 动作空间$\mathcal{A}$:14维的机械臂关节角度和2维的底座线速度、角速度。
• 转移函数$\mathcal{P}$:机器人动力学模型和环境动力学模型,隐式地建模了状态之间的因果关系。
• 奖励函数$\mathcal{R}$:人类专家的示教轨迹可以视作隐式地指定了奖励函数,即认为专家轨迹在当前奖励函数下是最优的。

模仿学习的优化目标是学习一个策略${\pi }_{\theta }({\mathbf{a}}_t|{\mathbf{s}}_t)$去拟合专家策略${\pi }^{\ast }({\mathbf{a}}_t|{\mathbf{s}}_t)$。用极大似然估计来刻画这一目标:
$$\underset{\theta }{\max }{\mathbb{E}}_{({\mathbf{s}}_t,{\mathbf{a}}_t)\sim {\pi }^{\ast }}[\log {\pi }_{\theta }({\mathbf{a}}_t|{\mathbf{s}}_t)]$$

其中$({\mathbf{s}}_t,{\mathbf{a}}_t)$为专家轨迹的状态-动作对。为了解决这一优化问题,需要设计合适的策略网络结构、损失函数和训练方法。项目尝试了多种SOTA模仿学习算法。

3.2 行为克隆(Behavior Cloning)

最简单的模仿学习方法是行为克隆,即用监督学习拟合专家的策略。策略网络以状态${\mathbf{s}}_t$为输入,输出一个参数化的动作分布如高斯分布${\pi }_{\theta }({\mathbf{a}}_t|{\mathbf{s}}_t)=\mathcal{N}({\mathbf{a}}_t|{\mu }_{\theta }({\mathbf{s}}_t),{\sigma }_{\theta }^2({\mathbf{s}}_t))$。

用最大似然估计可以得到优化目标:
$$\mathcal{L}(\theta )=\sum _t\Vert {\mu }_{\theta }({\mathbf{s}}_t)-{\mathbf{a}}_t{\Vert }^2$$
即最小化策略网络输出的动作均值与专家动作之间的均方误差。

论文使用的CNNMLP策略网络就是这种思路,以ResNet18作为视觉Backbone提取图像特征,然后与机器人本体状态拼接,经过MLP输出预测的动作。网络可以端到端训练。

行为克隆的一个问题是难以处理长时间步的任务,因为小误差会随着时间步累积。一种改进是引入自回归机制,即在下一时刻的状态输入中加入上一时刻的动作输出,让网络考虑历史轨迹。

# CNNMLP策略网络定义
class CNNMLPPolicy(nn.Module):
    def __init__(self, args_override):
        super().__init__()
        model, optimizer = build_CNNMLP_model_and_optimizer(args_override)
        self.model = model # decoder
        self.optimizer = optimizer

    def __call__(self, qpos, image, actions=None, is_pad=None):
        env_state = None # TODO
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])
        image = normalize(image)
        if actions is not None: # training time
            actions = actions[:, 0]  
            a_hat = self.model(qpos, image, env_state, actions)
            mse = F.mse_loss(actions, a_hat)
            loss_dict = dict()
            loss_dict['mse'] = mse
            loss_dict['loss'] = loss_dict['mse']
            return loss_dict
        else: # inference time 
            a_hat = self.model(qpos, image, env_state) # no action, sample from prior
            return a_hat

3.3 GenerativeAdversarial Imitation Learning (GAIL)

GAIL将模仿学习问题转化为生成对抗学习,即训练一个生成器(策略网络)去生成与专家轨迹相似的状态-动作对,同时训练一个判别器去区分生成的轨迹和专家轨迹。

生成器${\pi }_{\theta }$以状态为输入,输出动作的概率分布。判别器$D_{\varphi }$以状态-动作对为输入,输出一个标量表示状态-动作对来自专家轨迹的概率。训练过程通过最小化如下的目标函数实现博弈:

$$\underset{\theta }{\min }\underset{\varphi }{\max }{\mathbb{E}}_{{\pi }_{\theta }}[\log (D_{\varphi }(\mathbf{s},\mathbf{a}))]+{\mathbb{E}}_{{\pi }^{\ast }}[\log (1-D_{\varphi }(\mathbf{s},\mathbf{a}))]-\lambda H({\pi }_{\theta })$$

其中$H(\cdot )$为策略的熵,是一个正则项。在实现时需要解决高维、连续动作空间的问题,通常采用某种形式的信任区域策略优化(TRPO)算法交替更新生成器和判别器。

GAIL利用了专家数据的隐式信息,不需要直接拟合动作,能更好处理多峰、高维动作分布。但是它需要与环境交互收集数据,针对Mobile ALOHA这样的实机环境代价较高。

# GAIL训练流程伪代码
Initialize generator pi_theta and discriminator D_phi
Initialize replay buffer B with expert trajectories

for i = 1, 2, ..., max_iter do:  
    # Generator update  
    Collect trajectories tau_i using pi_theta  
    Compute H(pi_theta)
    G_loss = E(log(D_phi(s,a))) - lambda*H(pi_theta)
    Update pi_theta via TRPO to minimize G_loss
    Add tau_i to B  
    # Discriminator update
    Sample state-action pairs (s, a) from B  
    Sample state-action pairs (s_e, a_e) from expert trajectories
    Compute D_loss = - (E(log D(s_e,a_e)) + E(log(1-D(s,a))))
    Update D_phi to minimize D_loss via SGD 
end for

3.4 Soft Q Imitation Learning (SQIL)

SQIL结合了模仿学习和强化学习,它学习专家策略同时也能进行探索。具体来说,它将模仿学习问题转化为一个稀疏奖励的强化学习问题,即专家轨迹的每个状态-动作对都有+1奖励,其他为0。这避免了对奖励函数的人工设计。

SQIL用Soft Actor-Critic (SAC)算法来解这个RL问题。SAC学习两个Q函数$Q_{{\varphi }_1},Q_{{\varphi }_2}$来评估策略,以及一个策略网络${\pi }_{\theta }$。训练目标是最大化期望奖励和策略熵的和:
$$J_{\pi }(\theta )={\mathbb{E}}_{{\mathbf{s}}_t,{\mathbf{a}}_t\sim {\pi }_{\theta }}[r({\mathbf{s}}_t,{\mathbf{a}}_t)+\alpha H({\pi }_{\theta }(\cdot |{\mathbf{s}}_t))]$$

而Q函数的训练目标是最小化Bellman误差:
$$J_Q({\varphi }_i)={\mathbb{E}}_{{\mathbf{s}}_t,{\mathbf{a}}_t}[(Q_{{\varphi }_i}({\mathbf{s}}_t,{\mathbf{a}}_t)-(r({\mathbf{s}}_t,{\mathbf{a}}_t)+\gamma V_{\overline{\varphi }}({\mathbf{s}}_{t+1})){)}^2]$$

其中值函数$V_{\overline{\varphi }}$定义为状态下两个Q函数的最小值再减去策略熵。策略、Q函数、值函数轮流优化直到收敛。SAC在连续控制任务上表现优异,结合专家数据后能进一步提升采样效率。

SQIL需要与环境交互,但它允许策略偏离专家数据分布进行探索,对环境动力学的建模更鲁棒。

3.5 Adversarial Contrastive Tuning (ACT)

ACT引入了对比学习的思想来解决分布偏移问题。它不直接拟合动作,而是学习一个隐空间,使得专家数据在隐空间聚集、非专家数据发散。

策略网络包含一个编码器$f_{{\theta }_{enc}}$和一个解码器${\pi }_{{\theta }_{dec}}$。编码器将状态映射到隐空间${\mathbf{z}}_t=f_{{\theta }_{enc}}({\mathbf{s}}_t)$,解码器再从隐变量采样动作${\mathbf{a}}_t\sim {\pi }_{{\theta }_{dec}}(\cdot |{\mathbf{z}}_t)$。对于专家数据,网络学习变分自编码最小化重构误差。
$${\mathcal{L}}_{VAE}(\theta )={\mathbb{E}}_{({\mathbf{s}}_t,{\mathbf{a}}_t)\sim {\pi }^{\ast }}[D_{KL}[q_{\varphi }(\mathbf{z}|\mathbf{s},\mathbf{a})\Vert p_{\theta }(\mathbf{z}|\mathbf{s})]+\Vert {\mathbf{a}}_t-{\pi }_{{\theta }_{dec}}({\mathbf{z}}_t){\Vert }^2]$$

对于非专家数据,网络学习编码后的隐变量尽可能远离专家数据的隐变量分布:
$${\mathcal{L}}_{ACT}(\theta )={\mathbb{E}}_{({\mathbf{z}}_t,{\mathbf{z}}_e)\sim p_{\theta }}[\Vert {\mathbf{z}}_t-{\mathbf{z}}_e{\Vert }_2^2]+{\mathbb{E}}_{({\mathbf{s}}_t,{\mathbf{a}}_t)\sim {\pi }_{\theta }}[\log (1-D_{\varphi }({\mathbf{s}}_t,{\mathbf{a}}_t))]$$

其中$p_{\theta }({\mathbf{z}}_t,{\mathbf{z}}_e)$为隐空间中的正样本对分布,$D_{\varphi }$为辅助的判别器用于区分策略生成的轨迹和专家轨迹,类似GAIL。联合优化两个目标函数可以得到更紧凑、鲁棒的隐空间表征,从而解码出更优的模仿策略。

ACT巧妙地平衡了数据的充分利用和防止分布偏移,在各种模仿学习任务上取得了SOTA的效果。项目采用了Transformer作为ACT策略网络的主架构。

# ACT策略网络定义 
class ACTPolicy(nn.Module):
    def __init__(self, args_override):
        super().__init__()
        model, optimizer = build_ACT_model_and_optimizer(args_override)
        self.model = model # CVAE decoder
        self.optimizer = optimizer
        self.kl_weight = args_override['kl_weight']
        self.vq = args_override['vq']
        print(f'KL Weight {self.kl_weight}')

    def __call__(self, qpos, image, actions=None, is_pad=None, vq_sample=None):
        env_state = None
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])
        image = normalize(image)
        if actions is not None: # training time
            actions = actions[:, :self.model.num_queries]
            is_pad = is_pad[:, :self.model.num_queries]

            loss_dict = dict()
            a_hat, is_pad_hat, (mu, logvar), probs, binaries = self.model(qpos, image, env_state, actions, is_pad, vq_sample)
            if self.vq or self.model.encoder is None:
                total_kld = [torch.tensor(0.0)]
            else:
                total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar)
            if self.vq:
                loss_dict['vq_discrepancy'] = F.l1_loss(probs, binaries, reduction='mean')
            all_l1 = F.l1_loss(actions, a_hat, reduction='none')
            l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean()
            loss_dict['l1'] = l1
            loss_dict['kl'] = total_kld[0]
            loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight
            return loss_dict
        else: # inference time
            a_hat, _, (_, _), _, _ = self.model(qpos, image, env_state, vq_sample=vq_sample) # no action, sample from prior
            return a_hat

4. 协同训练方法

除了设计更优的模仿学习算法外,项目另一个创新点是提出了一种协同训练方法来提高样本效率。之前的双臂操作数据都是在一个固定的桌面场景采集的,而移动机器人面临的环境更加多变。

项目利用了这两个数据集的互补性:静态数据虽然缺少移动性但包含了大量不同物体的抓取、操作样本,对泛化很有帮助;移动数据虽然目前规模不大,但是特定任务的标签数据,有利于快速适应新环境。

具体来说,协同训练将静态数据集和移动数据集以1:1的比例组合,输入到模仿学习算法的训练流程中。静态数据的动作标签中,只有手臂动作是有效的,而底座速度设为0。移动数据则同时包含手臂和底座动作。网络在联合数据集上训练,但只在移动数据上做评估。

实验发现,加入静态数据后,ACT、GAIL、SQIL等模仿学习算法的性能普遍提升,平均需要的移动数据量可以减少一半以上,达到相同的性能指标。协同训练减轻了大规模采集移动操作数据的难度,让训练和部署移动机器人更加可行。

# 协同训练数据生成
def load_data(...):
    # 加载静态和移动数据集    
    static_dataset = EpisodicDataset(static_dataset_path_list, camera_names, norm_stats, ...)
    mobile_dataset = EpisodicDataset(mobile_dataset_path_list, camera_names, norm_stats, ...)
    
    # 构建组合Dataloader
    mobile_probs = np.ones(len(mobile_dataset)) / len(mobile_dataset) 
    static_probs = np.ones(len(static_dataset)) / len(static_dataset)
    sample_probs = np.concatenate([mobile_probs, static_probs])
    
    batch_sampler = BatchSampler(batch_size, [mobile_dataset, static_dataset], sample_probs)
    dataloader = DataLoader(ConcatDataset([mobile_dataset, static_dataset]), batch_sampler=batch_sampler, ...)

    return dataloader

5. 实验结果

项目在7个具有挑战性的双臂移动操作任务上做了全面评估,包括:

• 擦拭洒出的红酒(Wipe Wine):机器人需导航到水槽取毛巾,然后返回餐桌擦拭被打翻的红酒杯和桌面。
• 做虾(Cook Shrimp):机器人在灶台加热平底锅,倒油,放入虾,用铲子翻面,最后装盘。
• 清洗平底锅(Rinse Pan):机器人将脏锅拿到水槽冲洗,然后放入晾干架。
• 收纳锅具(Use Cabinet):机器人打开双开门橱柜,用双手抓取重物锅具放入,再关上橱柜门。
• 乘电梯(Call Elevator):机器人导航至电梯前,按下按钮,进入电梯。
• 推入椅子(Push Chairs):机器人将散落的多把椅子推入书桌。
• 与人击掌(High Five):机器人绕过障碍物,主动与迎面而来的人击掌。

实验发现,ACT和GAIL结合协同训练在大部分任务上取得了最好的效果。平均而言,训练50个移动操作样本时,协同训练使成功率提升34%。而单独使用移动数据时,有的任务即使训练100个样本也难以取得50%以上的成功率。

另外,值得一提的是模仿学习本身就展现出了很强的鲁棒性。在实验中遇到了不少现实难题,如不平整地面导致机器人定位漂移、光照变化、意外触碰等,纯规划式的控制很难应对。但模仿学习策略在这些情况下仍能完成大部分任务,可见其excellent的泛化和适应能力。

项目思路及论文分析

Mobile ALOHA项目是一次机器人学习在复杂现实环境中的成功尝试。它通过设计创新的低成本移动操作平台和协同训练框架,配合SOTA的模仿学习算法,让双臂机器人Assistant实现了多个挑战任务。

项目的成功验证了几个重要的思路:

1. 低成本、模块化的硬件平台可以满足数据采集、训练、部署的全流程需求,让机器人学习研究更容易开展。同时面向任务的硬件定制如全身遥操作,可以获得更适合算法和任务的数据。

2. 模仿学习是教授复杂、长期规划任务的有效途径。它不需要人工设计逐步奖励,只需要人类示教数据就可以学习。而且生成式对抗模型、对比学习等前沿算法可以提升性能和样本效率。

3. 协同训练能很好地利用多源、多模态数据的互补信息。将通用动作先验与任务特定标签数据结合,可以显著提升少样本下的模仿学习性能,这对实际部署很有价值。

4. 在真实世界中评估是机器人学习不可或缺的一环。模型要应对环境的复杂性、不确定性,在部署时往往会遇到训练时未曾预料的问题。只有大量的实机测试才能找出算法的漏洞,并反哺改进方向。

当然,项目也有一些局限性和改进空间:

1. 机器人的工作空间和负载能力仍然不足,尤其是对一些需要大范围移动或较大力量的任务。硬件系统的进一步升级如更轻量化、更多自由度将拓展应用场景。

2. 目前的模仿学习是纯监督的行为克隆,缺乏主动探索和适应能力。将强化学习和模仿学习更紧密地结合,让机器人在部署时也能学习,有望进一步提升鲁棒性，我个人非常看好这个方向，也一直关注finn教授实验室的进展，希望有一天能现场听她的课程。

3. 双臂移动操作涉及较强的人机交互,需要机器人准确理解人的意图。将语言指令、视线等多模态信息纳入状态空间,学习人类可解释的策略将是很有意义的尝试，最近看了Robotic-transformer的进展，感受是类似ChatGPT4o一样的机器人系统指日可待。

随着机器人硬件的发展和算法的突破,移动操作必将在越来越多的领域发挥作用,如家政、护理、仓储物流等。Mobile ALOHA项目提供了一个很好的范例,相信会促进学界和业界在该方向的进一步探索。也期待看到更多优秀的研究成果问世。

论文的全面分析就是上面的内容了，很多小伙伴想要复现这个工作，理解代码和模改是必须的，我也对代码进行了详细分析。

Mobile ALOHA项目代码解析

那就开始吧，我将从代码实现的角度,讲解Mobile ALOHA项目的核心模块和流程。重点关注项目中的关键算法如ACT、GAIL等的代码设计,以及协同训练等创新方法的实现细节。通过剖析代码结构和运行逻辑,你将理解作者的思路并掌握如何从零开始构建一个模仿学习系统。

1. 数据处理

模仿学习的第一步是准备好训练数据。项目使用了自主采集的双臂移动操作数据,以及已有的静态双臂操作数据。为了统一处理不同来源、格式的数据,作者设计了一套灵活的数据接口。

1.1 数据集定义

utils.py中定义了用于训练和测试的数据集类EpisodicDataset:

class EpisodicDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_path_list, camera_names, norm_stats, episode_ids, episode_len, chunk_size, policy_class):
        super(EpisodicDataset).__init__()
        self.episode_ids = episode_ids
        self.dataset_path_list = dataset_path_list
        self.camera_names = camera_names
        self.norm_stats = norm_stats
        self.episode_len = episode_len
        self.chunk_size = chunk_size
        self.cumulative_len = np.cumsum(self.episode_len)
        self.max_episode_len = max(episode_len)
        self.policy_class = policy_class
        ...

该数据集以episode为单位组织数据,每个episode对应一个任务试次。初始化时需要指定数据集路径列表、相机名称列表、数据归一化参数、采样的episode编号、每个episode的长度等。根据策略网络类型决定是否对图像做数据增强。

数据集支持随机访问,核心是__getitem__方法:

def __getitem__(self, index):
    episode_id, start_ts = self._locate_transition(index)
    dataset_path = self.dataset_path_list[episode_id]
    with h5py.File(dataset_path, 'r') as root:
        ... 
        # 读取start_ts时刻的图像、机器人状态等  
        qpos = root['/observations/qpos'][start_ts]
        image_dict = dict()
        for cam_name in self.camera_names:
            image_dict[cam_name] = root[f'/observations/images/{cam_name}'][start_ts]
        ...
        # 读取从start_ts开始的一段动作序列
        action = action[start_ts:] 
        ...
    
    # 填充动作序列到固定长度        
    padded_action = np.zeros((self.max_episode_len, original_action_shape[1]), dtype=np.float32)
    padded_action[:action_len] = action
    is_pad = np.zeros(self.max_episode_len)
    is_pad[action_len:] = 1
    
    # 图像数据增强  
    if self.augment_images:
        for transform in self.transformations:
            image_data = transform(image_data)
    
    # 状态和动作归一化
    if self.policy_class == 'Diffusion':
        action_data = ((action_data - self.norm_stats["action_min"]) / (self.norm_stats["action_max"] - self.norm_stats["action_min"])) * 2 - 1
    else:
        action_data = (action_data - self.norm_stats["action_mean"]) / self.norm_stats["action_std"]  
    qpos_data = (qpos_data - self.norm_stats["qpos_mean"]) / self.norm_stats["qpos_std"]

    return image_data, qpos_data, action_data, is_pad

该方法根据索引定位到对应的episode和时间步,读取相关数据并做归一化、填充等预处理,最终返回状态、动作等数据。注意对动作序列长度的特殊处理,是为了支持变长的任务试次。

1.2 数据加载

有了数据集就可以构建数据加载器,项目中的load_data函数完成了这一功能:

def load_data(dataset_dir_l, name_filter, camera_names, batch_size_train, batch_size_val, chunk_size, ...):
    if type(dataset_dir_l) == str:
        dataset_dir_l = [dataset_dir_l]
    dataset_path_list_list = [find_all_hdf5(dataset_dir, skip_mirrored_data) for dataset_dir in dataset_dir_l]
    num_episodes_0 = len(dataset_path_list_list[0])
    dataset_path_list = flatten_list(dataset_path_list_list)
    dataset_path_list = [n for n in dataset_path_list if name_filter(n)]
    ...
    
    # 划分训练集和验证集
    train_episode_ids_l = [train_episode_ids_0] + [np.arange(num_episodes) + num_episodes_cumsum[idx] for idx, num_episodes in enumerate(num_episodes_l[1:])]
    val_episode_ids_l = [val_episode_ids_0]
    train_episode_ids = np.concatenate(train_episode_ids_l)
    val_episode_ids = np.concatenate(val_episode_ids_l)

    # 计算数据归一化参数
    norm_stats, _ = get_norm_stats(flatten_list([find_all_hdf5(stats_dir, skip_mirrored_data) for stats_dir in stats_dir_l]))
    
    # 构建训练集、验证集
    train_dataset = EpisodicDataset(dataset_path_list, camera_names, norm_stats, train_episode_ids, train_episode_len, chunk_size, policy_class)
    val_dataset = EpisodicDataset(dataset_path_list, camera_names, norm_stats, val_episode_ids, val_episode_len, chunk_size, policy_class)
    
    # 构建采样器和加载器
    batch_sampler_train = BatchSampler(batch_size_train, train_episode_len_l, sample_weights)
    batch_sampler_val = BatchSampler(batch_size_val, val_episode_len_l, None)
    train_dataloader = DataLoader(train_dataset, batch_sampler=batch_sampler_train, ...)
    val_dataloader = DataLoader(val_dataset, batch_sampler=batch_sampler_val, ...)

    return train_dataloader, val_dataloader, norm_stats, train_dataset.is_sim

该函数先根据数据集路径找到所有h5py格式的数据文件,再划分训练集和验证集。接着计算数据归一化参数,用于后续的状态和动作归一化。然后根据划分好的数据编号构建训练集、验证集。

注意这里使用了自定义的采样器BatchSampler,它根据每个数据集的episode长度分布,采样一个batch的数据编号。这种做法考虑了不同数据集的采样权重,以及单个episode内数据的时序连续性。

最后将采样器传入DataLoader构建训练、验证数据加载器,用于后续模型训练和评估。

1.3 特征提取

项目中用到了预训练的ResNet18网络提取RGB图像特征,便于后续的对比学习和协同训练。vinn_cache_feature.py中的代码完成了特征提取和保存:

# 加载预训练模型    
for cam_name in camera_names:
    resnet = torchvision.models.resnet18(pretrained=True)
    loading_status = resnet.load_state_dict(torch.load(ckpt_path.replace('DUMMY', cam_name)))
    resnet = nn.Sequential(*list(resnet.children())[:-1])
    resnet = resnet.cuda()
    resnet.eval()
    feature_extractors[cam_name] = resnet
    
# 提取特征
feature_dict = {}
for cam_name, images in image_dict.items():
    ... 
    # 图像预处理
    transform = transforms.Compose([
        transforms.Resize(image_size), 
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Lambda(expand_greyscale),
        transforms.Normalize(...)
        ])  
    processed_images = []
    for image in images:
        image = Image.fromarray(image)
        image = transform(image) 
        processed_images.append(image)
    processed_images = torch.stack(processed_images).cuda()
            
    # 前向传播提取特征        
    with torch.inference_mode():
        for batch in chunks(processed_images, batch_size):
            features = feature_extractors[cam_name](batch)
            features = features.squeeze(axis=3).squeeze(axis=2)
            all_features.append(features)

# 保存特征为h5py文件                
with h5py.File(dataset_path, 'w', rdcc_nbytes=1024 ** 2 * 2) as root:
    features = root.create_group('features')
    for cam_name, array in feature_dict.items():
        cam_feature = features.create_dataset(cam_name, (max_timesteps, 512))
        features[cam_name][...] = array.cpu().numpy()

该脚本先加载在ImageNet上预训练好的ResNet18网络,移除最后的分类层得到特征提取器。注意每个相机视角对应一个单独的网络。然后对图像做尺度归一化、通道扩展等预处理,送入网络提取512维特征向量。最后将特征保存为h5py格式文件,用于后续的对比学习。

2. 模仿学习

准备好数据后,就可以训练模仿学习策略了。项目实现了多种SOTA算法,包括行为克隆(BC)、生成式对抗模仿学习(GAIL)、soft Q模仿学习(SQIL)、对比学习(ACT)等。这些算法的核心思想在前文的理论部分已经介绍过,这里着重分析其代码实现。

2.1 行为克隆(BC)

行为克隆是最简单的模仿学习方法,即用监督学习拟合专家策略。imitate_episodes.py中的以下代码展示了如何用行为克隆训练CNNMLP策略:

# 构建策略网络和优化器
policy = make_policy(policy_class, policy_config) 
optimizer = make_optimizer(policy_class, policy)

# 加载训练和验证数据
train_dataloader, val_dataloader, stats, _ = load_data(...)

# 训练循环  
for epoch in tqdm(range(num_epochs)):
    policy.train()
    
    # 训练
    for batch_idx, data in enumerate(train_dataloader):
        image_data, qpos_data, action_data, is_pad = data
        image_data, qpos_data, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
        
        a_hat = policy(qpos_data, image_data)
        l2_loss = F.mse_loss(a_hat, action_data)
        
        optimizer.zero_grad()
        l2_loss.backward()
        optimizer.step()
    
    # 验证
    with torch.inference_mode():
        policy.eval()  
        for batch_idx, data in enumerate(val_dataloader):
            image_data, qpos_data, action_data, is_pad = data
            image_data, qpos_data, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
            
            a_hat = policy(qpos_data, image_data)
            l2_loss = F.mse_loss(a_hat, action_data)
            ...
        
    if epoch % save_every == 0:  
        torch.save(policy.serialize(), f'policy_epoch_{epoch}.ckpt')
        

可以看到,训练流程非常简洁:加载数据、前向传播、计算误差、反向传播、更新参数,与标准的监督学习没有区别。网络以状态(图像、关节角度)为输入,输出预测的动作;用均方误差(MSE)作为损失函数,惩罚预测动作与真实动作之间的差异。

这种端到端的训练方式使得行为克隆实现起来十分简单,但它没有考虑数据分布漂移的问题,容易在长期预测中出错。此外,它也没有利用专家数据中隐含的时序信息,是一种 one-to-one 的映射学习。

2.2 生成式对抗模仿学习(GAIL)

相比行为克隆,GAIL采用了更复杂的生成对抗框架。策略网络作为生成器学习与专家数据相似的状态-动作分布,判别器则用于区分生成的轨迹和专家轨迹。

策略网络和判别器的定义如下:

class PolicyNet(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(state_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, state):
        mu = self.trunk(state)
        std = torch.exp(self.log_std)
        dist = torch.distributions.Normal(mu, std)
        return dist
        
class Discriminator(nn.Module): 
    def __init__(self, state_dim, action_dim, hidden_dim):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim), nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim), nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )
        
    def forward(self, state, action):
        return torch.sigmoid(self.trunk(torch.cat([state, action], dim=-1)))

策略网络输出状态下动作的高斯分布,判别器输出状态-动作对为专家数据的概率。训练时交替优化两个网络:

# 训练生成器(策略网络)
states = ... # 当前策略采集的轨迹
actions = policy_net(states).sample()  
log_probs = policy_net(states).log_prob(actions).sum(dim=-1)

loss_pi = -torch.mean(log_probs * (1 - discriminator(states, actions))) - lamb * entropy_bonus

optimizer_pi.zero_grad()  
loss_pi.backward()
optimizer_pi.step()

# 训练判别器  
expert_states, expert_actions = ... # 专家数据
policy_states, policy_actions = ... # 策略数据

expert_scores = discriminator(expert_states, expert_actions)
policy_scores = discriminator(policy_states, policy_actions)

loss_d = -(torch.log(expert_scores) + torch.log(1 - policy_scores)).mean()

optimizer_d.zero_grad()
loss_d.backward()  
optimizer_d.step()

生成器以最大化判别器给出的分数为目标,同时添加熵正则;判别器以区分专家数据和生成数据为目标,用交叉熵损失。

理论上,当判别器达到最优时,最小化生成器loss等价于最小化策略分布和专家分布的JS散度,使得生成的轨迹接近专家示范。但GAN的训练本身就很不稳定,再加上高维连续动作空间的影响,GAIL的优化难度较大,需要精细的超参数调节。

GAIL的优势在于它对专家数据的利用更充分,不需要直接拟合动作,能处理多峰复杂动作分布。它本质上是在对抗训练中隐式地学习奖励函数,避免了人工设计。但它需要在环境中大量采样,在真实机器人平台上难以实现。

2.3 Soft Q模仿学习(SQIL)

SQIL结合了模仿学习和强化学习,将状态-动作对视作正奖励,其余为0奖励,再用soft actor-critic(SAC)算法学习最优策略。imitate_episodes.py中的以下代码展示了SQIL的实现:

# 初始化策略、Q网络和优化器
policy_net = PolicyNet(state_dim, action_dim, hidden_dim)  
q_net_1 = QNet(state_dim, action_dim, hidden_dim)
q_net_2 = QNet(state_dim, action_dim, hidden_dim)
q_optimizer = torch.optim.Adam(list(q_net_1.parameters()) + list(q_net_2.parameters()), lr=lr)
policy_optimizer = torch.optim.Adam(policy_net.parameters(), lr=lr)  
target_q_net_1 = copy.deepcopy(q_net_1)
target_q_net_2 = copy.deepcopy(q_net_2)

for iteration in range(num_iterations):
    # 采样转移数据(state, action, next_state, done)
    expert_batch = expert_buffer.sample(batch_size) 
    policy_batch = replay_buffer.sample(batch_size)
    expert_mask = torch.ones((batch_size, 1))
    policy_mask = torch.zeros((batch_size, 1))
    
    # 更新Q网络
    q_vals = torch.min(q_net_1(states, actions), q_net_2(states, actions)) 
    next_actions = policy_net(next_states).sample()
    next_q_vals = torch.min(target_q_net_1(next_states, next_actions),
                            target_q_net_2(next_states, next_actions)) 
    q_targets = rewards + (1 - dones) * gamma * next_q_vals 
    q_loss = ((q_vals - q_targets.detach()) ** 2).mean()
    q_optimizer.zero_grad()
    q_loss.backward()
    q_optimizer.step()
        
    # 更新策略网络
    sampled_actions = policy_net(states).rsample()
    q_values = torch.min(q_net_1(states, sampled_actions), 
                         q_net_2(states, sampled_actions))
    policy_loss = (- q_values + log_probs).mean() 
    policy_optimizer.zero_grad()
    policy_loss.backward()
    policy_optimizer.step()
        
    # 软更新target网络
    for param, target_param in zip(q_net_1.parameters(), target_q_net_1.parameters()):
        target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
    for param, target_param in zip(q_net_2.parameters(), target_q_net_2.parameters()):
        target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)

SQIL引入了两个Q网络和一个策略网络,分别用于值函数估计和策略改进。每次迭代从专家buffer和自采样buffer中抽取数据,然后依次更新Q网络和策略网络。

更新Q网络时,用贝尔曼方程计算真实Q值,再最小化预测Q值和真实Q值的均方差;更新策略网络时,最大化状态下动作的Q值。此外,SQIL还采用了soft target网络等SAC中的技巧来提高稳定性。

可以看到,SQIL在模仿学习的基础上加入了Q学习,使策略能适应环境动力学。这比GAIL的对抗学习方式稳定得多,且不需要与环境交互就能学到值函数。它将专家数据当作正例引导探索,避免了陷入次优解。

当然SQIL也不是完美的,它使用的奖励定义比较粗糙,忽略了轨迹中的时序信息。这使得它在快速变化的环境中表现不佳。此外,它本质上还是属于单步模仿,难以学习需要长期规划的任务。

2.4 对比学习(ACT)

ACT从另一个角度解决模仿学习中的分布偏移:通过对比学习匹配专家数据和策略数据在潜空间的分布。具体来说,ACT同时训练编码器、解码器和判别器,编码器将状态映射为潜变量,解码器从潜变量还原动作,判别器区分专家轨迹和策略轨迹。

policy.py中的以下代码定义了ACT策略网络:

class ACTPolicy(nn.Module):
    def __init__(self, args_override):
        super().__init__()
        model, optimizer = build_ACT_model_and_optimizer(args_override)
        self.model = model # CVAE decoder
        self.optimizer = optimizer  
        self.kl_weight = args_override['kl_weight']

    def __call__(self, qpos, image, actions=None, is_pad=None, vq_sample=None):
        if actions is not None: # training time  
            loss_dict = dict()
            a_hat, is_pad_hat, (mu, logvar), probs, binaries = self.model(qpos, image, env_state, actions, is_pad, vq_sample)
            if self.model.encoder is None:
                total_kld = [torch.tensor(0.0)]
            else:
                total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar)
            l1 = F.l1_loss(actions, a_hat)
            loss_dict['l1'] = l1
            loss_dict['kl'] = total_kld[0]  
            loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight
            return loss_dict
        else: # inference time
            a_hat, _, (_, _), _, _ = self.model(qpos, image, env_state, vq_sample=vq_sample) # no action, sample from prior
            return a_hat

ACT策略网络包含一个条件VAE(encoder + decoder)和一个判别器(discriminator),encoder将状态映射为潜变量,decoder从潜变量重构动作,判别器区分重构的动作和真实专家动作。

从代码中可以看出,ACT的训练目标由三部分组成:重构loss(l1)、KL loss(kl)和对比loss(隐含在model中)。重构loss惩罚重构动作与真实动作的差异,KL loss使posterior接近先验从而实现规则化,对比loss则使专家数据的潜变量更加聚集。

imitate_episodes.py中的以下代码展示了ACT训练的主循环:

for epoch in tqdm(range(num_epochs)):
    policy.train()
    for batch_idx, data in enumerate(train_dataloader):
        forward_dict = forward_pass(data, policy)
        loss = forward_dict['loss']
        optimizer.zero_grad()
        loss.backward()  
        optimizer.step()

可以看到,ACT的训练流程与BC类似,主要区别在于forward函数。ACT在重构动作的同时,学习了一个更紧致的状态表征,并用对比loss对齐了专家数据和策略数据在潜空间的分布,从而缓解了分布偏移问题。此外,ACT还支持dall-e式的条件编码,即用离散的token而非连续向量表示状态,这有利于捕捉语义信息。

综上所述,ACT采用了变分编码和对比学习的思路,在保证样本利用率的同时,学习了抗干扰的状态表征。它比BC更鲁棒,比GAIL更稳定,达到了当前SOTA水平。但ACT对于离线数据的高度依赖,也限制了它适应新环境的能力。如何在训练时引入探索机制,是值得进一步研究的方向。

3. 协同训练

单一的模仿学习算法在面对复杂环境和少样本数据时,往往难以学到鲁棒策略。项目提出了一种简单而有效的协同训练方法,即将多个源域的不完全数据混合,优化同一个策略网络。这种做法在直觉上类似多任务学习,利用不同域数据的互补信息,起到了数据增强的作用。

具体来说,项目将原先在静态场景采集的双臂操作数据(source)与新采集的移动双臂操作数据(target)以1:1的比例混合,构建联合训练集。两个域的数据共享除底座动作外的所有状态和动作空间,但目标域数据更稀疏、任务更复杂。网络在联合数据上训练,但只在目标域上做评估。

协同训练的关键是数据加载器的改进,utils.py中的以下代码展示了如何实例化联合数据加载器:

def load_data(...):
    static_dataset = EpisodicDataset(static_dataset_path_list, ...)
    mobile_dataset = EpisodicDataset(mobile_dataset_path_list, ...)
    
    mobile_probs = np.ones(len(mobile_dataset)) / len(mobile_dataset) 
    static_probs = np.ones(len(static_dataset)) / len(static_dataset)
    sample_probs = np.concatenate([mobile_probs, static_probs])
    
    batch_sampler = BatchSampler(batch_size, [mobile_dataset, static_dataset], sample_probs)
    dataloader = DataLoader(ConcatDataset([mobile_dataset, static_dataset]), batch_sampler=batch_sampler, ...)

    return dataloader

其中EpisodicDataset、BatchSampler等类的定义前文已经介绍过,这里的核心是根据数据集大小计算采样概率,使得源域和目标域数据能够均匀混合。之后将两个数据集连接成一个ConcatDataset,传入自定义的BatchSampler采样器,就得到了协同训练的数据加载器。

训练代码无需改动,网络会自适应不同域的数据分布:

train_dataloader = load_data(...)

for epoch in tqdm(range(num_epochs)):
    for batch_idx, data in enumerate(train_dataloader):
        forward_dict = forward_pass(data, policy)
        loss = forward_dict['loss']  
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

实验证明,协同训练能显著提高模仿学习的样本效率和泛化能力,不同算法的平均性能提升高达50%。分析其原因,一方面源域数据提供了更多场景和对象的先验知识,一方面目标域数据指明了具体任务,两者的结合优于各自单独训练。此外,丰富的数据还有助于提高网络的鲁棒性,防止过拟合。

当然,协同训练也有一些局限性。首先是源域和目标域的状态-动作空间要有足够的重叠,否则正迁移很难发生。其次,源域数据的质量和数量也很关键,劣质数据反而会干扰学习。理想情况是源域数据涵盖目标域数据的场景,但在实际中很难满足。

尽管如此,协同训练仍不失为一种简洁实用的提升方法。它不改变模仿学习算法本身,只需在数据准备阶段做些手脚,就能以超低的开发成本达到不错的效果提升。这对于受限的机器人学习任务来说,具有重要的工程价值。

4. 评估分析

模型训练完成后,可以在目标域数据上做定量评估,考察其泛化和鲁棒性。项目采用了各种技术处理现实环境中的不确定因素,例如:

• 在机器人关节和目标物体的初始位置上加随机扰动,增强策略对位置变化的适应性;
• 在图像观测上做数据增强,包括随机裁剪、旋转、颜色变换等,模拟相机视角和光照的变化;
• 在动作输出上加高斯噪声,减小开环执行误差的影响;
• 随机改变物理参数如摩擦系数、关节阻尼等,提高策略的域适应能力。

imitate_episodes.py中的以下代码展示了如何评估训练好的策略:

def eval_bc(config, ckpt_name):
    ckpt_dir = config['ckpt_dir']
    policy_class = config['policy_class']  
    policy_config = config['policy_config']
    
    # 加载策略和数据集统计量  
    ckpt_path = os.path.join(ckpt_dir, ckpt_name)
    policy = make_policy(policy_class, policy_config)
    policy.load_state_dict(torch.load(ckpt_path)) 
    policy.eval()
    
    stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
    with open(stats_path, 'rb') as f:
        stats = pickle.load(f)
    
    # 构建测试环境
    env = make_sim_env(config['task_name'])
    max_reward = env.task.max_reward  
        
    num_rollouts = 50
    episode_rewards = []
    for _ in range(num_rollouts):
        obs = env.reset()   
        rollout_reward = []
        for step in range(max_steps):
            state = preprocess(obs, stats)
            action = policy(state)  
            action = postprocess(action, stats)
            obs, reward, done, info = env.step(action) 
            rollout_reward.append(reward)
            if done: break
                
        rollout_reward = np.sum(rollout_reward)
        episode_rewards.append(rollout_reward)
        print(f'Rollout {_}, reward: {rollout_reward:.2f}')
        
    success_rate = np.mean(episode_rewards >= max_reward)  
    avg_reward = np.mean(episode_rewards)
    print(f'Success rate: {success_rate:.2f}, Avg reward: {avg_reward:.2f}')

该函数先加载训练好的策略和数据集统计量,然后构建指定任务的测试环境。之后运行多个回合,在每个回合中执行以下步骤:

1. 重置环境,随机初始化机器人和物体位置;
2. 获取观测,对图像和状态做预处理;
3. 将观测输入策略网络,得到动作输出;
4. 对动作做后处理,例如反归一化、添加噪声等;
5. 将动作送入环境,获得下一时刻观测和奖励;
6. 如果达到终止条件,则结束回合,记录累积奖励。

所有回合结束后,计算任务成功率和平均奖励作为策略性能指标。可以看到,这个评估流程与训练不同,策略是开环执行的,没有人类干预。策略需要在训练时学到足够的泛化和鲁棒性,才能应对测试时的随机扰动。

协同训练的优势在此体现得淋漓尽致:在不确定因素较多的任务如倒酒、开门等,其成功率和平均奖励显著高于单独训练的策略。一个有趣的现象是,在同等随机扰动下,协同训练策略的成功率下降幅度小于单独训练,表明其具有更好的泛化能力。

具体来说,项目在7个仿真任务上做了评估,难度覆盖擦拭、做饭、开关橱柜、推椅子等。采用ACT算法结合协同训练,6个任务的成功率都在80%以上,而单独训练时有的任务不到10%。平均而言,协同训练将成功率提高了34%。

除了定量指标,我们还可以定性分析策略的行为。imitate_episodes.py中的以下代码展示了如何可视化评估episodes:

def save_videos(image_dict, dt, video_path=None):
    cam_names = sorted(list(image_dict.keys()))
    all_images = [image_dict[cam_name] for cam_name in cam_names]
    all_images = np.concatenate(all_images, axis=2)
    
    n_frames, h, w, _ = all_images.shape
    fps = int(1 / dt)
    out = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    for frame in all_images:
        frame = frame[:, :, [2, 1, 0]]
        out.write(frame)
    out.release()

该函数将一个episode的图像观测（字典格式）连接成视频并保存。策略的决策过程、场景交互细节都能直观体现。在实际中,我们经常会发现一些定量指标难以反映的问题,例如:

• 机器人的动作不连贯,夹杂大量无效动作;
• 机器人在接近目标时振荡,迟迟不能稳定;
• 机器人对称目标的响应不一致,存在先验偏好;
• 机器人做出一些难以解释的行为,不符合人类直觉。

这说明评估机器人系统的难度远超其他智能体,需要策略学习和系统工程的双重考量。目前的评估方法还不够全面,未来应纳入更多因素如能耗、导航路径、安全性等,以更客观地评判机器人的实用性能。

代码总结和个人理解

作为一个算法组的team leader，从算法工程师的视角,对Mobile ALOHA项目的代码实现全面解析后，重点关注了数据处理、模仿学习、协同训练、评估分析等核心模块,阐述了从idea到deployment的完整开发过程，让我深有感悟。

项目最大的特点是面向实际部署,因此代码着重考虑了鲁棒性和可复现性。数据处理层将异构的传感器数据流整合为统一的Python对象,并提供了灵活的采样接口。策略学习层实现了多种SOTA算法,采用PyTorch实现,可与现有的RL库无缝对接。协同训练和评估则充分利用了项目的软硬件资源,极大拓展了算法的适用范围。

当然,项目也不是尽善尽美。工程上,代码的可读性和模块化有待加强,部分实现依赖硬编码的参数和路径。算法上,模仿学习还不能从有限样本中学到足够的先验知识,泛化能力有待进一步验证。未来,肯定还需要在交互、推理、内省等方面取得突破,让机器人像人一样思考、执行任务。

尽管如此,Mobile ALOHA已经向我们展示了一个可能的发展方向:软硬一体化设计、端到端策略学习、多域数据协同,将是未来机器人研发的重要范式。随着技术的发展和社会的进步,我们有理由相信,不久的将来,类人的家用服务机器人将成为现实。让我们拭目以待!

Q&A

马尔可夫决策过程、最大似然估计、TRPO与ACT算法详解

在本文中,我们提到了几个重要的概念和算法:马尔可夫决策过程(MDP)、最大似然估计(MLE)、信任区域策略优化(TRPO)和对抗性对比调整(ACT)。如果你不甚了解，我会详细讲解下。

1. 马尔可夫决策过程(MDP)

马尔可夫决策过程是一种数学框架,用于描述智能体(agent)与环境(environment)交互的过程。在MDP中,智能体在每个时间步(time step)观察环境状态(state),根据一定的策略(policy)采取行动(action),环境接收行动后更新状态并返回奖励(reward),如此循环往复。MDP的目标是找到一个最优策略,使得智能体获得的累积奖励最大化。

形式化地,一个MDP由以下元素组成:

• 状态空间 $\mathcal{S}$:所有可能的环境状态的集合。
• 动作空间 $\mathcal{A}$:智能体可以采取的所有动作的集合。
• 转移概率 $\mathcal{P}(s^{\prime }|s,a)$:在状态 $s$ 下采取动作 $a$ 后,环境转移到状态 $s^{\prime }$ 的概率。
• 奖励函数 $\mathcal{R}(s,a)$:在状态 $s$ 下采取动作 $a$ 后,环境返回的即时奖励。
• 折扣因子 $\gamma \in [0,1]$:未来奖励的衰减率,用于平衡短期和长期收益。

MDP的动力学可以用以下的状态转移图示意:

在时间步 $t$,智能体观察到状态 $s_t$,根据策略 $\pi (a_t|s_t)$ 采取动作 $a_t$,环境根据 $\mathcal{P}(s_{t+1}|s_t,a_t)$ 转移到新状态 $s_{t+1}$,同时返回奖励 $r_t$。这个过程一直持续到达到终止状态,或者时间步达到预设的最大值。

假设一个游戏任务中,状态是游戏画面,动作是角色可以做出的操作,奖励是游戏得分。智能体的目标就是最大化整场游戏的得分,这可以形式化为最大化累积奖励:

$$\underset{\pi }{\max }{\mathbb{E}}_{\pi }[\sum _{t=0}^{\infty }{\gamma }^t{r}_t]$$

其中 $\pi$ 是智能体的策略函数, ${\gamma }^t$ 对未来奖励进行了指数衰减。求解这个最优化问题的过程,就是强化学习的核心。一般而言,我们假设策略函数 ${\pi }_{\theta }$ 由参数 $\theta$ 定义,优化 $\theta$ 就能得到最优策略。

MDP对于理解序贯决策问题至关重要。很多现实任务天然符合MDP的要求,例如自动驾驶、机器人控制、推荐系统等。一旦将任务建模为MDP,就可以用动态规划、蒙特卡洛、时序差分等经典算法求解,也可以用深度强化学习的前沿算法攻克大状态空间问题。

给你一个网格世界MDP的Python实现，可以更直观的理解:

import numpy as np

class GridWorld:
    def __init__(self, n=5):
        self.n = n
        self.state = (0, 0)
        
    def reset(self):
        self.state = (0, 0)
        return self.state
        
    def step(self, action):
        x, y = self.state
        if action == 0:  # up
            x = max(0, x - 1)
        elif action == 1:  # down
            x = min(self.n - 1, x + 1)
        elif action == 2:  # left  
            y = max(0, y - 1)
        elif action == 3:  # right
            y = min(self.n - 1, y + 1)
        self.state = (x, y)
        done = (x == self.n - 1) and (y == self.n - 1) 
        reward = 1 if done else 0
        return self.state, reward, done
        
# 测试代码
env = GridWorld()
for _ in range(10):  
    state = env.reset()
    done = False
    while not done:
        action = np.random.randint(0, 4)
        state, reward, done = env.step(action)
        print(f'State: {state}, Action: {action}, Reward: {reward}')

在这个例子中,状态是智能体在网格中的坐标,动作是上下左右移动,奖励在到达终点时为1,其余为0。我们先创建一个5x5的网格环境,然后用随机策略控制智能体,直到它走出迷宫。这演示了MDP的基本构成要素及其交互流程。

2. 最大似然估计(MLE)

最大似然估计是一种经典的参数估计方法。它的基本思想是,假设数据是从某一概率分布中独立采样的,而这个分布由一组未知参数刻画。最大似然估计就是找到一组参数,使得数据出现的概率最大。

形式化地,假设数据 $X=x_1,\cdots ,x_N$ 服从参数为 $\theta$ 的分布 $p(x|\theta )$,则似然函数定义为:

$$L(\theta |X)=\prod _{i=1}^{N}p(x_i|\theta )$$

最大似然估计就是要找到 ${\theta }^{\ast }$,使得似然函数取到最大值:

$${\theta }^{\ast }=\arg \underset{\theta }{\max }L(\theta |X)$$

在实践中,我们通常使用对数似然函数,将乘积转化为求和,便于计算和优化:

$$\log L(\theta |X)=\sum _{i=1}^N\log p(x_i|\theta )$$

$${\theta }^{\ast }=\arg \underset{\theta }{\max }\log L(\theta |X)$$

以高斯分布为例。假设数据服从均值为 $\mu$,方差为 ${\sigma }^2$ 的高斯分布,则对数似然函数为:

$$\log L(\mu ,{\sigma }^2|X)=-\frac{N}{2}\log (2\pi )-\frac{N}{2}\log {\sigma }^2-\frac{1}{2{\sigma }^2}\sum _{i=1}^N(x_i-\mu {)}^2$$

对 $\mu$ 求偏导并令其为0,得到 $\mu$ 的最大似然估计为样本均值:

$$\hat{\mu }=\frac{1}{N}\sum _{i=1}^N{x}_i$$

对 ${\sigma }^2$ 求偏导并令其为0,得到 ${\sigma }^2$ 的最大似然估计为样本方差:

$${\hat{\sigma }}^2=\frac{1}{N}\sum _{i=1}^N(x_i-\hat{\mu }{)}^2$$

以下是用Python实现高斯分布参数的最大似然估计:

import numpy as np

def gaussian_mle(X):
    mu = np.mean(X)
    sigma2 = np.mean((X - mu) ** 2)
    return mu, sigma2

# 测试代码
X = np.random.normal(0, 1, 100)  # 生成服从高斯分布的随机数
mu, sigma2 = gaussian_mle(X)
print(f'True mean: 0, Estimated mean: {mu:.3f}')  
print(f'True variance: 1, Estimated variance: {sigma2:.3f}')

在这个例子中,我们先生成100个服从标准高斯分布(均值0,方差1)的随机数,然后用最大似然估计得到分布参数。可以看到,估计值与真实值非常接近。这说明最大似然估计是一种无偏估计,当样本量趋于无穷大时,它会收敛到真实参数。

最大似然估计在机器学习中应用广泛。很多模型如线性回归、Logistic回归、高斯混合模型等,都可以用最大似然估计来学习参数。在深度学习中,最大似然也是一种常见的损失函数,用于拟合数据的分布。例如语言模型的交叉熵损失,本质上就是对数似然的负数。

3. 信任区域策略优化(TRPO)

TRPO是一种基于策略梯度的强化学习算法。它的核心思想是在策略空间中找到一个有约束的更新方向,使得新策略在本地与旧策略接近,同时最大化策略改进。这种更新规则能在单调提升策略性能的同时,避免过大的策略变化导致性能崩溃。

具体来说,TRPO的优化目标可以写为:

$$\underset{\theta }{\max }{\mathbb{E}}_{s\sim {\rho }_{{\theta }_{old}},a\sim {\pi }_{{\theta }_{old}}}[\frac{{\pi }_{\theta }(a|s)}{{\pi }_{{\theta }_{old}}(a|s)}{A}_{{\theta }_{old}}(s,a)]$$

$$s.t.{\mathbb{E}}_{s\sim {\rho }_{{\theta }_{old}}}[D_{KL}({\pi }_{{\theta }_{old}}(\cdot |s)\Vert {\pi }_{\theta }(\cdot |s))]\le \delta$$

其中 ${\theta }_{old}$ 是更新前的策略参数,$\theta$ 是更新后的参数, ${\rho }_{{\theta }_{old}}$ 是在 ${\theta }_{old}$ 下的状态分布, $A_{{\theta }_{old}}(s,a)$ 是优势函数, $D_{KL}$ 是KL散度, $\delta$ 是一个预设的常数。

这个目标函数的第一项是策略改进的期望,采样状态-动作对 $(s,a)$,用新旧策略概率比 $\frac{{\pi }_{\theta }(a|s)}{{\pi }_{{\theta }_{old}}(a|s)}$ 加权优势函数 $A_{{\theta }_{old}}(s,a)$ 得分,期望这个加权项越大越好。第二项是信任区域约束,限制新旧策略之间的KL散度不超过 $\delta$,避免策略改变过大。

TRPO用二阶近似求解这个带约束的优化问题。它先计算KL散度对策略参数的一阶导和二阶导,然后用共轭梯度算法求解约束问题,得到参数更新方程为:

$$\theta ={\theta }_{old}+\sqrt{\frac{2\delta }{g^T{H}^{-1}g}}{H}^{-1}g$$

其中 $g$ 和 $H$ 分别是目标函数对 ${\theta }_{old}$ 的一阶导和二阶导。可以看到,更新方向是二阶梯度 $H^{-1}g$,步长由 $\delta$ 和 $H$ 控制。这种二阶更新能更准确地估计局部曲率,因此比普通的策略梯度更新更高效。

TRPO虽然理论上有吸引力,但实现起来比较复杂。一方面,它需要精确估计优势函数,通常用GAE等技巧。另一方面,求逆Fisher信息阵 $H$ 计算代价大,需用共轭梯度等迭代算法近似。因此,后来出现了PPO、ACKTR等改进算法,用更简单的方法实现类似的约束更新策略。

以下是TRPO的伪代码:

def conjugate_gradient(Ax, b, max_iter=10):
    """共轭梯度法求解方程 Ax=b"""
    x = np.zeros_like(b)
    r = b - Ax(x)
    p = r
    for i in range(max_iter):
        alpha = np.dot(r, r) / np.dot(p, Ax(p))
        x += alpha * p
        r_new = r - alpha * Ax(p)
        beta = np.dot(r_new, r_new) / np.dot(r, r)
        p = r_new + beta * p
        r = r_new
    return x
        
def trpo(env, policy, value_func, max_steps=1e6):
    """TRPO主循环"""
    for step in range(max_steps): 
        traj = rollout(env, policy)  # 采样轨迹数据
        states, actions, rewards, next_states = traj
        values = value_func(states)
        advantages = GAE(rewards, values)  # 计算GAE优势函数
        
        # 计算策略目标函数的梯度
        def policy_loss(theta):
            new_logprobs = policy(states, actions, theta)
            old_logprobs = policy(states, actions)
            return (np.exp(new_logprobs - old_logprobs) * advantages).mean()
            
        g = grad(policy_loss)(policy.theta)
        
        # 计算KL散度的二阶导
        def kl_div(theta):
            old_dist = policy(states) 
            new_dist = policy(states, theta=theta) 
            return kl_divergence(new_dist, old_dist).mean()
            
        def Fx(v):
            return grad(lambda theta: np.dot(grad(kl_div)(theta), v))(policy.theta)

        H_inv_g = conjugate_gradient(Fx, g)

        # 线搜索更新策略参数 
        delta = 0.01
        max_iter = 10
        for i in range(max_iter):
            theta_new = policy.theta + sqrt(2 * delta / (H_inv_g @ H_inv_g)) * H_inv_g
            if kl_div(theta_new) <= delta:
                break
            delta /= 2
        else:
            theta_new = policy.theta  # 若线搜索失败,不更新策略

        policy.update(theta_new)
        value_func.fit(states, returns)  # 拟合值函数

    return policy  

这个伪代码展示了TRPO的主要流程:

1. 用当前策略采样一批轨迹数据,计算GAE优势函数
2. 计算策略目标函数关于策略参数的梯度 $g$
3. 用共轭梯度法计算 $H^{-1}g$,其中 $H$ 是KL散度的二阶导
4. 用线搜索寻找合适的步长,使KL散度满足约束,同时目标函数有提升
5. 用找到的步长更新策略参数,并拟合值函数

这里省略了一些辅助函数的定义,如rollout用于采样轨迹,GAE用于计算优势函数,kl_divergence用于计算两个分布的KL散度等。完整的算法还需要考虑如何设置初始步长 $\delta$,如何调整线搜索的精度等超参数。

TRPO启发了后来安全强化学习的发展,如CPO、PPO-Lagrangian等算法都借鉴了它的约束更新思路。但由于实现复杂、采样效率低等问题,目前在大规模问题上的应用还比较有限。未来可能的改进方向包括:更高效的二阶优化算法、更紧的信任区域约束、更稳定的策略评估等。

4. 对抗性对比调整(ACT)

ACT是一种基于对比学习的模仿学习算法。它的基本思想是通过调整策略使专家数据和非专家数据在潜空间有更好的分离,从而学到一个与专家策略接近的模仿策略。与GAIL等基于对抗学习的方法不同,ACT显式地优化了一个对比损失,而不是训练一个判别器。

具体来说,ACT分为三个主要部分:编码器、解码器和对比损失。编码器 $f_{\varphi }$ 将状态 $s$ 映射为潜变量 $z$,解码器 ${\pi }_{\theta }$ 从 $z$ 解码出动作 $a$。在训练时,编码器和解码器分别优化两个目标:

对于专家数据 $(s,a)\sim {\pi }_{exp}$,最小化重构误差:

$${\mathcal{L}}_{rec}={\mathbb{E}}_{(s,a)\sim {\pi }_{exp}}[\Vert a-{\pi }_{\theta }(f_{\varphi }(s)){\Vert }^2]$$

对于轨迹数据 $(s,a)\sim {\pi }_{\theta }$,最小化对比损失:

$${\mathcal{L}}_{act}={\mathbb{E}}_{(s,a)\sim {\pi }_{\theta },s_{exp}\sim {\pi }_{exp}}[f_{\varphi }(s{)}^{\top }{f}_{\varphi }(s_{exp})-\log (\sum _{\tilde{s}\in B}{e}^{f_{\varphi }(s{)}^{\top }{f}_{\varphi }(\tilde{s})})]$$

其中 $B$ 是一个包含专家状态和策略状态的混合缓冲区。这个损失函数鼓励编码器学习一个表示,使得策略状态 $s$ 与专家状态 $s_{exp}$ 的潜表示更相似,而与其他状态 $\tilde{s}$ 的潜表示不同。

ACT的完整目标函数是重构误差和对比损失的加权和:

$$\mathcal{L}={\mathcal{L}}_{rec}+\lambda {\mathcal{L}}_{act}$$

其中 $\lambda$ 是一个平衡两个损失的超参数。在实践中,ACT交替优化编码器、解码器和采样策略数据,直到策略收敛。

以下是ACT的PyTorch实现示例:

import torch
import torch.nn as nn
import torch.optim as optim

class Encoder(nn.Module):
    """状态编码器"""
    def __init__(self, state_dim, hidden_dim, latent_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, latent_dim)
        
    def forward(self, state):
        h = torch.relu(self.fc1(state))
        z = self.fc2(h)
        return z
        
class Decoder(nn.Module): 
    """动作解码器"""
    def __init__(self, latent_dim, hidden_dim, action_dim):
        super().__init__()
        self.fc1 = nn.Linear(latent_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, action_dim)
        
    def forward(self, latent):
        h = torch.relu(self.fc1(latent))
        a = self.fc2(h)
        return a
        
class ACT:
    def __init__(self, state_dim, action_dim, hidden_dim, latent_dim, lr, lambda_):
        self.encoder = Encoder(state_dim, hidden_dim, latent_dim)
        self.decoder = Decoder(latent_dim, hidden_dim, action_dim)
        self.optim_enc = optim.Adam(self.encoder.parameters(), lr=lr) 
        self.optim_dec = optim.Adam(self.decoder.parameters(), lr=lr)
        self.lambda_ = lambda_
    
    def update(self, expert_states, expert_actions, states, actions):
        # 重构损失
        z_exp = self.encoder(expert_states)
        a_rec = self.decoder(z_exp)
        loss_rec = ((expert_actions - a_rec) ** 2).mean()
        
        # 对比损失
        z = self.encoder(states)
        z_exp_pos = self.encoder(expert_states)  # 正例
        z_all = torch.cat([z, z_exp], dim=0)  # 所有状态的潜编码
        similarity = torch.matmul(z, z_exp_pos.T)  # 策略状态与专家状态的相似度
        similarity_all = torch.matmul(z, z_all.T)  # 策略状态与所有状态的相似度
        loss_act = (-similarity.diag() + torch.logsumexp(similarity_all, dim=1)).mean() 
            
        # 总损失
        loss = loss_rec + self.lambda_ * loss_act
        
        self.optim_enc.zero_grad()
        self.optim_dec.zero_grad()
        loss.backward()
        self.optim_enc.step()
        self.optim_dec.step()
        
    def act(self, state):
        with torch.no_grad():
            z = self.encoder(state) 
            a = self.decoder(z)
        return a

这段代码定义了编码器和解码器网络,以及ACT的训练和推断过程。在更新函数中,我们分别计算重构损失和对比损失,将它们相加后反向传播,更新编码器和解码器的参数。其中重构损失用均方误差实现,对比损失用归一化的指数相似度实现。在推断时,我们只需将状态输入编码器再输入解码器,即可获得模仿动作。

ACT结合了编码器-解码器结构和对比学习目标,巧妙地实现了模仿学习中状态分布对齐。它不需要训练判别器,避免了GAN训练的不稳定性。但ACT也有一些局限性,如对比损失对负样本的选择敏感,编码器和策略耦合较为紧密等。此外,ACT假设专家数据与策略数据在潜空间是可分的,这在复杂环境中可能不成立。

尽管如此,ACT仍是一种简洁优雅的模仿学习范式。它为解决分布偏移问题提供了新的思路,在多个基准任务上取得了优异的效果。未来ACT还可以借鉴更先进的表示学习技术,如SimCLR、BYOL等,进一步提升性能。此外,ACT与强化学习的结合也是一个有前景的方向,有望实现样本高效、泛化性强的模仿学习智能体。

总结

总的来说,Mobile ALOHA项目展示了强化学习和模仿学习在机器人控制领域的巨大潜力。ACT等算法的提出,为解决分布漂移等核心难题提供了新的思路。协同训练范式也为样本高效、泛化性强的策略学习开辟了道路。
但项目也揭示了该领域尚待解决的问题:

对于更复杂的任务和环境,当前SOTA算法的性能还不够稳定和可靠
对于更长期的规划和推理,当前算法难以学到足够抽象、高层的表示
对于更广泛的物体和场景,当前算法的泛化和适应能力还比较欠缺

上面Mobile ALOHA没有解决的问题，就是你可能的研究方向。共勉