通用智能体: Gato

November 4, 2022 · View on GitHub

开发环境

使用容器：las 虚拟环境：stable_rl_las

概述

尝试实现论文 A Generalist Agent的模型和实验. 智能体 Gato 使用大规模的 Transformer 模型，仅需训练一个通用的权重就能对多模态的不同任务产生通用的策略. 我们实现了粗略的模型，并在 Atari 和 Mujoco 游戏上进行实验.

模型

根据论文中的想法，我们实现了：

Tokenizer
Embedding and Encoding
Sequence
Transformer
Prompt

Tokenizer

针对不同模态的数据，使用不同的标记方法：

文本：通过 SentencePiece 方式编码，将 32000 个 token 编码为整数范围 [0, 32000)
图像：类似 ViT 的编码方式，将图像切成不重叠的 patches, 像素值归一化后除以根号下 patch 的宽度（目前图像的标记与嵌入同时实现，且未归一化）。
观测：
- 连续值: 通过 mu-law 编码到 [-1,1] 的范围后离散化为 1024 个 bins，并移动到范围 [32000, 33024)
- 离散值：[0, 1024) 的整数序列，移动到范围 [32000, 33024)
动作：
- 连续值: 取值范围在 [-1,1] 内，离散化为 1024 个 bins，并移动到范围 [33024, 34048)
- 离散值：[0, 1024) 的整数序列，移动到范围 [33024, 34048)

mu-law：

F(x)=\operatorname{sgn}(x) \frac{\log (|x| \mu+1.0)}{\log (M \mu+1.0)}

Embedding

使用参数化的 embedding 层来对每一个 token 进行嵌入，来生成最后的模型输入。embedding 层根据 token 模态的不同执行不同的操作：

文本、离散或连续观测、行动先通过一个查找表嵌入到可以学习的向量嵌入空间中，同时加上其时间步内不同顺序决定的可学习的位置编码.
$f(\cdot ; \theta)=\text { LookupTable }+\text { LocalPositionEncodings }$
图像块通过 ResNet block 来获得嵌入的向量，同时加上可学习的位置编码.
$f(\cdot ; \theta)=\text { ResNet }+\text { PatchPositionEncodings }$

Value Embedding

对文本、离散或连续观测、行动数据，构建大小为 32000+1024+1024+1 的查找表

Residual Embedding

Appendix C.2. Embedding Function

ResNet 块的结构为：

V2 ResNet architecture
GroupNorm
GELU

根据论文中的图示，我们只使用单一的残差块。

Position Encodings

Appendix C.3. Position Encodings

Patch Position Encodings 图像序列的位置编码
Local Observation Position Encodings 观测值的局部位置编码
Timestep Position Encodings 时间步的位置编码（考虑时间步的越界问题，全部时间步归零化）

Sequence

Appendix B. Agent Data Tokenization Details

针对嵌入数据，将每个 timestep 下的数据以 observation-action 的方式连接在一起，就构成了输入模型的时间序列。其序列化过程的细节如下：

Episodes 按照时间顺序（时间步）呈递给智能体
时间步按照以下顺序设置
- 观测： $\left[y_{1: k}, x_{1: m}, z_{1: n}\right]$ $[y_{1 : k}, x_{1 : m}, z_{1 : n}]$
  - 文本嵌入 token $y_{1: k}$
  - 图像块嵌入 token $x_{1: m}$
  - 张量（离散和连续观测）嵌入 token $z_{1: n}$
- 分隔符：'|' 放在观测后、行动前的分割嵌入 token
- 行动： $a_{1: A}$ 连续或离散值动作的嵌入 token

Token 的完整序列被给出为来自 T 个时间步长的数据的串联：

s_{1: L}=\left[\left[y_{1: k}^1, x_{1: m}^1, z_{1: n}^1,\left.\right|^{\prime}, a_{1: A}^1\right], \ldots,\left[y_{1: k}^T, x_{1: m}^T, z_{1: n}^T,\left.\right|^{\prime}, a_{1: A}^T\right]\right],

Transformer

使用 Decision Transformer 中的 GPT 架构，模型参数为：

Hyperparameters	Large(1.18B)	Baseline(364M)	Small(79M)
Transformer blocks	24	12	8
Attention heads	16	12	24
Layer width	2048	1536	768
Feedforward hidden size	8192	6144	3072
Key/value size	128	128	32

Prompt

训练过程中，对于每批中 25% 的序列，将提示序列加入到训练 batch 中。提示序列为从离线数据集中均匀随机抽取的轨迹的末端，自定义末端长度。

部署过程中，首先生成提示序列，重复与环境交互，将观测与分隔符连接到提示序列末端，自回归地生成动作。

环境需求

安装 Gym

pip install gym==0.23.1

安装 Mujoco 和 D4RL

下载 mujoco210.

mkdir ~/Downloads/
cd ~/Downloads/
wget https://github.com/deepmind/mujoco/releases/download/2.1.0/mujoco210-linux-x86_64.tar.gz
tar -zxvf mujoco210-linux-x86_64.tar.gz
mkdir ~/.mujoco
cp -r mujoco210 ~/.mujoco
gedit ~/.bashrc  
export MUJOCO_KEY_PATH=~/.mujoco${MUJOCO_KEY_PATH}
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/shenxi/.mujoco/mujoco210/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia
source ~/.bashrc

cd ~/.mujoco/mujoco210/bin/
./simulate ../model/humanoid.xml

安装 mujoco_py.

git clone https://github.com/openai/mujoco-py.git
cd mujoco-py
pip install -e .
pip install mujoco_py

sudo apt install libosmesa6-dev
sudo apt-get install libglew-dev glew-utils
sudo apt-get -y install patchelf
sudo apt install gcc

安装 dm_control

pip install dm_control

安装 D4RL

pip install absl-py
pip install matplotlib
git clone https://github.com/rail-berkeley/d4rl.git

找到 setup.py 文件，注释

install_requires=['gym',
                      'numpy',
                      # 'mujoco_py',
                      'pybullet',
                      'h5py',
                      'termcolor',  # adept_envs dependency
                      'click',  # adept_envs dependency
                      # 'dm_control' if 'macOS' in platform() else
                      # 'dm_control @ git+git://github.com/deepmind/dm_control@master#egg=dm_control',
                      'mjrl @ git+git://github.com/aravindr93/mjrl@master#egg=mjrl'],

安装并测试

# installing
pip install -e .

import gym
import d4rl # Import required to register environments

结果

固定随机种子进行训练，并比较不同模型的训练结果，待实现

Dataset	Environment	DT	TT	Gato

改进

采用其他方法对模型进行改进，待实现