注意

前往結尾下載完整的範例程式碼。

TorchRL 模組入門¶

注意

若要在筆記本中執行本教學課程，請在開頭新增一個包含以下內容的安裝儲存格

!pip install tensordict
!pip install torchrl

強化學習旨在建立能夠有效處理特定任務的策略。策略可以採用多種形式，從觀察空間到行動空間的可微分映射，到更臨時的方法，例如在為每個可能的行動計算的值列表上進行 argmax。策略可以是確定性的或隨機的，並且可能包含複雜的元素，例如遞迴神經網路 (RNN) 或轉換器。

適應所有這些情況可能相當複雜。在本簡潔的教學課程中，我們將深入研究 TorchRL 在策略建構方面的核心功能。我們將主要關注兩種常見情境中的隨機和 Q 值策略：使用多層感知器 (MLP) 或卷積神經網路 (CNN) 作為主幹。

TensorDictModules¶

與環境如何與 TensorDict 實例交互作用類似，用於表示策略和價值函數的模組也執行相同的操作。核心思想很簡單：將標準 Module (或任何其他函數) 封裝在一個類別中，該類別知道哪些條目需要被讀取並傳遞到模組，然後使用分配的條目記錄結果。為了說明這一點，我們將使用最簡單的策略：從觀察空間到行動空間的確定性映射。為了獲得最大的通用性，我們將使用在先前的教學課程中實例化的 Pendulum 環境的 LazyLinear 模組。

import torch

from tensordict.nn import TensorDictModule
from torchrl.envs import GymEnv

env = GymEnv("Pendulum-v1")
module = torch.nn.LazyLinear(out_features=env.action_spec.shape[-1])
policy = TensorDictModule(
    module,
    in_keys=["observation"],
    out_keys=["action"],
)

這就是執行我們策略所需的全部！使用延遲模組可以讓我們繞過獲取觀察空間形狀的需求，因為該模組會自動確定它。這個策略現在就可以在環境中運行了

rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

特殊化的封裝器¶

為了簡化 Actor、# ProbabilisticActor、# ActorValueOperator 或 # ActorCriticOperator 的整合，TorchRL 提供了一些特殊化的封裝器。例如，Actor 提供了 in_keys 和 out_keys 的預設值，使得與許多常見環境的整合變得簡單。

from torchrl.modules import Actor

policy = Actor(module)
rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

可用特殊化 TensorDictModule 的列表可在API 參考中找到。

網路¶

TorchRL 還提供了可以使用的常規模組，而無需使用 tensordict 功能。您將遇到的兩個最常見的網絡是 MLP 和 ConvNet (CNN) 模組。我們可以用這些模組之一替換我們的策略模組

from torchrl.modules import MLP

module = MLP(
    out_features=env.action_spec.shape[-1],
    num_cells=[32, 64],
    activation_class=torch.nn.Tanh,
)
policy = Actor(module)
rollout = env.rollout(max_steps=10, policy=policy)

TorchRL 還支援基於 RNN 的策略。由於這是一個更技術性的主題，因此在一個單獨的教學中進行了介紹。

機率策略¶

諸如 PPO 之類的策略優化算法要求策略具有隨機性：與上面的例子不同，該模組現在將從觀察空間到參數空間的映射進行編碼，該參數空間編碼了可能動作的分布。 TorchRL 通過將諸如從參數建立分布、從該分布中抽樣和檢索對數機率之類的各種操作分組在一個類別下，來促進此類模組的設計。在這裡，我們將使用三個組件來構建一個依賴於常規正態分布的 actor

一個 MLP 後端，讀取大小為 [3] 的觀察值並輸出大小為 [2] 的單個張量；
一個 NormalParamExtractor 模組，它將把這個輸出分成兩個區塊，大小為 [1] 的平均值和標準差；
一個 ProbabilisticActor，它將以 in_keys 的形式讀取這些參數，使用它們創建一個分布，並使用樣本和對數機率填充我們的 tensordict。

from tensordict.nn.distributions import NormalParamExtractor
from torch.distributions import Normal
from torchrl.modules import ProbabilisticActor

backbone = MLP(in_features=3, out_features=2)
extractor = NormalParamExtractor()
module = torch.nn.Sequential(backbone, extractor)
td_module = TensorDictModule(module, in_keys=["observation"], out_keys=["loc", "scale"])
policy = ProbabilisticActor(
    td_module,
    in_keys=["loc", "scale"],
    out_keys=["action"],
    distribution_class=Normal,
    return_log_prob=True,
)

rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        loc: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

關於這個 rollout 有幾點需要注意

由於我們在 actor 的構建過程中要求了它，因此當時給定分布的動作的對數機率也會被寫入。這對於像 PPO 這樣的算法是必需的。
分布的參數也會在輸出 tensordict 中以 "loc" 和 "scale" 條目的形式返回。

您可以控制動作的抽樣，以使用期望值或分布的其他屬性，而不是在您的應用程式需要時使用隨機樣本。這可以通過 set_exploration_type() 函數來控制

from torchrl.envs.utils import ExplorationType, set_exploration_type

with set_exploration_type(ExplorationType.DETERMINISTIC):
    # takes the mean as action
    rollout = env.rollout(max_steps=10, policy=policy)
with set_exploration_type(ExplorationType.RANDOM):
    # Samples actions according to the dist
    rollout = env.rollout(max_steps=10, policy=policy)

查看文檔字符串中的 default_interaction_type 關鍵字參數以了解更多信息。

探索¶

像這樣的隨機策略在某種程度上自然地權衡了探索和利用，但確定性策略則不會。幸運的是，TorchRL 也可以通過其探索模組來緩解這一點。我們將以 EGreedyModule 探索模組為例（另請查看 AdditiveGaussianModule 和 OrnsteinUhlenbeckProcessModule）。為了看到這個模組的實際效果，讓我們回到確定性策略

from tensordict.nn import TensorDictSequential
from torchrl.modules import EGreedyModule

policy = Actor(MLP(3, 1, num_cells=[32, 64]))

我們的 \(\epsilon\)-greedy 探索模組通常會使用一些退火幀和 \(\epsilon\) 參數的初始值進行自定義。 \(\epsilon = 1\) 的值意味著採取的每個動作都是隨機的，而 \(\epsilon=0\) 意味著根本沒有探索。為了退火（即降低）探索因子，需要調用 step()（有關示例，請參閱最後的教學）。

exploration_module = EGreedyModule(
    spec=env.action_spec, annealing_num_steps=1000, eps_init=0.5
)

為了構建我們的探索性策略，我們只需要將確定性策略模組與 TensorDictSequential 模組中的探索模組連接起來（這類似於 tensordict 領域中的 Sequential）。

exploration_policy = TensorDictSequential(policy, exploration_module)

with set_exploration_type(ExplorationType.DETERMINISTIC):
    # Turns off exploration
    rollout = env.rollout(max_steps=10, policy=exploration_policy)
with set_exploration_type(ExplorationType.RANDOM):
    # Turns on exploration
    rollout = env.rollout(max_steps=10, policy=exploration_policy)

由於它必須能夠在動作空間中採樣隨機動作，因此 EGreedyModule 必須配備環境中的 action_space，才能知道使用什麼策略來隨機採樣動作。

Q-Value actors¶

在某些情況下，策略不是一個獨立的模組，而是構建在另一個模組之上。**Q 值（Q-Value）行動者**就是這種情況。簡而言之，這些行動者需要一個動作價值（action value）的估計值（通常是離散的），並且會貪婪地選擇具有最高價值的動作。在某些情況下（有限的離散動作空間和有限的離散狀態空間），可以直接儲存一個狀態-動作對的 2D 表格，並選擇具有最高價值的動作。DQN 帶來的創新是通過利用神經網路對 Q(s, a) 價值映射進行編碼，從而將此方法擴展到連續狀態空間。讓我們考慮另一個具有離散動作空間的環境，以便更清楚地理解。

env = GymEnv("CartPole-v1")
print(env.action_spec)

OneHot(
    shape=torch.Size([2]),
    space=CategoricalBox(n=2),
    device=cpu,
    dtype=torch.int64,
    domain=discrete)

我們構建一個價值網路，當它從環境中讀取狀態時，會產生每個動作的值。

num_actions = 2
value_net = TensorDictModule(
    MLP(out_features=num_actions, num_cells=[32, 32]),
    in_keys=["observation"],
    out_keys=["action_value"],
)

我們可以通過在價值網路之後添加一個 QValueModule 來輕鬆構建我們的 Q 值行動者。

from torchrl.modules import QValueModule

policy = TensorDictSequential(
    value_net,  # writes action values in our tensordict
    QValueModule(spec=env.action_spec),  # Reads the "action_value" entry by default
)

讓我們看看！我們運行策略幾個步驟，並查看輸出。我們應該在獲得的 rollout 中找到一個 "action_value" 以及一個 "chosen_action_value" 條目。

rollout = env.rollout(max_steps=3, policy=policy)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.int64, is_shared=False),
        action_value: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        chosen_action_value: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

因為它依賴於 argmax 運算符，所以此策略是確定性的。在資料收集期間，我們需要探索環境。為此，我們再次使用 EGreedyModule。

policy_explore = TensorDictSequential(policy, EGreedyModule(env.action_spec))

with set_exploration_type(ExplorationType.RANDOM):
    rollout_explore = env.rollout(max_steps=3, policy=policy_explore)

這就是我們關於使用 TorchRL 構建策略的簡短教學！

您可以使用該函式庫做更多的事情。一個好的起點是查看模組的 API 參考。

下一步

了解如何在動作是複合動作時（例如，環境需要離散動作和連續動作），使用具有 CompositeDistribution 的複合分佈；
看看如何在策略中使用 RNN（一個教學）；
將此與 Decision Transformers 範例中轉換器的使用情況進行比較（請參閱 GitHub 上的 example 目錄）。

腳本的總運行時間：（0 分鐘 46.376 秒）

預估記憶體使用量： 320 MB

由 Sphinx-Gallery 生成的圖庫

TorchRL 模組入門¶

TensorDictModules¶

特殊化的封裝器¶

網路¶

機率策略¶

探索¶

Q-Value actors¶

文件

教學課程

資源