“世界模型”（World Models）是谷歌大腦研究科學家 David Ha 和 Swiss AI Lab 負責人 Jürgen Schmidhuber 共同提出的一個為強化學習環境構建的神經網絡模型。“世界模型”論文一經發布就引發了熱烈討論。近日，論文作者之一的 David Ha 公布了重現世界模型實驗的一種方法，本文帶來分步解讀。

“世界模型”（World Models）是谷歌大腦研究科學家 David Ha 和 Swiss AI Lab 負責人 Jürgen Schmidhuber 共同提出的一個為強化學習環境構建的神經網絡模型。世界模型可以通過無監督的方式快速訓練，讓人工智能在 “夢境” 中對外部環境的未來狀態進行預測，大幅提高了完成任務的效率。這篇論文一經發布就引發了熱烈討論。

近日，論文作者之一的 David Ha 公布了重現世界模型實驗的一種方法，本文帶來分步解讀。

已有其他人獨立實現了世界模型。Keras 有一個實現再現了 CarRacing-v0 實驗的一部分。PyTorch 也有一個項目，試圖在 OpenAI Retro Sonic 環境中實現此模型。

關于世界模型，我們已經在原始的交互式的論文中作了充分的討論，在此不再詳細解讀論文。

軟件設置

我使用以下設置測試代碼：

Ubuntu 16.04

Python 3.5.4

TensorFlow 1.8.0

NumPy 1.13.3

VizDoom Gym Levels

OpenAI Gym 0.9.4（Gym 1.0+ breaks VizDoom，可以適用于 CarRacing 任務）

cma 2.2.0

mpi4py 2，請參閱我們為此項目分配的 estool

Jupyter Notebook，用于模型測試和跟蹤進度

我使用 OS X 進行推理，但使用 Google Cloud VM 訓練模型。我在 P100 GPU 實例上訓練了 V 和 M 模型，使用 CMA-ES 在純 CPU 實例上訓練了具有 64 個 cpu-cores（n1-standard-64）的控制器 C。以下將概述哪部分訓練需要 GPU，哪些部分僅使用 CPU，并嘗試降低運行此實驗的成本。

運行預訓練模型

你只需要將 repo 克隆到以 CPU 模式運行的臺式電腦上，使用 repo 中提供的預訓練好的模型來重現結果。不需要 Clould VM 或 GPU。

CarRacing-v0

如果你使用的是 MacBook Pro，我建議將分辨率設置為 “更多空間”，因為 CarRacing-v0 環境呈現的分辨率較高，并且不符合默認的屏幕設置。

在命令行中，進入 carracing 子目錄。嘗試自己玩游戲，在終端中運行 python env.py。你可以使用鍵盤上的四個箭頭鍵來控制汽車。按（上，下）進行加速 / 制動，（左 / 右）進行轉向。

在這個環境下，每次運行都會生成一個新的隨機軌跡。雖然如果我非常仔細地開車，可以輕松獲得 800 分以上的分數，但 900 分以上就很難。斯坦福大學有些學生也發現難以持續高于 900 分。解決這一環境的要求是在 100 個連續的隨機路徑上平均獲得 900 分。

要運行一次預訓練的模型并以完全渲染模式查看 agent，請運行：

python model.py render log/carracing.cma.16.64.best.json

在無渲染模式下運行預訓練的模型 100 次（在無渲染模式下，由于需要使用 OpenGL 為該環境提取像素信息作為觀測值，因此仍然會在屏幕上渲染出更簡單的東西）：

python model.py norender log/carracing.cma.16.64.best.json

該命令將在 100 次試驗后輸出每 100 次試驗的分數，它也會輸出平均分和標準差。平均分數應該在 900 以上。

要在使用 M 生成的環境中運行預訓練好的控制器，并使用 V 進行可視化：

python dream_model.py log/carracing.cma.16.64.best.json

DoomTakeCover-v0

在 doomrnn 目錄中，在由 M 生成的環境中運行 python doomrnn.py

你可以按左、下或右在這個環境中玩游戲。要可視化在真實環境中運行的預訓練模型，請運行：

python model.py doomreal render log/doomrnn.cma.16.64.best.json

請注意，除了重構幀和游戲的實際幀之外，還會修改環境以顯示裁剪后的 64 x 64 px 的幀。要在實際環境中運行模型 100 次并計算平均分數，請運行：

python model.py doomreal norender log/doomrnn.cma.16.64.best.json

你應該在 100 次隨機 episodes 中獲得超過 900 時間步的平均分數。如果你想要獲取 agent 在生成環境中玩游戲的統計信息，則可以使用 doomrnn 代替 doomreal，上面這兩行代碼仍然工作。如果你想更改生成環境的溫度，請修改 doomrnn.py 中的常量 TEMPERATURE，當前的設置是 1.25。

要可視化在生成環境中玩游戲的模型，請運行：

python model.py doomrnn render log/doomrnn.cma.16.64.best.json

從頭開始訓練

DoomTakeCover-v0

我們先討論 VizDoom 實驗，因為這個實驗從頭開始實現需要的計算較少。 由于你可能會更新 repo 中的模型，因此我建議 fork 這個 repo 并在你的 fork 中 clone/update。我建議在 tmux 會話中運行任何命令，以便關閉 ssh 連接時作業仍將在后臺運行。

先創建一個具有～200GB 存儲容量和 220GB RAM 的 64 核 CPU 實例，并在該實例中克隆 repo。 在 doomrnn 目錄中，有一個名為 extract.py 的腳本，它將從 random policy 中提取 200 episodes，并在 doomrnn / record 中保存為 .npz 文件。名為 extract.bash 的 bash 腳本將運行 extract.py 64 次，因此通過運行 bash extract.bash，我們將在 doomrnn / record 中生成 12800 個 .npz 文件。有些情況可能會隨機失敗，所以我們產生了一些額外的數據，盡管最終我們只用了 10000 episodes 來訓練 V 和 M。這個過程需要幾個小時（可能不用 5 小時）。

在 record 子目錄中創建 .npz 文件之后，我創建了一個具有～200GB 存儲空間和 220GB RAM 的 P100 GPU 實例，并在其中克隆了該 repo。我使用 ssh copy 命令 scp 將 CPU 實例中的所有 .npz 文件復制到 GPU 實例的同一個 record 子目錄中。如果 scp 不起作用，你可以使用 gcloud 工具。如果兩個實例在同一個區域中，這應該會非常快，不到一分鐘。將. npz 文件復制到 GPU 機器后，請關閉 CPU 實例。

在 GPU 機器上，運行命令 bash gpu_jobs.bash 來訓練 VAE，預處理錄制的數據集并訓練 MDN-RNN。

這個 gpu_jobs.bash 將按順序運行以下三項：

1）python vae_train.py——這將訓練 VAE，并且在訓練后，模型將被保存在 tf_vae / vae.json 中

2）接下來，它將通過啟動：python series.py，使用預訓練的 VAE 預處理收集的數據。一個新的數據集將在一個名為 series 的子目錄中創建。

3）接下來，保存一個名為 series.npz 的數據集，腳本將使用以下命令啟動 MDN-RNN 訓練器：python rnn_train.py。這將在 tf_rnn / rnn.json 和 tf_initial_z / initial_z.json 中生成一個模型。initial_z.json 文件保存了生成環境時所需的 episode 的初始潛在變量（z）。整個過程可能需要 6~8 小時。

notebook “vae_test.ipynb” 將在訓練數據集上使用 VAE 對輸入 / 重建圖像進行可視化。

V 和 M 訓練完畢后，你將得到 3 個新的 json 文件，現在必須將 vae.json，initial_z.json 和 rnn.json 復制到 tf_modelssub 目錄并覆蓋以前的文件。你應該使用 git add doomrnn / tf_models / *. json 來將這些新模型更新到 git repo，并將更新提交給你的 fork。完成這個操作后，可以關閉 GPU 機器。你需要再次啟動 64-core CPU 實例，然后重新登錄機器。

現在，在 64-core CPU 實例上，運行基于 CMA-ES 的訓練，在 doomrnn 目錄中啟動 python train.py 命令。這將啟動 trainer 并繼續訓練，直到 Ctrl-C 這個工作。控制器 C 將在 M 的生成環境中以 1.25 的 temperature 進行訓練。你可以通過加載正在生成的日志文件的 plot_training_progress.ipynb notebook 來監視進度。經過 200 次迭代（或大約 4-5 小時）后，應該足以獲得可觀的結果，就可以停止了。我在接近 1800 次迭代后停止，盡管它在 200 次之后就沒有真正增加價值了，所以可以不用浪費錢。將 log / *. json 中的所有文件添加到已經 fork 的 repo 中，然后關閉實例。

圖：使用 CMA-ES 訓練 DoomRNN，在生成的環境中記錄 C 的性能

使用你的桌面實例，并再次 pull 已經 fork 的 repo，現在可以運行以下代碼來測試新訓練的 V，M 和 C 模型。

python model.py doomreal render log/doomrnn.cma.16.64.best.json

你可以用 doomrnn 替換 doomreal 或者渲染為 norender 來嘗試生成的環境，或者嘗試 agent 100 次。

CarRacing-v0

CarRacing-v0 的過程與前面的 VizDoom 示例幾乎相同，因此本節主要討論有差異的地方。

由于 CarRacing-v0 的環境是使用 OpenGL 構建的，因此即使在 gym 環境的無渲染模式下，它也依賴于圖形輸出，因此在 CloudVM box 中，必須用 headless X 服務器來包裝命令。你可以在 carracing 目錄下的 extract.bash 文件中看到，在真正的命令之前運行 xvfb-run -a -s "-screen 0 1400x900x24 +extension RANDR"。除此之外，收集數據和訓練 V 和 M 模型的過程與 VizDoom 相同。

請注意，在訓練 VAE 和 MDN-RNN 模型后，必須分別將 vae.json，initial_z.json 和 rnn.json 復制到 vae，initial_z 和 rnn 目錄（而不是像 DoomRNN 中復制到 tf_models），并覆蓋以前的文件，然后照常更新 forked repo。

圖：vae_test.ipynb 用于檢查在 CarRacing-v0 提取的數據上訓練的 VAE。

在這個環境下，我們使用 V 和 M 模型作為模型預測控制（MPC），并在實際環境中訓練控制器 C，而不是在生成環境里訓練。 因此，你需要運行 gce_train.bash，而不是運行 train.py。因為我們在實際環境中訓練，與 DoomRNN 相比，訓練速度較慢。由于是在 tmux 會話中運行訓練，你可以通過在另一個 tmux 會話中并行運行 Jupyter 來使用 plot_training_progress.ipynb notebook 監控進度，這會加載正在生成的日志文件。

圖：使用 CMA-ES 訓練 CarRacing-v0。記錄 C 在實際環境中的性能。

在 150-200 generations（或大約 3 天）之后，應該足以達到～880 的平均分數，非常接近想要的 900 分。如果你對 850 + 分就滿意了，建議這時就可以停下來。定性地說，與最終 agent 達到 900 + 分相比，850-870 的分數并沒有很糟糕，我不想在云上浪費血汗錢。獲得 900 + 分可能需要數周。最終模型以 log / *. json 格式保存，你可以按通常的方式進行測試和查看。

貢獻

有很多很酷的想法可以嘗試，例如，迭代訓練方法，遷移學習，intrinsic motivation，以及其他環境。

圖：A generative noisy pixel pendulum environment?

如果你想擴展代碼并嘗試新的東西，我建議修改代碼并嘗試解決特定的新環境，而不是嘗試改進代碼以同時在多個環境中工作。我發現對于研究工作，如果要嘗試解決困難的環境，通常需要特定的自定義修改。歡迎提交帶有獨立子目錄的 pull request，子目錄可以針對你嘗試解決的特定環境量身定制，并在子目錄的 README.md 文件中說明。