TorchRL 環境

作者：Vincent Moens

環境在強化學習（RL）設定中起著至關重要的作用，通常與監督和無監督設定中的資料集有些相似。RL 社群已經非常熟悉 OpenAI gym API，它提供了一種靈活的方式來構建、初始化和互動環境。然而，還存在許多其他庫，並且與它們互動的方式可能與 gym 所期望的非常不同。

讓我們從描述 TorchRL 如何與 gym 互動開始，這將作為介紹其他框架的引子。

Gym 環境

要執行本教程的這部分，您需要安裝一個最新版本的 gym 庫以及 atari suite。您可以透過安裝以下包來完成安裝：

為了統一所有框架，TorchRL 環境在 `__init__` 方法中構建，並使用一個私有方法 `_build_env`，該方法會將引數和關鍵字引數傳遞給根庫構建器。

使用 gym，這意味著構建一個環境就像：

import torch
from matplotlib import pyplot as plt
from tensordict import TensorDict

from torchrl.envs.libs.gym import GymEnv

env = GymEnv("Pendulum-v1")

可以透過此命令訪問可用環境的列表：

list(GymEnv.available_envs)[:10]

環境規範

與其他框架一樣，TorchRL 環境具有指示觀測、動作、完成和獎勵的空間的屬性。因為經常會檢索到不止一個觀測，所以我們期望觀測規範的型別為 `CompositeSpec`。獎勵和動作沒有這個限制。

print("Env observation_spec: \n", env.observation_spec)
print("Env action_spec: \n", env.action_spec)
print("Env reward_spec: \n", env.reward_spec)

這些規範附帶一系列有用的工具：可以斷言一個樣本是否在定義的空間內。我們還可以使用一些啟發式方法將樣本投影到空間內（如果它超出空間），並在該空間內生成隨機數（可能是均勻分佈的）。

action = torch.ones(1) * 3
print("action is in bounds?\n", bool(env.action_spec.is_in(action)))
print("projected action: \n", env.action_spec.project(action))

print("random action: \n", env.action_spec.rand())

在這些規範中，`done_spec` 值得特別關注。在 TorchRL 中，所有環境都會寫出至少兩種型別的軌跡結束訊號：`"terminated"`（表示馬爾可夫決策過程已達到最終狀態——__episode__ 已結束）和`"done"`，表示這是 __trajectory__ 的最後一步（但不一定是任務的結束）。通常，當 `terminal` 為 `False` 時為 `True` 的 `"done"` 條目是由 `"truncated"` 訊號引起的。Gym 環境考慮了這三個訊號。

print(env.done_spec)

環境還帶有一個型別為 `CompositeSpec` 的 `env.state_spec` 屬性，其中包含所有作為環境輸入的但非動作的規範。對於有狀態環境（例如 gym），這在大多數情況下是空的。對於無狀態環境（例如 Brax），這也應該包括前一個狀態的表示，或者環境的任何其他輸入（包括重置時的輸入）。

播種、重置和步進

環境的基本操作是（1）`set_seed`、（2）`reset` 和（3）`step`。

讓我們看看這些方法在 TorchRL 中是如何工作的。

torch.manual_seed(0)  # make sure that all torch code is also reproductible
env.set_seed(0)
reset_data = env.reset()
print("reset data", reset_data)

我們現在可以在環境中執行一步。由於我們沒有策略，我們可以生成一個隨機動作。

policy = TensorDictModule(env.action_spec.rand, in_keys=[], out_keys=["action"])


policy(reset_data)
tensordict_out = env.step(reset_data)

預設情況下，`step` 返回的 tensordict 與輸入相同……

assert tensordict_out is reset_data

……但帶有新鍵。

tensordict_out

我們剛才所做的（使用 `action_spec.rand()` 進行隨機步進）也可以透過簡單的快捷方式完成。

env.rand_step()

新鍵 `("next", "observation")`（如 `"next"` tensordict 下的所有鍵）在 TorchRL 中有一個特殊的作用：它們表示它們出現在沒有字首的同名鍵之後。

我們提供了一個函式 `step_mdp`，它在 tensordict 中執行一步：它返回一個更新後的新 tensordict，使得 *t < -t’*。

from torchrl.envs.utils import step_mdp

tensordict_out.set("some other key", torch.randn(1))
tensordict_tprime = step_mdp(tensordict_out)

print(tensordict_tprime)
print(
    (
        tensordict_tprime.get("observation")
        == tensordict_out.get(("next", "observation"))
    ).all()
)

我們可以注意到 `step_mdp` 移除了所有與時間相關的鍵值對，但保留了 `"some other key"`。此外，新觀測與之前的觀測匹配。

最後，請注意 `env.reset` 方法也接受一個 tensordict 進行更新。

data = TensorDict()
assert env.reset(data) is data
data

Rollouts

TorchRL 提供的通用環境類允許您輕鬆執行給定步數的 rollout。

tensordict_rollout = env.rollout(max_steps=20, policy=policy)
print(tensordict_rollout)

生成的 tensordict 的 `batch_size` 為 `[20]`，這是軌跡的長度。我們可以檢查觀測值是否與其下一個值匹配。

(
    tensordict_rollout.get("observation")[1:]
    == tensordict_rollout.get(("next", "observation"))[:-1]
).all()

`frame_skip`

在某些情況下，使用 `frame_skip` 引數將相同的動作用於幾個連續的幀很有用。

生成的 tensordict 將只包含序列中最後觀察到的幀，但獎勵將是多幀的總和。

如果環境在此過程中達到完成狀態，它將停止並返回截斷鏈的結果。

env = GymEnv("Pendulum-v1", frame_skip=4)
env.reset()

渲染

渲染在許多 RL 設定中起著重要作用，這就是為什麼 TorchRL 的通用環境類提供了一個 `from_pixels` 關鍵字引數，允許使用者快速請求基於影像的環境。

env = GymEnv("Pendulum-v1", from_pixels=True)

data = env.reset()
env.close()

plt.imshow(data.get("pixels").numpy())

讓我們看看 tensordict 中包含什麼。

data

我們仍然有一個 `"state"`，它描述了在上一例中 `"observation"` 所描述的內容（命名差異來自於 gym 現在返回一個字典，並且 TorchRL 從字典中獲取名稱（如果存在），否則會將步驟輸出命名為 `"observation"`：簡而言之，這是由於 gym 環境的 step 方法返回的物件型別不一致）。

也可以透過僅請求畫素來丟棄此附加輸出。

env = GymEnv("Pendulum-v1", from_pixels=True, pixels_only=True)
env.reset()
env.close()

有些環境只有基於影像的格式。

env = GymEnv("ALE/Pong-v5")
print("from pixels: ", env.from_pixels)
print("data: ", env.reset())
env.close()

DeepMind 控制環境

要執行本教程的這部分，請確保您已安裝 dm_control。

$ pip install dm_control

我們還為 DM Control suite 提供了一個包裝器。同樣，構建環境也很容易：首先讓我們看看可以訪問哪些環境。`available_envs` 現在返回一個包含環境和可能任務的字典。

from matplotlib import pyplot as plt

from torchrl.envs.libs.dm_control import DMControlEnv

DMControlEnv.available_envs

env = DMControlEnv("acrobot", "swingup")
data = env.reset()
print("result of reset: ", data)
env.close()

當然，我們也可以使用基於畫素的環境。

env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
data = env.reset()
print("result of reset: ", data)
plt.imshow(data.get("pixels").numpy())
env.close()

轉換環境

在策略讀取環境輸出或將其儲存在緩衝區之前對其進行預處理是很常見的。

在許多情況下，RL 社群已經採用了以下型別的包裝方案：

$ env_transformed = wrapper1(wrapper2(env))

來轉換環境。這有很多優點：它使訪問環境規範變得顯而易見（外部包裝器是外部世界的真實來源），並且它使得與向量化環境互動變得容易。然而，它也使得訪問內部環境變得困難：假設我們想從鏈中刪除一個包裝器（例如 `wrapper2`），這個操作需要我們收集

$ env0 = env.env.env

$ env_transformed_bis = wrapper1(env0)

TorchRL 採取使用轉換序列的立場，就像在其他 PyTorch 領域庫（例如 `torchvision`）中那樣。這種方法也類似於 `torch.distribution` 中轉換分佈的方式，其中 `TransformedDistribution` 物件圍繞 `base_dist` 分佈和（一系列）`transforms` 構建。

from torchrl.envs.transforms import ToTensorImage, TransformedEnv

# ToTensorImage transforms a numpy-like image into a tensor one,
env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
print("reset before transform: ", env.reset())

env = TransformedEnv(env, ToTensorImage())
print("reset after transform: ", env.reset())
env.close()

要組合轉換，只需使用 `Compose` 類。

from torchrl.envs.transforms import Compose, Resize

env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
env = TransformedEnv(env, Compose(ToTensorImage(), Resize(32, 32)))
env.reset()

也可以一次新增一個轉換。

from torchrl.envs.transforms import GrayScale

env.append_transform(GrayScale())
env.reset()

如預期，元資料也會更新。

print("original obs spec: ", env.base_env.observation_spec)
print("current obs spec: ", env.observation_spec)

如果需要，我們也可以連線張量。

from torchrl.envs.transforms import CatTensors

env = DMControlEnv("acrobot", "swingup")
print("keys before concat: ", env.reset())

env = TransformedEnv(
    env,
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation"),
)
print("keys after concat: ", env.reset())

此功能可以輕鬆修改應用於環境輸入和輸出的轉換集。事實上，轉換在執行步驟之前和之後都會執行：對於預步驟傳遞，`in_keys_inv` 鍵列表將被傳遞給 `_inv_apply_transform` 方法。這種轉換的一個例子是將浮點動作（來自神經網路的輸出）轉換為雙精度型別（被包裝環境需要）。在執行步驟之後，`_apply_transform` 方法將在 `in_keys` 鍵列表指定的鍵上執行。

環境轉換的另一個有趣之處在於，它們允許使用者檢索已包裝情況下的 `env.env` 的等效項，換句話說，父環境。父環境可以透過呼叫 `transform.parent` 來檢索：返回的環境將由一個 `TransformedEnvironment` 組成，其中包含當前轉換之前的所有轉換。例如，這可以在 `NoopResetEnv` 的情況下使用，它在重置時執行以下步驟：在執行該環境中的一定數量的隨機步驟之前重置父環境。

env = DMControlEnv("acrobot", "swingup")
env = TransformedEnv(env)
env.append_transform(
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation")
)
env.append_transform(GrayScale())

print("env: \n", env)
print("GrayScale transform parent env: \n", env.transform[1].parent)
print("CatTensors transform parent env: \n", env.transform[0].parent)

環境裝置

轉換可以在裝置上工作，當操作計算量適中或高時，這可以帶來顯著的加速。這些操作包括 `ToTensorImage`、`Resize`、`GrayScale` 等。

有人可能會合理地問這對被包裝的環境意味著什麼。對於常規環境來說，影響很小：操作仍將在應有的裝置上進行。TorchRL 中的環境裝置屬性指示傳入資料應該在哪個裝置上，以及傳出資料將在哪個裝置上。從該裝置到該裝置的轉換由 TorchRL 環境類負責。將資料儲存在 GPU 上的主要優點是（1）如上所述的轉換加速和（2）在多程序設定中共享資料。

from torchrl.envs.transforms import CatTensors, GrayScale, TransformedEnv

env = DMControlEnv("acrobot", "swingup")
env = TransformedEnv(env)
env.append_transform(
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation")
)

if torch.has_cuda and torch.cuda.device_count():
    env.to("cuda:0")
    env.reset()

並行執行環境

TorchRL 提供了並行執行環境的實用工具。預計各種環境讀取並返回形狀和資料型別相似的張量（但也可以設計掩碼函式，以便在這些張量形狀不同時使其可行）。建立這樣的環境非常容易。讓我們看看最簡單的情況。

from torchrl.envs import ParallelEnv


def env_make():
    return GymEnv("Pendulum-v1")


parallel_env = ParallelEnv(3, env_make)  # -> creates 3 envs in parallel
parallel_env = ParallelEnv(
    3, [env_make, env_make, env_make]
)  # similar to the previous command

`SerialEnv` 類與 `ParallelEnv` 類似，只是環境是按順序執行的。這主要用於除錯目的。

`ParallelEnv` 例項以惰性模式建立：環境將在被呼叫時才開始執行。這使得我們可以將 `ParallelEnv` 物件從一個程序移動到另一個程序，而無需過多擔心正在執行的程序。可以透過呼叫 `start`、`reset` 或簡單地呼叫 `step`（如果不需要先呼叫 `reset`）來啟動 `ParallelEnv`。

parallel_env.reset()

可以檢查並行環境是否具有正確的批處理大小。通常，`batch_size` 的第一部分表示批次，第二部分表示時間幀。讓我們用 `rollout` 方法來檢查這一點。

parallel_env.rollout(max_steps=20)

關閉並行環境

重要提示：在關閉程式之前，關閉並行環境很重要。通常，即使對於常規環境，用 `close` 呼叫結束函式也是一個好習慣。在某些情況下，如果未這樣做，TorchRL 會丟擲錯誤（通常會在程式結束時，當環境超出作用域時！）。

parallel_env.close()

播種

在播種並行環境時，我們面臨的困難是我們不希望為所有環境提供相同的種子。TorchRL 使用的啟發式方法是，我們給定輸入種子，以一種可以說是馬爾可夫的方式生成一個確定的種子鏈，使其可以從任何元素重建。所有 `set_seed` 方法都將返回下一個要使用的種子，這樣我們就可以輕鬆地根據最後一個種子保持鏈的連續性。當幾個收集器都包含 `ParallelEnv` 例項，並且我們希望每個子子環境都有不同的種子時，這很有用。

out_seed = parallel_env.set_seed(10)
print(out_seed)

del parallel_env

訪問環境屬性

有時會出現包裝環境具有感興趣的屬性的情況。首先，請注意，TorchRL 環境包裝器限制了訪問此屬性的工具。這是一個例子。

from time import sleep

from uuid import uuid1


def env_make():
    env = GymEnv("Pendulum-v1")
    env._env.foo = f"bar_{uuid1()}"
    env._env.get_something = lambda r: r + 1
    return env


env = env_make()

# Goes through env._env
env.foo

parallel_env = ParallelEnv(3, env_make)  # -> creates 3 envs in parallel

# env has not been started --> error:
try:
    parallel_env.foo
except RuntimeError:
    print("Aargh what did I do!")
    sleep(2)  # make sure we don't get ahead of ourselves

if parallel_env.is_closed:
    parallel_env.start()
foo_list = parallel_env.foo
foo_list  # needs to be instantiated, for instance using list

list(foo_list)

類似地，方法也可以被訪問。

something = parallel_env.get_something(0)
print(something)

parallel_env.close()
del parallel_env

並行環境的 kwargs

您可能希望為各種環境提供 kwargs。這可以在構造時或之後完成。

from torchrl.envs import ParallelEnv


def env_make(env_name):
    env = TransformedEnv(
        GymEnv(env_name, from_pixels=True, pixels_only=True),
        Compose(ToTensorImage(), Resize(64, 64)),
    )
    return env


parallel_env = ParallelEnv(
    2,
    [env_make, env_make],
    create_env_kwargs=[{"env_name": "ALE/AirRaid-v5"}, {"env_name": "ALE/Pong-v5"}],
)
data = parallel_env.reset()

plt.figure()
plt.subplot(121)
plt.imshow(data[0].get("pixels").permute(1, 2, 0).numpy())
plt.subplot(122)
plt.imshow(data[1].get("pixels").permute(1, 2, 0).numpy())
parallel_env.close()
del parallel_env

from matplotlib import pyplot as plt

轉換並行環境

有兩種等效的方式可以轉換並行環境：在每個程序中單獨轉換，或在主程序中轉換。甚至可以同時執行兩者。因此，可以仔細考慮轉換設計，以利用裝置功能（例如，CUDA 裝置上的轉換）並在可能的情況下在主程序上進行向量化操作。

from torchrl.envs import (
    Compose,
    GrayScale,
    ParallelEnv,
    Resize,
    ToTensorImage,
    TransformedEnv,
)


def env_make(env_name):
    env = TransformedEnv(
        GymEnv(env_name, from_pixels=True, pixels_only=True),
        Compose(ToTensorImage(), Resize(64, 64)),
    )  # transforms on remote processes
    return env


parallel_env = ParallelEnv(
    2,
    [env_make, env_make],
    create_env_kwargs=[{"env_name": "ALE/AirRaid-v5"}, {"env_name": "ALE/Pong-v5"}],
)
parallel_env = TransformedEnv(parallel_env, GrayScale())  # transforms on main process
data = parallel_env.reset()

print("grayscale data: ", data)
plt.figure()
plt.subplot(121)
plt.imshow(data[0].get("pixels").permute(1, 2, 0).numpy())
plt.subplot(122)
plt.imshow(data[1].get("pixels").permute(1, 2, 0).numpy())
parallel_env.close()
del parallel_env

VecNorm

在 RL 中，我們經常面臨在將資料輸入模型之前對其進行標準化的問​​題。有時，我們可以透過例如隨機策略（或演示）從環境中收集的資料來獲得標準統計資料的良好近似值。然而，有時建議“即時”標準化資料，逐步更新標準化常數以反映到目前為止已觀察到的情況。當期望標準化統計資料隨著任務效能的變化而變化，或者環境由於外部因素而演變時，這一點尤其有用。

警告：此功能應謹慎用於離策略學習，因為舊資料將因其使用以前有效的標準化統計資料進行標準化而被“棄用”。在同策略設定中，此功能也會使學習不穩定，並可能產生意外效果。因此，建議使用者謹慎使用此功能，並將其與使用固定版本標準化常數的資料標準化進行比較。

在常規設定中，使用 VecNorm 非常簡單。

from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.transforms import TransformedEnv, VecNorm

env = TransformedEnv(GymEnv("Pendulum-v1"), VecNorm())
data = env.rollout(max_steps=100)

print("mean: :", data.get("observation").mean(0))  # Approx 0
print("std: :", data.get("observation").std(0))  # Approx 1

在 **並行環境** 中，事情會稍微複雜一些，因為我們需要在程序之間共享執行統計資料。我們建立了一個 `EnvCreator` 類，該類負責檢視環境建立方法，檢索要在環境類中共享的 tensordicts，並指向每個程序到正確（公共、共享）資料一旦建立。

from torchrl.envs import EnvCreator, ParallelEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.transforms import TransformedEnv, VecNorm

make_env = EnvCreator(lambda: TransformedEnv(GymEnv("CartPole-v1"), VecNorm(decay=1.0)))
env = ParallelEnv(3, make_env)
print("env state dict:")
sd = TensorDict(make_env.state_dict())
print(sd)
# Zeroes all tensors
sd *= 0

data = env.rollout(max_steps=5)

print("data: ", data)
print("mean: :", data.get("observation").view(-1, 3).mean(0))  # Approx 0
print("std: :", data.get("observation").view(-1, 3).std(0))  # Approx 1

計數略高於步數（因為我們沒有使用任何衰減）。兩者之間的差異是由於 `ParallelEnv` 建立了一個虛擬環境來初始化共享的 `TensorDict`，該 `TensorDict` 用於從分派的環境中收集資料。這個微小的差異通常會在訓練過程中被吸收。

print(
    "update counts: ",
    make_env.state_dict()["_extra_state"]["observation_count"],
)

env.close()
del env