“世界模型”（World Models）是谷歌大腦研究科學(xué)家 David Ha 和 Swiss AI Lab 負責人 Jürgen Schmidhuber 共同提出的一個為強化學(xué)習(xí)環(huán)境構(gòu)建的神經(jīng)網(wǎng)絡(luò)模型?！笆澜缒Ｐ汀闭撐囊唤?jīng)發(fā)布就引發(fā)了熱烈討論。近日，論文作者之一的 David Ha 公布了重現(xiàn)世界模型實驗的一種方法，本文帶來分步解讀。

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

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

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

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

軟件設(shè)置

我使用以下設(shè)置測試代碼：

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 任務(wù)）

cma 2.2.0

mpi4py 2，請參閱我們?yōu)榇隧椖糠峙涞?estool

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

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

運行預(yù)訓(xùn)練模型

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

CarRacing-v0

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

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

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

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

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

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

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

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

要在使用 M 生成的環(huán)境中運行預(yù)訓(xùn)練好的控制器，并使用 V 進行可視化：

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

DoomTakeCover-v0

在 doomrnn 目錄中，在由 M 生成的環(huán)境中運行 python doomrnn.py

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

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

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

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

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

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

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

從頭開始訓(xùn)練

DoomTakeCover-v0

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

先創(chuàng)建一個具有～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 文件。有些情況可能會隨機失敗，所以我們產(chǎn)生了一些額外的數(shù)據(jù)，盡管最終我們只用了 10000 episodes 來訓(xùn)練 V 和 M。這個過程需要幾個小時（可能不用 5 小時）。

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

在 GPU 機器上，運行命令 bash gpu_jobs.bash 來訓(xùn)練 VAE，預(yù)處理錄制的數(shù)據(jù)集并訓(xùn)練 MDN-RNN。

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

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

2）接下來，它將通過啟動：python series.py，使用預(yù)訓(xùn)練的 VAE 預(yù)處理收集的數(shù)據(jù)。一個新的數(shù)據(jù)集將在一個名為 series 的子目錄中創(chuàng)建。

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

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

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

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

圖：使用 CMA-ES 訓(xùn)練 DoomRNN，在生成的環(huán)境中記錄 C 的性能

使用你的桌面實例，并再次 pull 已經(jīng) fork 的 repo，現(xiàn)在可以運行以下代碼來測試新訓(xùn)練的 V，M 和 C 模型。

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

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

CarRacing-v0

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

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

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

圖：vae_test.ipynb 用于檢查在 CarRacing-v0 提取的數(shù)據(jù)上訓(xùn)練的 VAE。

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

圖：使用 CMA-ES 訓(xùn)練 CarRacing-v0。記錄 C 在實際環(huán)境中的性能。

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

貢獻

有很多很酷的想法可以嘗試，例如，迭代訓(xùn)練方法，遷移學(xué)習(xí)，intrinsic motivation，以及其他環(huán)境。

圖：A generative noisy pixel pendulum environment?

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