Package: gato-tf
Unofficial Gato: A Generalist Agent[Deepmind Publication][arXiv Paper]This repository contains Deepmind's Gato architecture imitation in TensorFlow.Since Deepmind only mentions parts of the architecture in its paper, We still don't know much about the model.However, I believe the paper is enough to imitate the architecture, I'm trying to do that with the open source community's help.Currently, the repository supports the following operations:Gato (viaGato)Transformer (viaTransformer)Patch Position Encodings (viaPatchPositionEncoding)Embedding Function (viaResidualEmbedding)Local Observation Position Encodings (viaLocalPositionEncoding)Tokenizing Continuous Values (viaContinuousValueTokenizer)Shared Embedding (viaDiscreteEmbedding)Action tokens are still a mystery in the paper, I need your help.However, the repository lacks the following miscellaneous.Datasets (most important, Issue:#1,ThomasRochefortB/torch-gato)Pre-trained tokenizers(No longer required because of E2E model)Training strategy (E2E, WIP)But, you can still explore the basic architecture of the Gato based on the paper.Usage$pipinstallgato-tfimporttensorflowastffromgatoimportGato,GatoConfig# Create model instanceconfig=GatoConfig.small()gato=Gato(config)# Fake inputs for Gatoinput_dim=config.input_diminput_ids=tf.concat([# ...# observation 1tf.random.uniform((1,1,input_dim)),# image patch 0tf.random.uniform((1,1,input_dim)),# image patch 1tf.random.uniform((1,1,input_dim)),# image patch 2# ...tf.random.uniform((1,1,input_dim)),# image patch 19tf.fill((1,1,input_dim),value=0.25),# continuous valuetf.fill((1,1,input_dim),value=624.0),# discrete (actions, texts)# observation 2tf.random.uniform((1,1,input_dim)),# image patch 0tf.random.uniform((1,1,input_dim)),# image patch 1tf.random.uniform((1,1,input_dim)),# image patch 2# ...tf.random.uniform((1,1,input_dim)),# image patch 19tf.fill((1,1,input_dim),value=0.12),# continuous valuetf.fill((1,1,input_dim),value=295.0)# discrete (actions, texts)# ...],axis=1)encoding=tf.constant([# 0 - image patch embedding# 1 - continuous value embedding# 2 - discrete embedding (actions, texts)[0,0,0,0,1,2,0,0,0,0,1,2]])row_pos=(tf.constant([[0.00,0.25,0.50,0.75,0,0,0.00,0.25,0.50,0.75,0,0]]),# pos_fromtf.constant([[0.25,0.50,0.75,1.00,0,0,0.25,0.50,0.75,1.00,0,0]])# pos_to)col_pos=(tf.constant([[0.00,0.00,0.00,0.80,0,0,0.00,0.00,0.00,0.80,0,0]]),# pos_fromtf.constant([[0.20,0.20,0.20,1.00,0,0,0.20,0.20,0.20,1.00,0,0]])# pos_to)obs=(tf.constant([[0,1,2,19,20,21,0,1,2,19,20,21]]),# obs tokentf.constant([[1,1,1,1,1,0,1,1,1,1,1,0]])# obs token masking (for action tokens))hidden_states=gato((input_ids,(encoding,row_pos,col_pos),obs))Dataset and Model ArchitecturePaper ReviewsFull Episode SequenceArchitecture VariantsAppendix C.1. Transformer HyperparametersIn the paper, Deepmind tested Gato with 3 architecture variants,1.18B,364M, and79M.I have named them aslarge(),baseline()andsmall()respectively inGatoConfig.HyperparametersLarge(1.18B)Baseline(364M)Small(79M)Transformer blocks24128Attention heads161224Layer width20481536768Feedforward hidden size819261443072Key/value size12812832Residual EmbeddingAppendix C.2. Embedding FunctionThere are no mentions that how many residual networks must be stacked for token embeddings.Therefore, I remain configurable inGatoConfig.Whatever how many residual layers are existing, full-preactivation is a key.The blocks are consisted of:Version 2 ResNet architecture (based on ResNet50V2)GroupNorm (instead of LayerNorm)GeLU (instead of ReLU)Position EncodingsAppendix C.3. Position EncodingsPatch Position EncodingsLikeVision Transformer (ViT)by Google, Gato takes the input images as raster-ordered 16x16 patches.Unlike the Vision Transformer model, however, Gato divides its patch encoding strategy into 2 phases, training and evaluation.For high-performance computation in TensorFlow, I have used the following expressions.$C$ and $R$ mean column and row-wise, and $F$ and $T$ meanfromandtorespectively.$$
\begin{align}
v^R_F &= \begin{bmatrix}
0 & 32 & 64 & 96
\end{bmatrix} \
v^R_T &= \begin{bmatrix}
32 & 64 & 96 & 128
\end{bmatrix} \
v^C_F &= \begin{bmatrix}
0 & 26 & 51 & 77 & 102
\end{bmatrix} \
v^C_T &= \begin{bmatrix}
26 & 51 & 77 & 102 & 128
\end{bmatrix} \
\
P_R &= \begin{cases}
\mathsf{if} \ \mathsf{training} & v^R_F + \mathsf{uniform}(v^R_T - v^R_F) \
\mathsf{otherwise} & \mathsf{round}(\frac{v^R_F + v^R_T}{2})
\end{cases} \
P_C &= \begin{cases}
\mathsf{if} \ \mathsf{training} & v^C_F + \mathsf{uniform}(v^C_T - v^C_F) \
\mathsf{otherwise} & \mathsf{round}(\frac{v^C_F + v^C_T}{2})
\end{cases} \
\
E^R_P &= P_R \cdot 1^{\mathsf{T}}_C \
E^C_P &= 1^{\mathsf{T}}_R \cdot P_C \
\
\therefore E &= E_I + E^R_P + E^C_P
\end{align}
$$Local Observation Position EncodingsIn the definition of Appendix B., text tokens, image patch tokens, and discrete & continuous values are observation tokensWhen Gato receives those values, they must be encoded with their own (local) time steps.Requirementspipinstalltensorflow>=2.11.0ContributingThis repository is still a work in progress.Currently, no downloads and no executables are provided.I welcome many contributors who can help.LicenseLicensed under theMIT license.