TorchRL 環境

作者: Vincent Moens

環境在強化學習 (RL) 設定中起著至關重要的作用，通常在某種程度上類似於監督學習和無監督學習中的資料集。RL 社群對 OpenAI gym API 相當熟悉，該 API 提供了一種靈活的方式來構建、初始化環境並與之互動。然而，還有許多其他庫，與它們互動的方式可能與使用 gym 的預期截然不同。

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

Gym 環境

要執行本教程的這一部分，您需要安裝最新版本的 gym 庫以及 atari 套件。您可以透過安裝以下包來安裝它

為了統一所有框架，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"（表示馬爾可夫決策過程已達到最終狀態 - __回合__結束）和 "done"，表示這是__軌跡__的最後一步（但不一定是任務的結束）。通常，當 "terminal" 為 False 時，"done" 條目為 True 是由 "truncated" 訊號引起的。Gym 環境考慮了這三個訊號

print(env.done_spec)

環境還包含一個 env.state_spec 屬性，型別為 CompositeSpec，其中包含作為環境輸入但不作為動作的所有規範。對於有狀態環境（例如 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

軌跡推演

TorchRL 提供的通用環境類允許您輕鬆執行給定步數的軌跡推演

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()

渲染

渲染在許多強化學習設定中扮演著重要角色，這就是 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 Control 環境

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

$ pip install dm_control

我們還為 DM Control 套件提供了包裝器。同樣，構建環境也很容易：首先讓我們看看可以訪問哪些環境。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 採取了使用 transforms 序列的方式，這與其他 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()

要組合 transforms，只需使用 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()

Transforms 也可以一次新增一個

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())

此功能使得修改應用於環境輸入和輸出的 transforms 集合變得容易。實際上，transforms 在執行步進之前和之後都會執行：對於步進前處理，in_keys_inv 鍵列表將傳遞給 _inv_apply_transform 方法。這類 transform 的一個例子是將浮點動作（神經網路的輸出）轉換為 double 資料型別（由包裝環境要求）。執行步進後，_apply_transform 方法將在 in_keys 鍵列表指示的鍵上執行。

環境 transforms 的另一個有趣特性是，它們允許使用者檢索包裝情況下的 env.env 的等價物，換句話說就是父環境。透過呼叫 transform.parent 可以檢索父環境：返回的環境將由一個 TransformedEnvironment 組成，包含直到（但不包括）當前 transform 的所有 transforms。這可以例如用於 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)

環境裝置

Transforms 可以在裝置上工作，這在操作計算需求中度或高度密集時可以帶來顯著的速度提升。這包括 ToTensorImage, Resize, GrayScale 等。

人們可能會合理地問，這對包裝環境端意味著什麼。對於常規環境來說影響很小：操作仍然會在它們應該發生的裝置上進行。torchrl 中的環境裝置屬性指示傳入資料應位於哪個裝置以及輸出資料將位於哪個裝置。從該裝置進行型別轉換是 torchrl 環境類的責任。將資料儲存在 GPU 上的主要優點是 (1) 如上所述的 transforms 加速，以及 (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 來啟動 ParallelEnv（如果不需要先呼叫 reset）。

parallel_env.reset()

可以檢查並行環境是否具有正確的 batch 大小。按照慣例，batch_size 的第一部分表示 batch，第二部分表示時間幀。讓我們用 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

轉換並行環境

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

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

在強化學習中，我們經常面臨在將資料輸入模型之前進行標準化的 問題。有時，我們可以從環境中收集的資料（例如使用隨機策略或演示資料）中獲得對標準化統計資料的良好近似。然而，建議對資料進行“動態”標準化，根據目前觀察到的情況逐步更新標準化常數。當期望標準化統計資料隨任務效能變化而改變，或當環境因外部因素而演變時，這尤其有用。

注意：在離策略學習中應謹慎使用此功能，因為舊資料由於使用了先前有效的標準化統計資料進行標準化，將會“過時”。在在策略設定中，此功能也會使學習不穩定並可能產生意外影響。因此建議使用者謹慎依賴此功能，並將其與給定固定版本的標準化常數進行資料標準化的方法進行比較。

在常規設定中，使用 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。這個小差異通常會在整個訓練過程中被吸收。

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

env.close()
del env