匯出 TorchRL 模組

作者: Vincent Moens

注意

要在 notebook 中執行本教程，請在開頭新增一個安裝單元格，包含

!pip install tensordict
!pip install torchrl
!pip install "gymnasium[atari,accept-rom-license]"<1.0.0

簡介

如果在實際場景中無法部署所學到的策略，那麼學習策略的價值就很小。正如其他教程所示，TorchRL 非常注重模組化和可組合性：得益於 tensordict，庫的元件可以以最通用的方式編寫，只需將其簽名抽象為對輸入 TensorDict 的一組操作即可。這可能會給人一種印象，即該庫僅用於訓練，因為典型的低階執行硬體（邊緣裝置、機器人、arduino、Raspberry Pi）不執行 python 程式碼，更不用說安裝 pytorch、tensordict 或 torchrl 了。

幸運的是，PyTorch 提供了一整套生態系統解決方案，用於將程式碼和訓練好的模型匯出到裝置和硬體上，並且 TorchRL 完全能夠與其互動。可以選擇多種後端，包括本教程中示例的 ONNX 或 AOTInductor。本教程簡要介紹瞭如何將訓練好的模型隔離並作為獨立可執行檔案分發，以便匯出到硬體上。

要點回顧

• 訓練後匯出任何 TorchRL 模組；

• 使用各種後端；

• 測試匯出的模型。

快速回顧：一個簡單的 TorchRL 訓練迴圈

在本節中，我們重現了上一個入門教程中的訓練迴圈，並稍作修改以用於由 gymnasium 庫渲染的 Atari 遊戲。我們將繼續使用 DQN 示例，並稍後展示如何使用輸出價值分佈的策略來代替它。

import time
from pathlib import Path

import numpy as np

import torch

from tensordict.nn import (
    TensorDictModule as Mod,
    TensorDictSequential,
    TensorDictSequential as Seq,
)

from torch.optim import Adam

from torchrl._utils import timeit
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer

from torchrl.envs import (
    Compose,
    GrayScale,
    GymEnv,
    Resize,
    set_exploration_type,
    StepCounter,
    ToTensorImage,
    TransformedEnv,
)

from torchrl.modules import ConvNet, EGreedyModule, QValueModule

from torchrl.objectives import DQNLoss, SoftUpdate

torch.manual_seed(0)

env = TransformedEnv(
    GymEnv("ALE/Pong-v5", categorical_action_encoding=True),
    Compose(
        ToTensorImage(), Resize(84, interpolation="nearest"), GrayScale(), StepCounter()
    ),
)
env.set_seed(0)

value_mlp = ConvNet.default_atari_dqn(num_actions=env.action_spec.space.n)
value_net = Mod(value_mlp, in_keys=["pixels"], out_keys=["action_value"])
policy = Seq(value_net, QValueModule(spec=env.action_spec))
exploration_module = EGreedyModule(
    env.action_spec, annealing_num_steps=100_000, eps_init=0.5
)
policy_explore = Seq(policy, exploration_module)

init_rand_steps = 5000
frames_per_batch = 100
optim_steps = 10
collector = SyncDataCollector(
    env,
    policy_explore,
    frames_per_batch=frames_per_batch,
    total_frames=-1,
    init_random_frames=init_rand_steps,
)
rb = ReplayBuffer(storage=LazyTensorStorage(100_000))

loss = DQNLoss(value_network=policy, action_space=env.action_spec, delay_value=True)
optim = Adam(loss.parameters())
updater = SoftUpdate(loss, eps=0.99)

total_count = 0
total_episodes = 0
t0 = time.time()
for data in collector:
    # Write data in replay buffer
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > init_rand_steps:
        # Optim loop (we do several optim steps
        # per batch collected for efficiency)
        for _ in range(optim_steps):
            sample = rb.sample(128)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()
            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())
            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    if max_length > 200:
        break

匯出基於 TensorDictModule 的策略

TensorDict 允許我們以極大的靈活性構建策略：從一個輸出觀察對應的動作價值的常規 Module，我們添加了一個 QValueModule 模組，該模組讀取這些價值並使用某種啟發式方法（例如 argmax 呼叫）計算動作。

然而，在我們的案例中有一個小小的技術細節：環境（實際的 Atari 遊戲）不返回灰度、84x84 的影像，而是原始螢幕尺寸的彩色影像。我們附加到環境的 transforms 確保影像可以被模型讀取。我們可以看到，從訓練的角度來看，環境和模型之間的界限是模糊的，但在執行時事情就清晰多了：模型應該負責將輸入資料（影像）轉換成可以被我們的 CNN 處理的格式。

在這裡，tensordict 的魔力再次為我們掃清障礙：碰巧大多數區域性的（非遞迴的）TorchRL transforms 既可以用作環境 transforms，也可以用作 Module 例項內的預處理塊。讓我們看看如何將它們前置到我們的策略中

policy_transform = TensorDictSequential(
    env.transform[
        :-1
    ],  # the last transform is a step counter which we don't need for preproc
    policy_explore.requires_grad_(
        False
    ),  # Using the explorative version of the policy for didactic purposes, see below.
)

我們建立一個偽輸入，並將其與策略一起傳遞給 export()。這將生成一個“原始”的 python 函式，該函式將讀取我們的輸入張量並輸出一個動作，不包含任何對 TorchRL 或 tensordict 模組的引用。

一個好的做法是呼叫 select_out_keys() 來告知模型我們只需要特定的輸出集合（以防策略返回多個張量）。

fake_td = env.base_env.fake_tensordict()
pixels = fake_td["pixels"]
with set_exploration_type("DETERMINISTIC"):
    exported_policy = torch.export.export(
        # Select only the "action" output key
        policy_transform.select_out_keys("action"),
        args=(),
        kwargs={"pixels": pixels},
        strict=False,
    )

視覺化策略可能會非常有啟發性：我們可以看到第一個操作是 permute、div、unsqueeze、resize，接著是卷積層和 MLP 層。

print("Deterministic policy")
exported_policy.graph_module.print_readable()

作為最後檢查，我們可以使用一個虛擬輸入執行策略。輸出（對於單個影像）應該是一個從 0 到 6 的整數，表示在遊戲中要執行的動作。

output = exported_policy.module()(pixels=pixels)
print("Exported module output", output)

關於匯出 TensorDictModule 例項的更多詳細資訊，請參閱 tensordict 文件。

注意

匯出接受和輸出巢狀鍵的模組是完全可以的。對應的 kwargs 將是鍵的 “_”.join(key) 版本，即 (“group0”, “agent0”, “obs”) 鍵將對應於 “group0_agent0_obs” 關鍵字引數。鍵衝突（例如 (“group0_agent0”, “obs”) 和 (“group0”, “agent0_obs”)）可能導致未定義的行為，應不惜一切代價避免。顯然，鍵名也應始終產生有效的關鍵字引數，即它們不應包含空格或逗號等特殊字元。

torch.export 還有許多其他特性，我們將在下文進一步探討。在此之前，讓我們先簡要探討一下在測試時推理環境下的探索和隨機策略，以及迴圈策略。

使用隨機策略

您可能已經注意到，上面我們使用了 set_exploration_type 上下文管理器來控制策略的行為。如果策略是隨機的（例如，策略輸出動作空間的分佈，就像 PPO 或其他類似的 on-policy 演算法中那樣）或具有探索性（附加了探索模組，如 E-Greedy、加性高斯或 Ornstein-Uhlenbeck），我們可能希望或不希望在匯出的版本中使用該探索策略。幸運的是，匯出工具可以理解該上下文管理器，並且只要匯出發生在正確的上下文管理器中，策略的行為就應該與所示一致。為了演示這一點，讓我們嘗試另一種探索型別

with set_exploration_type("RANDOM"):
    exported_stochastic_policy = torch.export.export(
        policy_transform.select_out_keys("action"),
        args=(),
        kwargs={"pixels": pixels},
        strict=False,
    )

與之前的版本不同，我們匯出的策略現在應該在呼叫棧的末尾包含一個隨機模組。實際上，最後的三個操作是：生成一個介於 0 到 6 之間的隨機整數，使用一個隨機掩碼，並根據掩碼中的值選擇網路輸出或隨機動作。

print("Stochastic policy")
exported_stochastic_policy.graph_module.print_readable()

使用迴圈策略

另一種典型的用例是迴圈策略，它將輸出一個動作以及一個或多個迴圈狀態。LSTM 和 GRU 是基於 CuDNN 的模組，這意味著它們的行為會與常規 Module 例項不同（匯出工具可能無法很好地跟蹤它們）。幸運的是，TorchRL 提供了這些模組的 Python 實現，可以在需要時與 CuDNN 版本互換使用。

為了展示這一點，讓我們編寫一個依賴於 RNN 的原型策略

from tensordict.nn import TensorDictModule
from torchrl.envs import BatchSizeTransform
from torchrl.modules import LSTMModule, MLP

lstm = LSTMModule(
    input_size=32,
    num_layers=2,
    hidden_size=256,
    in_keys=["observation", "hidden0", "hidden1"],
    out_keys=["intermediate", "hidden0", "hidden1"],
)

如果 LSTM 模組不是基於 Python 而是 CuDNN 的（LSTM），則可以使用 make_python_based() 方法來使用 Python 版本。

lstm = lstm.make_python_based()

現在我們來建立策略。我們將兩個修改輸入形狀的層（unsqueeze/squeeze 操作）與 LSTM 和一個 MLP 結合起來。

recurrent_policy = TensorDictSequential(
    # Unsqueeze the first dim of all tensors to make LSTMCell happy
    BatchSizeTransform(reshape_fn=lambda x: x.unsqueeze(0)),
    lstm,
    TensorDictModule(
        MLP(in_features=256, out_features=5, num_cells=[64, 64]),
        in_keys=["intermediate"],
        out_keys=["action"],
    ),
    # Squeeze the first dim of all tensors to get the original shape back
    BatchSizeTransform(reshape_fn=lambda x: x.squeeze(0)),
)

和之前一樣，我們選擇相關的鍵

recurrent_policy.select_out_keys("action", "hidden0", "hidden1")
print("recurrent policy input keys:", recurrent_policy.in_keys)
print("recurrent policy output keys:", recurrent_policy.out_keys)

現在我們準備好匯出了。為此，我們構建偽輸入並將其傳遞給 export()

fake_obs = torch.randn(32)
fake_hidden0 = torch.randn(2, 256)
fake_hidden1 = torch.randn(2, 256)

# Tensor indicating whether the state is the first of a sequence
fake_is_init = torch.zeros((), dtype=torch.bool)

exported_recurrent_policy = torch.export.export(
    recurrent_policy,
    args=(),
    kwargs={
        "observation": fake_obs,
        "hidden0": fake_hidden0,
        "hidden1": fake_hidden1,
        "is_init": fake_is_init,
    },
    strict=False,
)
print("Recurrent policy graph:")
exported_recurrent_policy.graph_module.print_readable()

AOTInductor：將策略匯出到不依賴 PyTorch 的 C++ 二進位制檔案

AOTInductor 是一個 PyTorch 模組，允許您將模型（策略或其他）匯出到不依賴 PyTorch 的 C++ 二進位制檔案。當您需要在沒有安裝 PyTorch 的裝置或平臺上部署模型時，這特別有用。

這是一個如何使用 AOTInductor 匯出策略的示例，靈感來自 AOTI 文件

from tempfile import TemporaryDirectory

from torch._inductor import aoti_compile_and_package, aoti_load_package

with TemporaryDirectory() as tmpdir:
    path = str(Path(tmpdir) / "model.pt2")
    with torch.no_grad():
        pkg_path = aoti_compile_and_package(
            exported_policy,
            # Specify the generated shared library path
            package_path=path,
        )
    print("pkg_path", pkg_path)

    compiled_module = aoti_load_package(pkg_path)

print(compiled_module(pixels=pixels))

使用 ONNX 匯出 TorchRL 模型

注意

要執行指令碼的這一部分，請確保已安裝 pytorch onnx

!pip install onnx-pytorch
!pip install onnxruntime

您還可以在 PyTorch 生態系統中找到更多關於使用 ONNX 的資訊 在此。以下示例基於此文件。

在本節中，我們將展示如何以不依賴 PyTorch 的方式匯出我們的模型，使其可以在沒有安裝 PyTorch 的環境中執行。

網路上有很多資源解釋瞭如何使用 ONNX 將 PyTorch 模型部署到各種硬體和裝置上，包括 Raspberry Pi、NVIDIA TensorRT、iOS 和 Android。

我們訓練所用的 Atari 遊戲可以使用 ALE 庫 在沒有 TorchRL 或 gymnasium 的情況下獨立執行，因此為我們提供了一個關於使用 ONNX 可以實現什麼的良好示例。

讓我們看看這個 API 的樣子

from ale_py import ALEInterface, roms

# Create the interface
ale = ALEInterface()
# Load the pong environment
ale.loadROM(roms.Pong)
ale.reset_game()

# Make a step in the simulator
action = 0
reward = ale.act(action)
screen_obs = ale.getScreenRGB()
print("Observation from ALE simulator:", type(screen_obs), screen_obs.shape)

from matplotlib import pyplot as plt

plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
plt.imshow(screen_obs)
plt.title("Screen rendering of Pong game.")

匯出到 ONNX 與上面的 Export/AOTI 非常相似

import onnxruntime

with set_exploration_type("DETERMINISTIC"):
    # We use torch.onnx.dynamo_export to capture the computation graph from our policy_explore model
    pixels = torch.as_tensor(screen_obs)
    onnx_policy_export = torch.onnx.dynamo_export(policy_transform, pixels=pixels)

現在我們可以將程式儲存到磁碟並載入它

with TemporaryDirectory() as tmpdir:
    onnx_file_path = str(Path(tmpdir) / "policy.onnx")
    onnx_policy_export.save(onnx_file_path)

    ort_session = onnxruntime.InferenceSession(
        onnx_file_path, providers=["CPUExecutionProvider"]
    )

onnxruntime_input = {ort_session.get_inputs()[0].name: screen_obs}
onnx_policy = ort_session.run(None, onnxruntime_input)

使用 ONNX 執行 rollout

我們現在有了一個可以執行我們策略的 ONNX 模型。讓我們將其與原始的 TorchRL 例項進行比較：由於 ONNX 版本更輕量，它應該比 TorchRL 版本更快。

def onnx_policy(screen_obs: np.ndarray) -> int:
    onnxruntime_input = {ort_session.get_inputs()[0].name: screen_obs}
    onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
    action = int(onnxruntime_outputs[0])
    return action


with timeit("ONNX rollout"):
    num_steps = 1000
    ale.reset_game()
    for _ in range(num_steps):
        screen_obs = ale.getScreenRGB()
        action = onnx_policy(screen_obs)
        reward = ale.act(action)

with timeit("TorchRL version"), torch.no_grad(), set_exploration_type("DETERMINISTIC"):
    env.rollout(num_steps, policy_explore)

print(timeit.print())

請注意，ONNX 也提供了直接最佳化模型的可能性，但這超出了本教程的範圍。

結論

在本教程中，我們學習瞭如何使用各種後端匯出 TorchRL 模組，例如 PyTorch 內建的匯出功能、AOTInductor 和 ONNX。我們演示瞭如何匯出在 Atari 遊戲上訓練的策略，並使用 ALE 庫在不依賴 PyTorch 的環境中執行它。我們還比較了原始 TorchRL 例項與匯出的 ONNX 模型的效能。

主要收穫

• 匯出 TorchRL 模組可以在未安裝 PyTorch 的裝置上進行部署。

• AOTInductor 和 ONNX 提供了用於匯出模型的替代後端。

• 最佳化 ONNX 模型可以提高效能。

進一步閱讀和學習步驟

• 查閱 PyTorch 關於 匯出功能、AOTInductor 和 ONNX 的官方文件，瞭解更多資訊。

• 嘗試在不同裝置上部署匯出的模型。

• 探索 ONNX 模型的最佳化技術以提高效能。