TinyWorlds 是一个基于 Google Deepmind 的 Genie 架构 构建的最小化自回归 (autoregressive) 世界模型 (World Models)。

世界模型无法像 VEO3 那样利用无动作的互联网视频进行扩展。Deepmind 的 Genie 通过不使用任何先验动作数据来推断帧之间的动作，从而解决了这个问题。

TinyWorlds 旨在帮助人们理解 Deepmind 可能用于实现可扩展世界模型的那种巧妙的自回归、无监督方法。

目录

• 入门
• 概述
• 构建模块
  • 时空 Transformer
  • 变分自编码器
  • 有限标量量化
• 架构
  • 视频 Tokenizer
  • 动作 Tokenizer
  • 动力学模型
  • TinyWorlds 推理
  • 数据
  • 训练/推理加速
  • 形状标注键
• 贡献

入门

# Installation
git clone https://github.com/AlmondGod/tinyworlds.git
cd tinyworlds
pip install -r requirements.txt
export WANDB_API_KEY=<YOUR_WANDB_API_KEY>
export PYTHONPATH="/workspace/tinyworlds:$PYTHONPATH"

# Training
# 1. download data from huggingface
python scripts/download_assets.py datasets --pattern "zelda_frames.h5"
# 2. run training
python scripts/full_train.py --config configs/training.yaml -- --dataset=ZELDA

# Inference
# 1. pull pretrained sonic checkpoints from huggingface
python scripts/download_assets.py models --suite-name sonic
# 2. run inference
python scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC

概述

为什么需要世界模型？

塑造世界，生成它

世界模型 (World Models) 是一个将环境当前状态映射到环境下一状态的函数。

为了准确预测环境的下一个状态，该函数必须将世界中所有信息压缩为一组定律。

因此，世界模型捕捉了世界的所有固有结构和涌现现象。

所有深度学习，以及所有智能，都在 试图将宇宙压缩为一个模型。一个通过学习关于宇宙如何运行的启发式规则来预测宇宙下一状态重要方面的模型。

我们的宇宙也可以被视为一个世界模型。它是一个遵循一组定律在每一刻执行的状态到状态的映射。人类在这些基础定律之上体验着许多层的涌现行为。

截至 2025 年，基于视频的世界模型已实际应用为：

1. 赋予机器人物理世界理解的皮层
2. 供模型完全在线与物理交互的模拟器
3. 供人类交互的新现实结构的体验

但人类才刚刚处于为我们自己的世界建模的起步阶段。

TinyWorlds 旨在帮助你更好地理解世界建模，并通过贡献学习。

架构概述

TinyWorlds 是一个基于离散 token 的自回归 Transformer，因此我们也可以使用 SOTA（最先进）LLM（大语言模型）技术来改进我们的世界模型。

为什么是离散 token？离散化使我们的动力学预测问题变得容易得多，因为不需要预测图像所在的近无限连续空间，它只需要从我们词汇表（即码本）中的约 1000 个 token 中选择一个。

TinyWorlds 由三个模块组成：

视频 Tokenizer： 这个 Tokenizer（标记器）重建视频序列，中间有一个小的离散瓶颈（我们的视频 token）。这一层将重要信息从视频压缩到 token。

动作 Tokenizer： 这个 Tokenizer 推断两个帧之间的离散动作 token。它通过使用前一帧和一个能看到下一帧的离散动作 token 来重建下一帧的方式进行训练。

动力学模型： 给定过去的动作和帧 token，此模型预测我们的下一帧 token。这应该捕捉我们微型电子游戏世界的物理规律。

构建模块

时空 Transformer

时空 Transformer (Space-Time Transformer) (STT) 是一种用于视频的 Transformer。每个 STT 块包含一个空间注意力层、一个时间注意力层和一个前馈网络 (FFN)。关于自注意力机制的复习，请参阅 Karpathy 的 从零开始构建 GPT 视频。

在空间层中，每个 token 关注同一帧中的所有其他 token。在时间层中，每个 token 关注相同位置但在之前时间步的 token。

FFN 是每个嵌入向量上的多层感知机。受神圣仁慈的启发，我使用了 SwiGLU 作为 FFN。SwiGLU 在 Swish 基础上增加了一个门控线性单元 (GLU)，其计算方式为

$x_t = W_3[\sigma(W_1x + b1) * W_2x + b2] + b3$ (参见 SwiGLU 图表以获取清晰度)

对于常规 STT，我使用了 均方根归一化 (RMSNorm) 作为归一化器，它比零方差归一化对极端异常值更不敏感。在 RMS 中，我们将输入除以

$\sqrt(\epsilon + x / \sum x^2)$.

对于基于动作条件化的 STT，我使用了 特征线性调制 (FiLM)。FiLM 将每个时间步的动作通过 FFN 传递，以将每个动作潜在变量转换为 Gamma ($\gamma$) 和 Beta ($\beta$) 向量。我们的归一化公式则为

$(x - \mu) / \sigma * (1 + \gamma) + \beta$

变分自编码器

*VAE 很复杂，但以下是概述，省略了许多细节

变分自编码器 (Variational Autoencoders) (VAEs) 定义如下：

1. 一个编码器网络 (Encoder)，用于参数化给定数据 $x$ 下潜在变量 $z$ 的近似后验分布 (posterior distribution) $q(z | x)$
2. 一个解码器网络 (Decoder)，用于参数化给定潜在变量 $z$ 下输入数据 $x$ 的似然 (likelihood) $p(x | z)$

VAEs 最大化 $log(p(x | z))$，即在给定编码器提供的潜在变量 $z$ 的情况下，解码器精确重构输入 $x$ 的可能性。

重要的结论是 $z$ 是低维的，因此对于重构，它将压缩来自 $x$ 的所有重要信息。

有限标量量化

由于我们需要一组离散的 token（标记），我们将连续的 $z$ 量化为有限个可能的 $z$ 中的一个。

如果向量是高维空间中的点，有限标量量化（FSQ）是一种将空间划分为超立方体的量化方法，向量落入的超立方体即成为其量化表示。

具体来说，我们在 FSQ 中通过以下步骤对向量进行量化：

1. tanh(x) 将其限制在 [-1,1] 范围内
2. 缩放/平移至 [0, L]
3. 四舍五入到最近的整数（量化步长）
4. 缩放/平移回 [-1,1]

Token 词汇表的大小为 ${L^{D}}$，其中 $L$ 是每个维度的分箱（bins）数量，$D$ 是超立方体的维度。如果有 3 个维度且每个维度有 2 个层级，那么立方体内将有 8 个区域，token 词汇表大小为 8。

FSQ VAE（变分自编码器）允许我们学习结构化的超立方体，用作编码输入信息的 token 词汇表。在我们的上下文中，这些超立方体中的一个可能代表向左移动，另一个代表跳跃，另一个代表蹲下，等等。

为了让梯度流向编码器（因为量化是不可微的），我们将量化后的梯度直接传递给量化前的层。

具体来说，解码器以 $z + stopgrad(z_q - z)$ 作为输入，其中 stopgrad 在 PyTorch 中对应 .detach()。解码器仅使用 $z_q$（因为 $z - z = 0$），但梯度仅针对 $z$ 计算。

架构

视频 Tokenizer

视频 Tokenizer 是一个 FSQ VAE（变分自编码器），它将视频压缩为离散 token。它在实现高质量视频生成的同时降低了动态的维度。

它使用像素混合 2D 卷积将 patch 转换为嵌入。

然后它使用 STTransformer 处理这些嵌入以生成量化 token。

每个视频 token 包含关于其自身 patch 以及同一位置或时间步其他 patch 的信息。

最后，它将视频 token 解码为重建图像。

动作 Tokenizer

动作 Tokenizer 也是一个 FSQ VAE（变分自编码器），它是可扩展性的关键。它通过学习推断两帧之间的动作，使我们能够在没有动作标签的情况下进行训练。然后我们基于这些动作来条件化动态模型。

编码器接收一系列帧并输出帧之间的动作 token。

解码器接收所有之前的帧 $(x_1...x_t-1)$ 和量化动作潜在向量 $(a_1...a_t-1)$ 作为输入，并预测下一帧 $(x_2...x_t)$。

动作 token 应学习编码过去帧与当前帧之间最有意义的变化，这应对应于某种高级动作。

在实践中，动作解码器试图忽略动作，仅从图像中推断。为了抵消这种情况，

1. 我们掩码除第一帧外的所有帧，因此解码器必须学习使用动作序列作为重建信号
2. 我们通过辅助损失鼓励编码器中的批次方差

在推理时，我们将每个键映射到一个动作 token，该 token 用于条件化动态模型，供用户影响视频生成。

动态模型

在时间步 $t$，动态模型应接收 tokenized 视频 token $z_{1..t - 1}$ 和动作 token $a_{1..t - 1}$，并预测下一帧 token $z{t}$。

在实践中，我们像 MaskGIT 和 BERT 一样训练动态模型。

我们掩码一部分 token，并训练我们的模型在给定所有当前和之前帧及动作 token 的条件下预测被掩码的 token。

为了推断给定步骤的动态，我们首先将一个完全掩码的帧附加到我们的上下文序列中。然后，对于 T 个步骤，我们执行以下操作：

1. 预测每个掩码位置的 logits
2. 使用 Softmax 计算 token 概率
3. 从未掩码的位置中采样 k 个最可能的 token
4. 将它们放入上下文张量中，移除对应的掩码 token
5. 重复此过程

我选择了 k 的指数调度（第一步采样约 1 个 token，然后约 2 个，然后约 5 个，然后约 20 个，然后约 50 个，以此类推）

TinyWorlds 推理

给定来自训练分布的初始上下文帧，我们首先将它们 tokenized。

然后我们运行以下循环：

1. 玩家通过在 $[0, |A|]$ 中选择整数指定要使用的 n_actions 个动作 token 之一
2. 使用上下文窗口 c 上的视频 token t-c...t 和动作 token t-c..t 条件化动态模型，并运行动态推理
3. 将预测的视频 token 反 tokenize 为用户的新视频帧

随着动作被传递给模型，tokens 由动态模型预测，我们将它们反 tokenize 为帧展示给用户，我们沿时间维度自回归地重复此过程。

此过程还允许我们一次性预测多个未来帧（受内存和训练分布限制），这可以提高推理质量。

数据

数据从游戏 .mp4s 处理和下采样为 .h5 文件。 您可以在 开始指南 中使用 datasets 命令从 Huggingface TinyWorlds Datasets 下载现有数据集。

可用数据集包括：

1. PicoDoom (picodoom_frames.h5)，Doom 的极简版本
2. Pong (pong_frames.h5)，经典游戏
3. Zelda Ocarina of Time (zelda_frames.h5)，原始 2D 塞尔达游戏之一
4. Pole Position (pole_position_frames.h5)，像素赛车游戏
5. Sonic (sonic_frames.h5)，原始游戏

要创建新数据集，请在 datasets.py 中创建新的 dataclass 并指定 mp4 路径。提交 PR 或直接私信我，将你的数据集上传到 HF 仓库，以便其他人使用 :)

训练/推理加速

TinyWorlds 支持以下 torch 特性以加速训练和/或推理：

1. Torch compile，允许我们对某些预优化的操作（如 attention 和 matmuls）使用更快的 CUDA kernels
2. 分布式数据并行 (DDP)，允许我们通过在每台 GPU 上使用相同模型不同数据的方式使用多 GPU 进行训练
3. 自动混合精度 (AMP)，根据当前节点使用的浮点范围将某些操作从 FP32 扩展到 BF16
4. TF32 训练，允许我们对张量核心优化的 matmuls 和 convolutions 使用 NVIDIA TensorFloat32

形状标注键

所有张量均进行了形状标注，并使用 einops（张量操作库）进行张量操作，缩写如下：

B： 批次大小
T： 时间/序列维度（帧数）
P： 每帧的 patch token 数量
E： 嵌入维度 (d_model)
L： 视频分词器潜在维度
A： 动作分词器潜在维度（动作维度）
D： 每个视频分词器维度的箱数
L^D： 视频分词器的词汇表大小
C： 图像通道
H： 像素网格高度
W： 像素网格宽度
Hp： Patch 网格高度
Wp： Patch 网格宽度
S： Patch 大小

贡献指南

仍有许多待办事项可能带来显著的性能提升……

• 在前馈网络 (Feedforward Network) 中实现混合专家 (Mixture of Experts)
• 尝试使用 RoPE（旋转位置编码）/AliBi 位置编码 (Position Embeddings)
• 尝试不同的优化器 (optimizers) (Muon, SOAP)
• 添加更多数据集（Terraria, Street Fighter, <your favorite retro videogame>）
• 尝试使用 AdaLN-Zero 替代 FiLM（增加预缩放参数）
• 为 MaskGIT 添加新的调度器，如余弦和 Halton
• 将动作分词器中的 mean pool + concat 替换为 length-2 windowed attention + mean
• 通过生成、保存和加载预处理后的图像 patch 嵌入而不是完整帧来加速动力学训练
• 扩展规模：通过在更多 GPU 上训练并添加 FSDP（全分片数据并行）支持扩展到数十亿参数 —— 由 alekseymalakhov11 完成

杂项

TinyWorlds（不包括数据集和外部资源）采用 MIT 许可证。TinyWorlds 是一个独立的研究项目，与 DeepMind 或 Google 无关联、未获其认可或赞助。

aesthetic 设计灵感来源于 Tinygrad 和 Tinygpu

TinyWorlds 快速上手指南

TinyWorlds 是一个基于 Google DeepMind Genie 架构 构建的最小化自回归世界模型。它通过离散 token 预测视频帧，无需先验动作数据即可推断动作，旨在帮助开发者理解可扩展的世界模型技术。

环境准备

• 系统要求: Linux / macOS / Windows (WSL)
• 硬件建议: NVIDIA GPU (CUDA 支持)，显存建议 8GB+ (推理可更低，训练需更高)
• 前置依赖:
  • Python 3.8+
  • Git
  • PyTorch (与 CUDA 版本匹配)
  • Weights & Biases (WandB) 账号 (用于实验追踪)

网络提示: 由于涉及从 GitHub 克隆仓库及从 HuggingFace 下载数据集/模型，国内用户建议使用代理或配置镜像源（如 HF Mirror）以确保连接稳定。

安装步骤

1. 克隆项目

   git clone https://github.com/AlmondGod/tinyworlds.git
   cd tinyworlds
   

2. 安装依赖

   pip install -r requirements.txt
   

3. 配置环境变量 替换 <YOUR_WANDB_API_KEY> 为你的实际 API Key：

   export WANDB_API_KEY=<YOUR_WANDB_API_KEY>
   export PYTHONPATH="/workspace/tinyworlds:$PYTHONPATH"
   

基本使用

1. 模型推理 (Inference)

这是最快体验模型效果的方式。

1. 下载预训练检查点 (以 Sonic 为例):

   python scripts/download_assets.py models --suite-name sonic
   

2. 运行推理:

   python scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC
   

2. 模型训练 (Training)

如需自定义训练，请确保已准备好数据。

1. 下载训练数据 (以 Zelda 为例):

   python scripts/download_assets.py datasets --pattern "zelda_frames.h5"
   

2. 开始训练:

   python scripts/full_train.py --config configs/training.yaml -- --dataset=ZELDA
   

---

注：详细架构原理及更多配置项请参考项目根目录 README 文档。