目錄
  • 文件 >
  • 鐘擺：使用 TorchRL 編寫您的環境和變換
快捷方式
tutorials/pendulum

注意

跳至末尾 下載完整的示例程式碼。

Pendulum：使用 TorchRL 編寫你的環境和變換

作者：Vincent Moens

建立環境（模擬器或物理控制系統的介面）是強化學習和控制工程中不可或缺的一部分。

TorchRL 提供了一套工具，可在多種上下文中完成此操作。本教程將演示如何使用 PyTorch 和 TorchRL 從頭開始編寫一個 Pendulum 模擬器。該教程深受 OpenAI-Gym/Farama-Gymnasium 控制庫 的 Pendulum-v1 實現的啟發。

Pendulum

簡單擺錘

主要學習內容

  • 如何在 TorchRL 中設計環境：- 編寫規範（輸入、觀察和獎勵）；- 實現行為：設定種子、重置和步進。

  • 轉換環境的輸入和輸出，並編寫自己的變換；

  • 如何使用 TensorDict 在 codebase 中傳遞任意資料結構。

    在此過程中，我們將涉及 TorchRL 的三個關鍵元件：

為了展示 TorchRL 環境的潛力，我們將設計一個無狀態環境。有狀態環境會跟蹤遇到的最新物理狀態，並依賴此狀態來模擬狀態到狀態的轉換，而無狀態環境則期望在每一步都接收當前狀態以及執行的動作。TorchRL 支援這兩種型別的環境，但無狀態環境更通用，因此涵蓋了 TorchRL 環境 API 的更廣泛功能。

對無狀態環境進行建模可使使用者完全控制模擬器的輸入和輸出：可以在任何階段重置實驗或從外部主動修改動態。但是，這假定我們對任務有一定的控制，而事實並非總是如此：解決我們無法控制當前狀態的問題更具挑戰性，但具有更廣泛的應用集。

無狀態環境的另一個優點是它們可以實現轉換模擬的批處理執行。如果後端和實現允許，代數運算可以無縫地在標量、向量或張量上執行。本教程提供了此類示例。

本教程將按以下結構進行：

  • 我們將首先熟悉環境屬性：其形狀（batch_size）、其方法（主要是 step()、reset() 和 set_seed()）以及最終的規範。

  • 在編寫完模擬器後，我們將演示如何在訓練過程中使用變換。

  • 我們將探索 TorchRL API 帶來的一些新途徑，包括：轉換輸入的可能性、模擬的向量化執行以及透過模擬圖進行反向傳播的可能性。

  • 最後，我們將訓練一個簡單的策略來解決我們實現的系統。

此類中包含此環境的內建版本：~torchrl.envs.PendulumEnv。

from collections import defaultdict

import numpy as np
import torch
import tqdm
from tensordict import TensorDict, TensorDictBase
from tensordict.nn import TensorDictModule
from torch import nn

from torchrl.data import Bounded, Composite, Unbounded
from torchrl.envs import (
    CatTensors,
    EnvBase,
    Transform,
    TransformedEnv,
    UnsqueezeTransform,
)
from torchrl.envs.transforms.transforms import _apply_to_composite
from torchrl.envs.utils import check_env_specs, step_mdp

DEFAULT_X = np.pi
DEFAULT_Y = 1.0

在設計新的環境類時，有四件事你需要注意：

  • EnvBase._reset()，它負責在（可能隨機的）初始狀態下重置模擬器；

  • EnvBase._step()，它負責狀態轉換的動態；

  • EnvBase._set_seed()，它實現了播種機制；

  • 環境規範。

讓我們先描述一下手頭的問題：我們想模擬一個簡單的擺錘，我們可以控制其固定點的扭矩。我們的目標是將擺錘放置在向上位置（約定為角度為 0），並使其在該位置靜止。為了設計我們的動態系統，我們需要定義兩個方程：根據動作（施加的扭矩）的運動方程和構成我們目標函式的獎勵方程。

對於運動方程，我們將根據以下公式更新角速度：

\[\dot{\theta}_{t+1} = \dot{\theta}_t + (3 * g / (2 * L) * \sin(\theta_t) + 3 / (m * L^2) * u) * dt\]

其中 \(\dot{\theta}\) 是角速度（以弧度/秒為單位），\(g\) 是重力加速度，\(L\) 是擺錘長度，\(m\) 是其質量，\(\theta\) 是其角度位置，\(u\) 是扭矩。然後根據以下公式更新角度位置：

\[\theta_{t+1} = \theta_{t} + \dot{\theta}_{t+1} dt\]

我們將獎勵定義為：

\[r = -(\theta^2 + 0.1 * \dot{\theta}^2 + 0.001 * u^2)\]

當角度接近 0（擺錘處於向上位置）、角速度接近 0（無運動）且扭矩也為 0 時，這將最大化。

編碼動作的效果：_step()

步進方法是首先要考慮的，因為它將編碼我們感興趣的模擬。在 TorchRL 中，EnvBase 類有一個 EnvBase.step() 方法，該方法接收一個 tensordict.TensorDict 例項，其中包含一個 "action" 條目，指示要執行的動作。

為了方便讀寫 tensordict，並確保鍵與庫的期望一致，模擬部分已委託給私有的抽象方法 _step()，該方法從 tensordict 讀取輸入資料，並寫入一個新的 tensordict 包含輸出資料。

_step() 方法應執行以下操作：

  1. 讀取輸入鍵（如 "action"）並基於這些鍵執行模擬；

  2. 檢索觀察值、完成狀態和獎勵；

  3. 將觀察值、獎勵和完成狀態集寫入新的 TensorDict 中的相應條目。

接下來，step() 方法將 step() 的輸出合併到輸入 tensordict 中，以強制執行輸入/輸出一致性。

通常，對於有狀態環境，它看起來像這樣：

>>> policy(env.reset())
>>> print(tensordict)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
>>> env.step(tensordict)
>>> print(tensordict)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

請注意，根 tensordict 沒有改變，唯一的修改是出現了一個新的 "next" 條目，其中包含新資訊。

在 Pendulum 示例中，我們的 _step() 方法將從輸入 tensordict 中讀取相關條目，並計算在施加了 "action" 鍵編碼的力後，擺錘的位置和速度。我們計算新的擺錘角度 "new_th"，它等於先前位置 "th" 加上新的速度 "new_thdot" 在時間間隔 dt 上的值。

由於我們的目標是使擺錘向上並保持靜止，因此我們的 cost（負獎勵）函式對於接近目標的位置和低速度來說較低。確實，我們想阻止那些偏離“向上”位置太遠以及/或速度偏離 0 太遠的情況。

在我們的示例中，EnvBase._step() 被編碼為一個靜態方法，因為我們的環境是無狀態的。在有狀態環境中，需要 self 引數，因為狀態需要從環境中讀取。

def _step(tensordict):
    th, thdot = tensordict["th"], tensordict["thdot"]  # th := theta

    g_force = tensordict["params", "g"]
    mass = tensordict["params", "m"]
    length = tensordict["params", "l"]
    dt = tensordict["params", "dt"]
    u = tensordict["action"].squeeze(-1)
    u = u.clamp(-tensordict["params", "max_torque"], tensordict["params", "max_torque"])
    costs = angle_normalize(th) ** 2 + 0.1 * thdot**2 + 0.001 * (u**2)

    new_thdot = (
        thdot
        + (3 * g_force / (2 * length) * th.sin() + 3.0 / (mass * length**2) * u) * dt
    )
    new_thdot = new_thdot.clamp(
        -tensordict["params", "max_speed"], tensordict["params", "max_speed"]
    )
    new_th = th + new_thdot * dt
    reward = -costs.view(*tensordict.shape, 1)
    done = torch.zeros_like(reward, dtype=torch.bool)
    out = TensorDict(
        {
            "th": new_th,
            "thdot": new_thdot,
            "params": tensordict["params"],
            "reward": reward,
            "done": done,
        },
        tensordict.shape,
    )
    return out


def angle_normalize(x):
    return ((x + torch.pi) % (2 * torch.pi)) - torch.pi

重置模擬器：_reset()

我們需要關注的第二個方法是 _reset() 方法。與 _step() 類似，它應該在輸出的 tensordict 中寫入觀察條目以及可能的完成狀態（如果省略完成狀態，則父方法 reset() 會將其填充為 False）。在某些情況下，_reset 方法需要接收呼叫它的函式的命令（例如，在多代理設定中，我們可能需要指示哪些代理需要被重置）。這就是為什麼 _reset() 方法也期望一個 tensordict 作為輸入，儘管它可以是空的或 None。

父方法 EnvBase.reset() 進行一些簡單的檢查，就像 EnvBase.step() 所做的那樣，例如確保 "done" 狀態在輸出 tensordict 中返回，並且形狀與規範期望的相匹配。

對我們來說，唯一需要考慮的重要事情是 EnvBase._reset() 是否包含所有預期的觀察值。再次強調，由於我們正在使用無狀態環境，我們將擺錘的配置放在一個名為 "params" 的巢狀 tensordict 中。

在此示例中，我們不傳遞完成狀態，因為這對於 _reset() 不是必需的，並且我們的環境是非終止的，所以我們總是期望它為 False。

def _reset(self, tensordict):
    if tensordict is None or tensordict.is_empty():
        # if no ``tensordict`` is passed, we generate a single set of hyperparameters
        # Otherwise, we assume that the input ``tensordict`` contains all the relevant
        # parameters to get started.
        tensordict = self.gen_params(batch_size=self.batch_size)

    high_th = torch.tensor(DEFAULT_X, device=self.device)
    high_thdot = torch.tensor(DEFAULT_Y, device=self.device)
    low_th = -high_th
    low_thdot = -high_thdot

    # for non batch-locked environments, the input ``tensordict`` shape dictates the number
    # of simulators run simultaneously. In other contexts, the initial
    # random state's shape will depend upon the environment batch-size instead.
    th = (
        torch.rand(tensordict.shape, generator=self.rng, device=self.device)
        * (high_th - low_th)
        + low_th
    )
    thdot = (
        torch.rand(tensordict.shape, generator=self.rng, device=self.device)
        * (high_thdot - low_thdot)
        + low_thdot
    )
    out = TensorDict(
        {
            "th": th,
            "thdot": thdot,
            "params": tensordict["params"],
        },
        batch_size=tensordict.shape,
    )
    return out

環境元資料：env.*_spec

規範定義了環境的輸入和輸出域。準確定義執行時接收的張量很重要，因為它們通常用於在多程序和分散式環境中傳遞環境資訊。它們還可以用於例項化惰性定義的神經網路和測試指令碼，而無需實際查詢環境（例如，對於真實世界的物理系統，這可能會很昂貴）。

有四種規範是我們必須在環境中編寫的：

  • EnvBase.observation_spec：這將是一個 CompositeSpec 例項，其中每個鍵都是一個觀察值（CompositeSpec 可以看作是規範的字典）。

  • EnvBase.action_spec：它可以是任何型別的規範，但要求它對應於輸入 tensordict 中的 "action" 條目；

  • EnvBase.reward_spec：提供有關獎勵空間的資訊；

  • EnvBase.done_spec：提供有關完成標誌空間的資訊。

TorchRL 規範分為兩個通用容器：input_spec，其中包含 step 函式讀取的資訊的規範（分為 action_spec，包含動作，以及 state_spec，包含其餘所有資訊），以及 output_spec，它編碼 step 輸出的規範（observation_spec、reward_spec 和 done_spec）。通常，您不應直接與 output_spec 和 input_spec 互動，而應僅與其內容互動：observation_spec、reward_spec、done_spec、action_spec 和 state_spec。原因是規範在 output_spec 和 input_spec 中以非平凡的方式組織，並且這些都不應被直接修改。

換句話說，observation_spec 和相關屬性是輸出和輸入規範容器內容的便捷快捷方式。

TorchRL 提供了多個 TensorSpec 子類 來編碼環境的輸入和輸出特徵。

規範形狀

環境規範的前導維度必須與環境的批次大小匹配。這樣做是為了強制使環境的每個元件（包括其變換）都能準確地表示預期的輸入和輸出形狀。在有狀態設定中，應準確編寫此項。

對於非批次鎖定環境，例如我們示例中的環境（見下文），這無關緊要，因為環境的批次大小很可能為空。

def _make_spec(self, td_params):
    # Under the hood, this will populate self.output_spec["observation"]
    self.observation_spec = Composite(
        th=Bounded(
            low=-torch.pi,
            high=torch.pi,
            shape=(),
            dtype=torch.float32,
        ),
        thdot=Bounded(
            low=-td_params["params", "max_speed"],
            high=td_params["params", "max_speed"],
            shape=(),
            dtype=torch.float32,
        ),
        # we need to add the ``params`` to the observation specs, as we want
        # to pass it at each step during a rollout
        params=make_composite_from_td(td_params["params"]),
        shape=(),
    )
    # since the environment is stateless, we expect the previous output as input.
    # For this, ``EnvBase`` expects some state_spec to be available
    self.state_spec = self.observation_spec.clone()
    # action-spec will be automatically wrapped in input_spec when
    # `self.action_spec = spec` will be called supported
    self.action_spec = Bounded(
        low=-td_params["params", "max_torque"],
        high=td_params["params", "max_torque"],
        shape=(1,),
        dtype=torch.float32,
    )
    self.reward_spec = Unbounded(shape=(*td_params.shape, 1))


def make_composite_from_td(td):
    # custom function to convert a ``tensordict`` in a similar spec structure
    # of unbounded values.
    composite = Composite(
        {
            key: make_composite_from_td(tensor)
            if isinstance(tensor, TensorDictBase)
            else Unbounded(dtype=tensor.dtype, device=tensor.device, shape=tensor.shape)
            for key, tensor in td.items()
        },
        shape=td.shape,
    )
    return composite

可復現的實驗：播種

在初始化實驗時，為環境設定種子是一項常見操作。EnvBase._set_seed() 的唯一目標是設定包含的模擬器的種子。如果可能，此操作不應呼叫 reset() 或與環境執行進行互動。父方法 EnvBase.set_seed() 包含一個機制，允許使用不同的偽隨機且可復現的種子為多個環境設定種子。

def _set_seed(self, seed: int | None) -> None:
    rng = torch.manual_seed(seed)
    self.rng = rng

將各部分整合在一起：EnvBase 類

我們終於可以將各部分整合在一起，設計我們的環境類。規範初始化需要在環境構建過程中執行，因此我們必須在 PendulumEnv.__init__() 中呼叫 _make_spec() 方法。

我們添加了一個靜態方法 PendulumEnv.gen_params()，該方法可確定性地生成一組將在執行過程中使用的超引數。

def gen_params(g=10.0, batch_size=None) -> TensorDictBase:
    """Returns a ``tensordict`` containing the physical parameters such as gravitational force and torque or speed limits."""
    if batch_size is None:
        batch_size = []
    td = TensorDict(
        {
            "params": TensorDict(
                {
                    "max_speed": 8,
                    "max_torque": 2.0,
                    "dt": 0.05,
                    "g": g,
                    "m": 1.0,
                    "l": 1.0,
                },
                [],
            )
        },
        [],
    )
    if batch_size:
        td = td.expand(batch_size).contiguous()
    return td

我們將環境定義為非batch_locked 的，將 homonymous 屬性設定為 False。這意味著我們不會強制輸入的 tensordict 的 batch-size 與環境的相匹配。

以下程式碼將把我們上面編寫的部分整合在一起。

class PendulumEnv(EnvBase):
    metadata = {
        "render_modes": ["human", "rgb_array"],
        "render_fps": 30,
    }
    batch_locked = False

    def __init__(self, td_params=None, seed=None, device="cpu"):
        if td_params is None:
            td_params = self.gen_params()

        super().__init__(device=device, batch_size=[])
        self._make_spec(td_params)
        if seed is None:
            seed = torch.empty((), dtype=torch.int64).random_().item()
        self.set_seed(seed)

    # Helpers: _make_step and gen_params
    gen_params = staticmethod(gen_params)
    _make_spec = _make_spec

    # Mandatory methods: _step, _reset and _set_seed
    _reset = _reset
    _step = staticmethod(_step)
    _set_seed = _set_seed

測試我們的環境

TorchRL 提供了一個簡單的函式 check_env_specs() 來檢查（轉換後的）環境的輸入/輸出結構是否與其規範匹配。讓我們嘗試一下：

env = PendulumEnv()
check_env_specs(env)

我們可以檢視我們的規範，以獲得環境簽名的視覺化表示：

print("observation_spec:", env.observation_spec)
print("state_spec:", env.state_spec)
print("reward_spec:", env.reward_spec)

我們還可以執行一些命令來檢查輸出結構是否與預期相符。

td = env.reset()
print("reset tensordict", td)

我們可以執行 env.rand_step() 來從 action_spec 域中隨機生成一個動作。必須傳入一個包含超引數和當前狀態的 tensordict，因為我們的環境是無狀態的。在有狀態的上下文中，env.rand_step() 也能完美工作。

td = env.rand_step(td)
print("random step tensordict", td)

轉換環境

為無狀態模擬器編寫環境變換比為有狀態模擬器稍微複雜一些：轉換一個需要在下一個迭代中讀取的輸出條目，需要在下一個步驟呼叫 meth.step() 之前應用逆變換。這是展示 TorchRL 變換所有功能的理想場景！

例如，在以下變換後的環境中，我們 unsqueeze 條目 ["th", "thdot"]，以便能夠沿最後一個維度堆疊它們。我們還將它們作為 in_keys_inv 傳遞，以便在它們作為下一個迭代的輸入時將它們擠壓回原始形狀。

env = TransformedEnv(
    env,
    # ``Unsqueeze`` the observations that we will concatenate
    UnsqueezeTransform(
        dim=-1,
        in_keys=["th", "thdot"],
        in_keys_inv=["th", "thdot"],
    ),
)

編寫自定義變換

TorchRL 的變換可能無法涵蓋您在環境執行後想要執行的所有操作。編寫變換並不需要太多精力。與環境設計一樣，編寫變換有兩個步驟：

  • 正確處理動態（正向和反向）；

  • 調整環境規範。

變換可以在兩種設定中使用：本身可以作為 Module 使用。它也可以作為 TransformedEnv 的一個附加部分使用。該類的結構允許自定義不同上下文下的行為。

一個 Transform 的骨架可以總結如下：

class Transform(nn.Module):
    def forward(self, tensordict):
        ...
    def _apply_transform(self, tensordict):
        ...
    def _step(self, tensordict):
        ...
    def _call(self, tensordict):
        ...
    def inv(self, tensordict):
        ...
    def _inv_apply_transform(self, tensordict):
        ...

有三個入口點（forward()、_step() 和 inv()），它們都接收 tensordict.TensorDict 例項。前兩個最終會透過 in_keys 指定的鍵，並對每個鍵呼叫 _apply_transform()。如果提供了 Transform.out_keys（否則 in_keys 將被轉換後的值更新），結果將寫入由 Transform.out_keys 指向的條目中。如果需要執行逆變換，類似的邏輯將執行，但使用 Transform.inv() 和 Transform._inv_apply_transform() 方法，並跨 in_keys_inv 和 out_keys_inv 鍵列表進行。下圖總結了環境和回放緩衝區的資料流。

變換 API

在某些情況下，變換不會以單一方式作用於鍵的子集，而是會執行一些操作來作用於父環境或使用整個輸入 tensordict。在這些情況下，應重寫 _call() 和 forward() 方法，並可以跳過 _apply_transform() 方法。

讓我們編寫新的變換，計算位置角度的 sine 和 cosine 值，因為與原始角度值相比，這些值對於學習策略更有用：

class SinTransform(Transform):
    def _apply_transform(self, obs: torch.Tensor) -> None:
        return obs.sin()

    # The transform must also modify the data at reset time
    def _reset(
        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
    ) -> TensorDictBase:
        return self._call(tensordict_reset)

    # _apply_to_composite will execute the observation spec transform across all
    # in_keys/out_keys pairs and write the result in the observation_spec which
    # is of type ``Composite``
    @_apply_to_composite
    def transform_observation_spec(self, observation_spec):
        return Bounded(
            low=-1,
            high=1,
            shape=observation_spec.shape,
            dtype=observation_spec.dtype,
            device=observation_spec.device,
        )


class CosTransform(Transform):
    def _apply_transform(self, obs: torch.Tensor) -> None:
        return obs.cos()

    # The transform must also modify the data at reset time
    def _reset(
        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
    ) -> TensorDictBase:
        return self._call(tensordict_reset)

    # _apply_to_composite will execute the observation spec transform across all
    # in_keys/out_keys pairs and write the result in the observation_spec which
    # is of type ``Composite``
    @_apply_to_composite
    def transform_observation_spec(self, observation_spec):
        return Bounded(
            low=-1,
            high=1,
            shape=observation_spec.shape,
            dtype=observation_spec.dtype,
            device=observation_spec.device,
        )


t_sin = SinTransform(in_keys=["th"], out_keys=["sin"])
t_cos = CosTransform(in_keys=["th"], out_keys=["cos"])
env.append_transform(t_sin)
env.append_transform(t_cos)

將觀察值連線到一個“observation”條目。 del_keys=False 確保我們保留這些值供下次迭代使用。

cat_transform = CatTensors(
    in_keys=["sin", "cos", "thdot"], dim=-1, out_key="observation", del_keys=False
)
env.append_transform(cat_transform)

再次，讓我們檢查一下我們的環境規範是否與接收到的相符：

check_env_specs(env)

執行一次軌跡執行

執行一次軌跡執行是一系列簡單的步驟：

  • 重置環境

  • 當某個條件不滿足時

    • 給定一個策略計算一個動作

    • 給定此動作執行一個步進

    • 收集資料

    • 執行一個 MDP 步進

  • 收集資料並返回

這些操作已方便地封裝在 rollout() 方法中，我們將在下面提供其簡化版本。

def simple_rollout(steps=100):
    # preallocate:
    data = TensorDict(batch_size=[steps])
    # reset
    _data = env.reset()
    for i in range(steps):
        _data["action"] = env.action_spec.rand()
        _data = env.step(_data)
        data[i] = _data
        _data = step_mdp(_data, keep_other=True)
    return data


print("data from rollout:", simple_rollout(100))

批處理計算

我們教程中最後一個未探索的方面是 TorchRL 中批處理計算的能力。因為我們的環境不對輸入資料的形狀做出任何假設，所以我們可以無縫地在資料批次上執行它。更好的是：對於非批次鎖定環境（如我們的 Pendulum），我們可以動態更改批次大小而無需重新建立環境。為此，我們只需生成具有所需形狀的引數。

batch_size = 10  # number of environments to be executed in batch
td = env.reset(env.gen_params(batch_size=[batch_size]))
print("reset (batch size of 10)", td)
td = env.rand_step(td)
print("rand step (batch size of 10)", td)

使用資料批次執行軌跡執行需要我們在軌跡執行函式外部重置環境，因為我們需要動態定義批次大小，而 rollout() 不支援此功能。

rollout = env.rollout(
    3,
    auto_reset=False,  # we're executing the reset out of the ``rollout`` call
    tensordict=env.reset(env.gen_params(batch_size=[batch_size])),
)
print("rollout of len 3 (batch size of 10):", rollout)

訓練一個簡單的策略

在本例中，我們將使用獎勵作為可微分目標（例如負損失）來訓練一個簡單的策略。我們將利用我們的動態系統是完全可微分的事實，透過軌跡回報進行反向傳播，並直接調整我們策略的權重以最大化此值。當然，在許多情況下，我們所做的假設（例如可微分系統和對底層機制的完全訪問）並不成立。

儘管如此，這是一個非常簡單的示例，展示瞭如何使用 TorchRL 中的自定義環境編寫訓練迴圈。

讓我們先編寫策略網路：

torch.manual_seed(0)
env.set_seed(0)

net = nn.Sequential(
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(1),
)
policy = TensorDictModule(
    net,
    in_keys=["observation"],
    out_keys=["action"],
)

以及我們的最佳化器：

optim = torch.optim.Adam(policy.parameters(), lr=2e-3)

訓練迴圈

我們將依次執行以下操作：

  • 生成一個軌跡

  • 累加獎勵

  • 透過操作定義的圖進行反向傳播

  • 剪輯梯度範數並執行最佳化步驟

  • 重複

在訓練迴圈結束時，我們應該得到一個接近 0 的最終獎勵，這表明擺錘已經向上並且靜止，正如我們所期望的。

batch_size = 32
n_iter = 1000  # set to 20_000 for a proper training
pbar = tqdm.tqdm(range(n_iter // batch_size))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, n_iter)
logs = defaultdict(list)

for _ in pbar:
    init_td = env.reset(env.gen_params(batch_size=[batch_size]))
    rollout = env.rollout(100, policy, tensordict=init_td, auto_reset=False)
    traj_return = rollout["next", "reward"].mean()
    (-traj_return).backward()
    gn = torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
    optim.step()
    optim.zero_grad()
    pbar.set_description(
        f"reward: {traj_return: 4.4f}, "
        f"last reward: {rollout[..., -1]['next', 'reward'].mean(): 4.4f}, gradient norm: {gn: 4.4}"
    )
    logs["return"].append(traj_return.item())
    logs["last_reward"].append(rollout[..., -1]["next", "reward"].mean().item())
    scheduler.step()


def plot():
    import matplotlib
    from matplotlib import pyplot as plt

    is_ipython = "inline" in matplotlib.get_backend()
    if is_ipython:
        from IPython import display

    with plt.ion():
        plt.figure(figsize=(10, 5))
        plt.subplot(1, 2, 1)
        plt.plot(logs["return"])
        plt.title("returns")
        plt.xlabel("iteration")
        plt.subplot(1, 2, 2)
        plt.plot(logs["last_reward"])
        plt.title("last reward")
        plt.xlabel("iteration")
        if is_ipython:
            display.display(plt.gcf())
            display.clear_output(wait=True)
        plt.show()


plot()

結論

在本教程中，我們學會了如何從頭開始編寫無狀態環境。我們探討了以下主題：

  • 編寫環境時需要注意的四個基本元件（step、reset、播種和構建規範）。我們看到了這些方法和類如何與 TensorDict 類進行互動；

  • 如何使用 check_env_specs() 測試環境是否已正確編寫；

  • 如何在無狀態環境的上下文中附加變換以及如何編寫自定義變換；

  • 如何在一個完全可微分的模擬器上訓練策略。

由 Sphinx-Gallery 生成的畫廊

文件

訪問全面的 PyTorch 開發者文件

檢視文件

教程

為初學者和高階開發者提供深入的教程

檢視教程

資源

查詢開發資源並讓您的問題得到解答

檢視資源

為了分析流量並最佳化您的體驗，我們在此網站上使用 Cookie。點選或瀏覽即表示您同意我們使用 Cookie。作為本網站的當前維護者，Facebook 的 Cookie 政策適用。瞭解更多資訊，包括可用的控制元件： Cookie 政策。