万事通，亦有所精：一个多功能 Transformer 智能体

引言

我们很高兴能与大家分享“万事通”（Jack of All Trades, JAT）项目，该项目旨在朝着通用智能体的方向迈进。这个项目始于对 Gato (Reed et al., 2022) 工作的开源复现，Gato 提出训练一个能够同时执行视觉语言和决策任务的 Transformer 模型。因此，我们首先构建了一个 Gato 数据集的开源版本。然后，我们在此基础上训练了多模态 Transformer 模型，并引入了多项对 Gato 的改进，以更好地处理序列数据和连续值。

总的来说，该项目取得了以下成果：

• 在多种多样的任务上发布了大量的**专家级强化学习 (RL) 智能体**。
• 发布了 **JAT 数据集**，这是首个用于通用智能体训练的数据集。它包含了数十万条由专家智能体收集的专家轨迹。
• 发布了 **JAT 模型**，这是一个基于 Transformer 的智能体，能够玩视频游戏、控制机器人执行各种任务、在简单的导航环境中理解并执行指令等等！

数据集与专家策略

专家策略

传统强化学习 (RL) 通常是在单一环境中训练策略。利用这些专家策略是构建一个多功能智能体的有效途径。我们选择了多种不同性质和难度的环境，包括 Atari、BabyAI、Meta-World 和 MuJoCo。对于每种环境，我们都训练一个智能体，直到其达到业界顶尖水平。（对于 BabyAI，我们使用了 BabyAI bot）。这些训练出的智能体被称为专家智能体，并已发布在 🤗 Hub 上。您可以在 JAT 数据集卡片中找到所有智能体的列表。

JAT 数据集

我们发布了 JAT 数据集，这是首个用于通用智能体训练的数据集。JAT 数据集包含由上述专家智能体收集的数十万条专家轨迹。要使用此数据集，只需像从 🤗 Hub 加载任何其他数据集一样加载它即可。

>>> from datasets import load_dataset
>>> dataset = load_dataset("jat-project/jat-dataset", "metaworld-assembly")
>>> first_episode = dataset["train"][0]
>>> first_episode.keys()
dict_keys(['continuous_observations', 'continuous_actions', 'rewards'])
>>> len(first_episode["rewards"])
500
>>> first_episode["continuous_actions"][0]
[6.459120273590088, 2.2422609329223633, -5.914587020874023, -19.799840927124023]

除了强化学习数据外，我们还加入了文本数据集，以便为用户提供独特的交互界面。因此，您还会找到用于 Wikipedia、Oscar、OK-VQA 和 Conceptual-Captions 的子集。

JAT 智能体架构

JAT 的架构基于 Transformer，使用了 EleutherAI 的 GPT-Neo 实现。JAT 的独特之处在于其嵌入机制，该机制被设计为能内在地处理序列决策任务。我们将观测值嵌入与动作嵌入以及相应的奖励交错排列。

因此，每个嵌入要么对应一个观测值（与奖励相关），要么对应一个动作。但 JAT 是如何编码这些信息的呢？这取决于数据的类型。如果数据（观测值或动作）是图像（如 Atari 游戏），那么 JAT 使用 CNN。如果是连续向量，JAT 使用线性层。最后，如果是离散值，JAT 使用线性投影层。同样的原理也用于模型输出，具体取决于要预测的数据类型。预测是因果的，将观测值移动一个时间步。这样，智能体必须根据之前所有的观测值和动作来预测下一个动作。

此外，我们觉得训练我们的智能体执行自然语言处理 (NLP) 和计算机视觉 (CV) 任务会很有趣。为此，我们还让编码器能够接受文本和图像数据作为输入。对于文本数据，我们使用 GPT-2 的分词策略进行分词；对于图像，我们使用 ViT 类型的编码器。

考虑到数据的模态会因环境不同而改变，JAT 是如何计算损失的呢？它会为每种模态分别计算损失。对于图像和连续值，它使用 MSE 损失。对于离散值，它使用交叉熵损失。最终的损失是序列中每个元素损失的平均值。等等，这是否意味着我们对预测动作和预测观测值赋予了相同的权重？实际上不是，但我们将在下面详细讨论这一点。

实验与结果

我们在所有 157 个训练任务上评估了 JAT。我们收集了 10 个回合的数据并记录了总奖励。为了便于阅读，我们按领域汇总了结果。

如果用一个数字来总结这些结果，那就是 65.8%，这是与 4 个领域的 JAT 专家相比的平均性能。这表明 JAT 能够在各种各样的任务上模仿专家的表现。让我们更深入地了解一下细节：

• 对于 Atari 57，智能体达到了专家分数的 14.1%，相当于人类水平的 37.6%。它在 21 个游戏中超过了人类水平。
• 对于 BabyAI，智能体达到了专家分数的 99.0%，并且仅在 1 个任务上未能超过专家分数的 50%。
• 对于 Meta-World，智能体达到了专家分数的 65.5%。
• 对于 MuJoCo，智能体达到了专家分数的 84.8%。

最令人印象深刻的是，JAT 仅使用**单一网络**就在所有领域取得了这样的性能。为了衡量这一性能，让我们看看 JAT 在几个任务上的表现：

想试试吗？你可以的！JAT 模型已在 🤗 Hub 上提供！

对于文本任务，我们的模型表现出初步的能力，我们建议读者参阅论文以获取更多细节。

预测观测值的惊人好处

在训练强化学习 (RL) 智能体时，主要目标是最大化未来奖励。但是，如果我们还要求智能体预测它未来会观测到什么呢？这个额外的任务会帮助还是妨碍学习过程？

关于这个问题，有两种对立的观点。一方面，学习预测观测值可以提供对环境更深入的理解，从而实现更好、更快的学习。另一方面，这可能会分散智能体对其主要目标的注意力，导致在观测值和动作预测方面的表现都平平无奇。

为了解决这个争论，我们进行了一项实验，使用了一个结合了观测值损失和动作损失的损失函数，并用一个加权参数 $\kappa$ 来平衡这两个目标。

结果非常值得注意。当 $\kappa$ 太高（0.5）时，预测观测值的额外目标似乎妨碍了学习过程。但当 $\kappa$ 较低时，对学习的影响可以忽略不计，智能体的性能与不将观测值预测作为目标时获得的性能相似。

然而，我们发现在 $\kappa= 0.005$ 附近有一个最佳点，在这里，学习预测观测值实际上提高了智能体的学习效率。我们的研究表明，在学习过程中加入观测值预测是有益的，只要平衡得当。这一发现对这类智能体的设计具有重要意义，凸显了辅助目标在提高学习效率方面的潜在价值。

所以，下次你训练强化学习智能体时，不妨考虑让它预测未来会观测到什么。这可能会带来更好的性能和更快的学习！

结论

在这项工作中，我们介绍了 JAT，一个多功能 Transformer 智能体，能够掌握多种序列决策任务，并在自然语言处理 (NLP) 和计算机视觉 (CV) 任务中展现出初步的能力。对于所有这些任务，JAT 都使用单一网络。我们的贡献包括发布了专家级强化学习智能体、JAT 数据集和 JAT 模型。我们希望这项工作能激励通用智能体领域的未来研究，并为开发更多功能、更强大的 AI 系统做出贡献。

下一步是什么？对研究的展望

我们相信 JAT 项目为通用智能体领域的研究开辟了一个新方向，而我们仅仅触及了皮毛。以下是一些未来工作的想法：

• 改进数据：尽管具有开创性，JAT 数据集仍处于早期阶段。专家轨迹仅来自每个环境的一个专家智能体，这可能会导致一些偏差。尽管我们已尽最大努力达到业界顶尖水平，但一些环境仍然具有挑战性。我们相信收集更多数据和训练更多专家智能体将**大有帮助**。

• 使用离线强化学习 (Offline RL)：JAT 智能体使用基础的行为克隆进行训练。这意味着两件事：（1）我们无法利用次优轨迹；（2）JAT 智能体无法超越专家。我们选择这种方法是为了简单，但我们相信使用离线强化学习可以**真正帮助**提高智能体的性能，而且实现起来不会太复杂。

• 释放更智能的多任务采样策略的全部潜力：目前，JAT 智能体从所有任务中均匀采样数据，但这种方法可能会限制其发展。通过动态调整采样率以专注于最具挑战性的任务，我们可以极大地加速智能体的学习过程，并解锁**显著的性能提升**。

链接

• 📄 论文
• 💻 源代码
• 🗂️ JAT 数据集
• 🤖 JAT 模型

引用

@article{gallouedec2024jack,
    title = {{Jack of All Trades, Master of Some, a Multi-Purpose Transformer Agent}},
    author = {Gallouédec, Quentin and Beeching, Edward and Romac, Clément and Dellandréa, Emmanuel},
    journal = {arXiv preprint arXiv:2402.09844},
    year = {2024},
    url = {https://arxiv.org/abs/2402.09844}
}