LAS1520 / Gato-A-Generalist-Agent

Minimal code for A Generalist Agent

通用智能体: Gato

开发环境

使用容器：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]$
    • 文本嵌入 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

改进

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