torchrl.envs package

TorchRL 提供了一個 API，用於處理不同後端的環境，例如 gym、dm-control、dm-lab、基於模型的環境以及自定義環境。目標是能夠在實驗中輕鬆切換環境，即使這些環境是使用不同的庫模擬的。TorchRL 在 torchrl.envs.libs 下提供了一些開箱即用的環境包裝器，我們希望這些包裝器可以輕鬆地被其他庫模仿。父類 EnvBase 是 torch.nn.Module 的子類，它使用 tensordict.TensorDict 作為資料組織者實現了一些典型的環境方法。這使得這個類具有通用性，並且能夠處理任意數量的輸入和輸出，以及巢狀或批次的資料結構。

每個環境都有以下屬性

• env.batch_size：表示批次處理的環境數量的 torch.Size。

• env.device：輸入和輸出 tensordict 預期所在的裝置。環境裝置不意味著實際的 step 操作將在裝置上計算（這是後端的責任，TorchRL 對此能做的很少）。環境的裝置僅表示資料在輸入環境或從環境檢索時預期所在的裝置。TorchRL 負責將資料對映到所需的裝置。這對於 transforms 尤其有用（見下文）。對於引數化環境（例如基於模型的環境），裝置確實代表了將用於計算操作的硬體。

• env.observation_spec：一個 Composite 物件，包含所有觀察（observation）的鍵-spec 對。

• env.state_spec：一個 Composite 物件，包含所有輸入的鍵-spec 對（除了 action）。對於大多數有狀態環境，這個容器將為空。

• env.action_spec：一個 TensorSpec 物件，表示 action spec。

• env.reward_spec：一個 TensorSpec 物件，表示 reward spec。

• env.done_spec：一個 TensorSpec 物件，表示 done 標誌 spec。請參閱下面的軌跡終止部分。

• env.input_spec：一個 Composite 物件，包含所有輸入鍵（"full_action_spec" 和 "full_state_spec"）。

• env.output_spec：一個 Composite 物件，包含所有輸出鍵（"full_observation_spec"、"full_reward_spec" 和 "full_done_spec"）。

如果環境攜帶非張量（non-tensor）資料，可以使用 NonTensor 例項。

環境 Specs：鎖和批次大小

環境 spec 預設是鎖定的（透過傳遞給環境建構函式的 spec_locked 引數）。鎖定 spec 意味著對 spec（或其子項，如果它是 Composite 例項）的任何修改都需要先解鎖。這可以透過 set_spec_lock_() 方法完成。預設鎖定 spec 的原因是這樣可以輕鬆快取值，例如 action 或 reset 鍵等。只有當預期 spec 會經常被修改時才應該解鎖環境（原則上應該避免這種情況）。允許進行諸如 env.observation_spec = new_spec 的 spec 修改：在底層，如果環境之前是鎖定的，TorchRL 將清除快取，解鎖 spec，進行修改，然後重新鎖定 spec。

重要的是，環境 spec 的形狀應該包含批次大小，例如，env.batch_size == torch.Size([4]) 的環境的 env.action_spec 應該具有 torch.Size([4, action_size]) 的形狀。這在預分配張量、檢查形狀一致性等方面很有幫助。

環境方法

基於這些，實現了以下方法

• env.reset()：一個 reset 方法，它可能（但不一定需要）接受 tensordict.TensorDict 輸入。它返回 rollout 的第一個 tensordict，通常包含一個 "done" 狀態和一組 observations。如果不存在，將使用 0 和適當的形狀例項化一個 “reward” 鍵。

• env.step()：一個 step 方法，它接受包含輸入 action 以及其他輸入（例如，對於基於模型的或無狀態環境）的 tensordict.TensorDict 輸入。

• env.step_and_maybe_reset()：執行一個 step，並在需要時（部分）重置環境。它返回帶有包含下一步資料的 "next" 鍵的更新輸入，以及一個包含下一步輸入資料（即 reset 或結果或 step_mdp() ）的 tensordict。這是透過讀取 done_keys 併為每個 done 狀態分配一個 "_reset" 訊號來完成的。這個方法可以輕鬆編寫不間斷的 rollout 函式

  >>> data_ = env.reset()
  >>> result = []
  >>> for i in range(N):
  ...     data, data_ = env.step_and_maybe_reset(data_)
  ...     result.append(data)
  ...
  >>> result = torch.stack(result)
  

• env.set_seed()：一個 seeding 方法，它將返回在多環境設定中使用的下一個種子。這個下一個種子是根據前一個種子確定性地計算出來的，這樣就可以用不同的種子為多個環境播種，而不會在連續的實驗中出現種子重疊的風險，同時仍然獲得可復現的結果。

• env.rollout()：在環境中執行一個 rollout，最大步數（max_steps=N），並使用 policy（policy=model）。policy 應該使用 tensordict.nn.TensorDictModule （或任何其他與 tensordict.TensorDict 相容的模組）編寫。生成的 tensordict.TensorDict 例項將標記一個尾部的 "time" 命名維度，其他模組可以使用這個維度來適當地處理這個批次維度。

下圖總結了 TorchRL 中如何執行 rollout。

使用 TensorDict 的 TorchRL Rollouts。

簡而言之，TensorDict 由 reset() 方法建立，然後由 policy 用 action 填充，再傳遞給 step() 方法，後者將 observations、done 標誌和 reward 寫入 "next" 條目下。此呼叫的結果被儲存以供後續使用，並且 "next" 條目由 step_mdp() 函式收集。

注意

通常，所有 TorchRL 環境在其輸出 tensordict 中都有 "done" 和 "terminated" 條目。如果設計中未包含它們，EnvBase 元類將確保每個 done 或 terminated 都伴隨其對應的條目。在 TorchRL 中，"done" 嚴格指代所有軌跡結束訊號的聯合，應解釋為“軌跡的最後一步”或等同於“指示需要重置的訊號”。如果環境提供了截斷資訊（例如 Gymnasium），截斷條目也會在 EnvBase.step() 輸出中寫入一個 "truncated" 條目下。如果環境攜帶一個單獨的值，預設情況下它將被解釋為 "terminated" 訊號。預設情況下，TorchRL 的 collectors 和 rollout 方法將查詢 "done" 條目來判斷環境是否應該重置。

注意

torchrl.collectors.utils.split_trajectories 函式可用於分割相鄰的軌跡。它依賴於輸入 tensordict 中的 "traj_ids" 條目，如果 "traj_ids" 缺失，則依賴於 "done" 和 "truncated" 鍵的組合。

注意

在某些上下文中，標記軌跡的第一步會很有用。TorchRL 透過 InitTracker transform 提供了此功能。

我們的環境 教程 提供了更多關於如何從頭設計自定義環境的資訊。

EnvBase(*args, **kwargs)	抽象環境父類。
GymLikeEnv(*args, **kwargs)	類似 gym 的環境是一個環境。
EnvMetaData(*, tensordict, specs, ...)	一個用於在多程序設定中儲存和傳遞環境元資料的類。

向量化環境

向量化（或者更確切地說：並行）環境是強化學習中的一個常見特性，其中執行環境 step 可能佔用大量 CPU 資源。一些庫，例如 gym3 或 EnvPool，提供了介面來同時執行批次環境。雖然它們通常提供非常有競爭力的計算優勢，但它們不一定能擴充套件到 TorchRL 支援的各種環境庫。因此，TorchRL 提供了自己的通用 ParallelEnv 類，用於並行執行多個環境。由於此類繼承自 SerialEnv，它享有與其他環境完全相同的 API。當然，ParallelEnv 將具有與其環境數量相對應的批次大小

注意

考慮到該庫的許多可選依賴項（例如 Gym、Gymnasium 等許多庫），在多程序/分散式設定中，警告可能很快變得相當煩人。預設情況下，TorchRL 在子程序中過濾掉這些警告。如果仍然希望看到這些警告，可以透過設定 torchrl.filter_warnings_subprocess=False 來顯示它們。

重要的是，你的環境 spec 必須與它傳送和接收的輸入和輸出相匹配，因為 ParallelEnv 將根據這些 spec 建立緩衝區以便與 spawned 程序通訊。使用 check_env_specs() 方法進行健全性檢查。

並行環境

     >>> def make_env():
     ...     return GymEnv("Pendulum-v1", from_pixels=True, g=9.81, device="cuda:0")
     >>> check_env_specs(env)  # this must pass for ParallelEnv to work
     >>> env = ParallelEnv(4, make_env)
     >>> print(env.batch_size)
     torch.Size([4])

ParallelEnv 允許檢索其所包含環境的屬性：可以直接呼叫

並行環境屬性

     >>> a, b, c, d = env.g  # gets the g-force of the various envs, which we set to 9.81 before
     >>> print(a)
     9.81

TorchRL 使用私有的 "_reset" 鍵來指示環境哪些元件（子環境或智慧體）應該被重置。

這允許只重置部分元件。

1. 在呼叫 _reset() 期間，輸入 tensordict 中可能存在或不存在 "_reset" 鍵。TorchRL 的約定是，在給定 "done" 級別上缺少 "_reset" 鍵表示該級別的完全重置（除非在更高級別上找到 "_reset" 鍵，詳見下文）。如果存在，則預期僅那些 "_reset" 條目為 True 的條目和元件（沿鍵和形狀維度）將被重置。

   環境在其 _reset() 方法中處理 "_reset" 鍵的方式特定於其類。根據 "_reset" 輸入設計環境行為是開發者的責任，因為 TorchRL 無法控制 _reset() 的內部邏輯。

2. 儘管如此，在設計該方法時應牢記以下幾點。呼叫 _reset() 後，輸出將用 "_reset" 條目進行掩碼，並且將前一個 step() 的輸出寫入 "_reset" 為 False 的位置。實際上，這意味著如果 "_reset" 修改了未被它暴露的資料，這個修改將會丟失。進行此掩碼操作後，"_reset" 條目將從 reset() 輸出中刪除。

必須指出，"_reset" 是一個私有鍵，它只能在編寫內部使用的特定環境功能時使用。換句話說，這不應該在庫外部使用，並且只要不影響 TorchRL 的內部測試，開發者將保留修改透過設定 "_reset" 實現部分重置邏輯的權利，恕不事先保證。

最後，在設計重置功能時，需要記住以下假設：

• 每個 "_reset" 都與一個 "done" 條目配對（+ "terminated"，並且可能還有 "truncated"）。這意味著以下結構是不允許的：TensorDict({"done": done, "nested": {"_reset": reset}}, [])，因為 "_reset" 所在的巢狀級別與 "done" 不同。

• 在一個級別上的重置並不排除在較低級別上存在 "_reset"，但它會使其效果失效。原因很簡單，根級別的 "_reset" 是否對應於對巢狀 "done" 條目的 all()、any() 或自定義呼叫無法提前得知，並且明確假定根級別的 "_reset" 被放置在那裡是為了取代巢狀值（例如，請檢視 PettingZooWrapper 的實現，其中每個組關聯有一個或多個 "done" 條目，這些條目根據任務以 any 或 all 邏輯在根級別進行聚合）。

• 當呼叫 env.reset(tensordict)() 並帶有將重置部分而非全部已完成子環境的部分 "_reset" 條目時，輸入資料應包含__未__被重置的子環境的資料。這個限制的原因在於，env._reset(data) 的輸出只能針對被重置的條目進行預測。

對於其他條目，TorchRL 無法提前知道它們是否有意義。例如，可以完全用填充值填充未重置元件的值，在這種情況下，未重置的資料將是無意義的，應該被丟棄。

>>> # single reset at the root
>>> data = TensorDict({"val": [1, 1], "_reset": [False, True]}, [])
>>> env.reset(data)
>>> print(data.get("val"))  # only the second value is 0
tensor([1, 0])
>>> # nested resets
>>> data = TensorDict({
...     ("agent0", "val"): [1, 1], ("agent0", "_reset"): [False, True],
...     ("agent1", "val"): [2, 2], ("agent1", "_reset"): [True, False],
... }, [])
>>> env.reset(data)
>>> print(data.get(("agent0", "val")))  # only the second value is 0
tensor([1, 0])
>>> print(data.get(("agent1", "val")))  # only the first value is 0
tensor([0, 2])
>>> # nested resets are overridden by a "_reset" at the root
>>> data = TensorDict({
...     "_reset": [True, True],
...     ("agent0", "val"): [1, 1], ("agent0", "_reset"): [False, True],
...     ("agent1", "val"): [2, 2], ("agent1", "_reset"): [True, False],
... }, [])
>>> env.reset(data)
>>> print(data.get(("agent0", "val")))  # reset at the root overrides nested
tensor([0, 0])
>>> print(data.get(("agent1", "val")))  # reset at the root overrides nested
tensor([0, 0])

下面，我們給出 "_reset" 鍵對重置後返回零的環境產生預期效果的一些示例

     >>> tensordict = TensorDict({"_reset": [[True], [False], [True], [True]]}, [4])
     >>> env.reset(tensordict)  # eliminates the "_reset" entry
     TensorDict(
         fields={
             terminated: Tensor(torch.Size([4, 1]), dtype=torch.bool),
             done: Tensor(torch.Size([4, 1]), dtype=torch.bool),
             pixels: Tensor(torch.Size([4, 500, 500, 3]), dtype=torch.uint8),
             truncated: Tensor(torch.Size([4, 1]), dtype=torch.bool),
         batch_size=torch.Size([4]),
         device=None,
         is_shared=True)

注意

並行環境重置

注意

關於效能的注意事項：啟動 ParallelEnv 可能需要相當長的時間，因為它需要啟動與程序數量相同的 python 例項。由於執行 import torch （以及其他 imports）所需的時間，啟動並行環境可能成為瓶頸。例如，這就是 TorchRL 測試如此緩慢的原因。一旦環境啟動，應該會觀察到顯著的加速。

TorchRL 需要精確的 specs：另一個需要考慮的問題是，ParallelEnv （以及資料收集器）將根據環境 spec 建立資料緩衝區，以便在不同程序之間傳遞資料。這意味著 spec 錯誤（輸入、觀察或 reward）將在執行時導致程式中斷，因為資料無法寫入預分配的緩衝區。通常，在使用 ParallelEnv 之前，應該使用 check_env_specs() 測試函式對環境進行測試。當預分配的緩衝區與收集到的資料不匹配時，此函式將引發一個斷言錯誤。

我們還提供了 SerialEnv 類，它享有完全相同的 API，但以序列方式執行。這主要用於測試目的，當想要評估 ParallelEnv 的行為而不啟動子程序時非常有用。

除了提供基於程序的並行性的 ParallelEnv 之外，我們還提供了使用 MultiThreadedEnv 建立多執行緒環境的方法。此類底層使用 EnvPool 庫，這可以實現更高的效能，但同時限制了靈活性 - 只能建立在 EnvPool 中實現的環境。這涵蓋了許多流行的強化學習環境型別（Atari、經典控制等），但不能使用任意的 TorchRL 環境，不像 ParallelEnv 那樣靈活。執行 benchmarks/benchmark_batched_envs.py 來比較並行化批次環境的不同方法的效能。	SerialEnv(*args, **kwargs)
在同一程序中建立一系列環境。批次環境允許使用者查詢遠端執行環境的任意方法/屬性。	ParallelEnv(*args, **kwargs)
為每個程序建立一個環境。	EnvCreator(create_env_fn[, ...])

環境建立器類。

自定義原生 TorchRL 環境

TorchRL 提供了一系列自定義的內建環境。	ChessEnv(*args, **kwargs)
一個遵循 TorchRL API 的國際象棋環境。	PendulumEnv(*args, **kwargs)
一個無狀態的擺錘環境。	TicTacToeEnv(*args, **kwargs)
一個井字棋實現。	LLMHashingEnv(*args, **kwargs)

一個文字生成環境，使用雜湊模組識別獨特的觀察。

多智慧體環境

TorchRL 開箱即用地支援多智慧體學習。在單智慧體學習管線中使用的相同類可以無縫地用於多智慧體場景，無需任何修改或專用的多智慧體基礎設施。

• 觀測（observation） 可以是每個智慧體的，也可以包含一些共享部分

• 獎勵（reward） 可以是每個智慧體的，也可以是共享的

• 完成（done）（以及 "truncated" 或 "terminated"）可以是每個智慧體的，也可以是共享的。

TorchRL 憑藉其 tensordict.TensorDict 資料載體，能夠適應所有這些可能的範例。特別是在多智慧體環境中，按智慧體的鍵將儲存在巢狀的“agents” TensorDict 中。這個 TensorDict 將具有額外的智慧體維度，從而將每個智慧體不同的資料分組。另一方面，共享鍵將保留在第一層，就像單智慧體情況中一樣。

讓我們看一個例子來更好地理解這一點。在本例中，我們將使用 VMAS，這是一個同樣基於 PyTorch 的多機器人任務模擬器，它在裝置上執行並行批處理模擬。

我們可以建立一個 VMAS 環境，並檢視隨機步驟的輸出是什麼樣子

多智慧體步驟 tensordict 示例

     >>> from torchrl.envs.libs.vmas import VmasEnv
     >>> env = VmasEnv("balance", num_envs=3, n_agents=5)
     >>> td = env.rand_step()
     >>> td
     TensorDict(
         fields={
             agents: TensorDict(
                 fields={
                     action: Tensor(shape=torch.Size([3, 5, 2]))},
                 batch_size=torch.Size([3, 5])),
             next: TensorDict(
                 fields={
                     agents: TensorDict(
                         fields={
                             info: TensorDict(
                                 fields={
                                     ground_rew: Tensor(shape=torch.Size([3, 5, 1])),
                                     pos_rew: Tensor(shape=torch.Size([3, 5, 1]))},
                                 batch_size=torch.Size([3, 5])),
                             observation: Tensor(shape=torch.Size([3, 5, 16])),
                             reward: Tensor(shape=torch.Size([3, 5, 1]))},
                         batch_size=torch.Size([3, 5])),
                     done: Tensor(shape=torch.Size([3, 1]))},
                 batch_size=torch.Size([3]))},
         batch_size=torch.Size([3]))

我們可以觀察到，所有智慧體共享的鍵，例如 done，存在於根 tensordict 中，其批大小為 (num_envs,)，這代表了模擬的環境數量。

另一方面，智慧體之間不同的鍵，例如 action、reward、observation 和 info，存在於巢狀的“agents” tensordict 中，其批大小為 (num_envs, n_agents)，這代表了額外的智慧體維度。

多智慧體張量規範將遵循與 tensordict 相同的風格。與智慧體之間變化的值相關的規範需要巢狀在“agents”條目中。

這裡有一個示例，說明如何在只有完成標誌在智慧體之間共享的多智慧體環境中建立規範（如在 VMAS 中）

多智慧體規範建立示例

     >>> action_specs = []
     >>> observation_specs = []
     >>> reward_specs = []
     >>> info_specs = []
     >>> for i in range(env.n_agents):
     ...    action_specs.append(agent_i_action_spec)
     ...    reward_specs.append(agent_i_reward_spec)
     ...    observation_specs.append(agent_i_observation_spec)
     >>> env.action_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"action": torch.stack(action_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.reward_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"reward": torch.stack(reward_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.observation_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"observation": torch.stack(observation_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.done_spec = Categorical(
     ...    n=2,
     ...    shape=torch.Size((1,)),
     ...    dtype=torch.bool,
     ... )

如你所見，這非常簡單！按智慧體的鍵將具有巢狀的複合規範，而共享鍵將遵循單智慧體標準。

注意

由於 reward、done 和 action 鍵可能帶有額外的“agent”字首（例如，(“agents”,”action”)），因此在其他 TorchRL 元件的引數中使用的預設鍵（例如，“action”）將無法完全匹配。因此，TorchRL 提供了 env.action_key、env.reward_key 和 env.done_key 屬性，它們會自動指向要使用的正確鍵。確保將這些屬性傳遞給 TorchRL 中的各種元件，以告知它們正確的鍵（例如，loss.set_keys() 函式）。

注意

TorchRL 抽象化了這些巢狀規範，以便於使用。這意味著，如果訪問的規範是 Composite，則訪問 env.reward_spec 將始終返回葉規範。因此，如果在上面的示例中，我們在建立環境後執行 env.reward_spec，我們將獲得與 torch.stack(reward_specs)} 相同的輸出。要獲取帶有“agents”鍵的完整複合規範，可以執行 env.output_spec[“full_reward_spec”]。動作和完成規範也是如此。請注意，env.reward_spec == env.output_spec[“full_reward_spec”][env.reward_key]。

MarlGroupMapType(value)	Marl 分組對映型別。
check_marl_grouping(group_map, agent_names)	檢查 MARL 分組對映。

自動重置環境

自動重置環境是指，當環境在 rollout 期間達到 "done" 狀態時，不需要呼叫 reset()，因為重置會自動發生。通常，在這種情況下，隨同 done 和 reward（這些實際上是在環境中執行動作的結果）傳遞的觀測是新回合的第一個觀測，而不是當前回合的最後一個觀測。

為了處理這些情況，torchrl 提供了一個 AutoResetTransform，它會將呼叫 step 產生的觀測複製到下一次 reset，並在 rollouts 期間（在 rollout() 和 SyncDataCollector 迭代中）跳過 reset 的呼叫。這個變換類還對無效觀測採取的行為提供了細粒度控制，這些觀測可以用 “nan” 或任何其他值進行遮蔽，或者完全不遮蔽。

要告知 torchrl 某個環境是自動重置的，只需在構造時提供一個 auto_reset 引數即可。如果提供，auto_reset_replace 引數還可以控制回合中最後一個觀測的值是否應被某些佔位符替換。

>>> from torchrl.envs import GymEnv
>>> from torchrl.envs import set_gym_backend
>>> import torch
>>> torch.manual_seed(0)
>>>
>>> class AutoResettingGymEnv(GymEnv):
...     def _step(self, tensordict):
...         tensordict = super()._step(tensordict)
...         if tensordict["done"].any():
...             td_reset = super().reset()
...             tensordict.update(td_reset.exclude(*self.done_keys))
...         return tensordict
...
...     def _reset(self, tensordict=None):
...         if tensordict is not None and "_reset" in tensordict:
...             return tensordict.copy()
...         return super()._reset(tensordict)
>>>
>>> with set_gym_backend("gym"):
...     env = AutoResettingGymEnv("CartPole-v1", auto_reset=True, auto_reset_replace=True)
...     env.set_seed(0)
...     r = env.rollout(30, break_when_any_done=False)
>>> print(r["next", "done"].squeeze())
tensor([False, False, False, False, False, False, False, False, False, False,
        False, False, False,  True, False, False, False, False, False, False,
        False, False, False, False, False,  True, False, False, False, False])
>>> print("observation after reset are set as nan", r["next", "observation"])
observation after reset are set as nan tensor([[-4.3633e-02, -1.4877e-01,  1.2849e-02,  2.7584e-01],
        [-4.6609e-02,  4.6166e-02,  1.8366e-02, -1.2761e-02],
        [-4.5685e-02,  2.4102e-01,  1.8111e-02, -2.9959e-01],
        [-4.0865e-02,  4.5644e-02,  1.2119e-02, -1.2542e-03],
        [-3.9952e-02,  2.4059e-01,  1.2094e-02, -2.9009e-01],
        [-3.5140e-02,  4.3554e-01,  6.2920e-03, -5.7893e-01],
        [-2.6429e-02,  6.3057e-01, -5.2867e-03, -8.6963e-01],
        [-1.3818e-02,  8.2576e-01, -2.2679e-02, -1.1640e+00],
        [ 2.6972e-03,  1.0212e+00, -4.5959e-02, -1.4637e+00],
        [ 2.3121e-02,  1.2168e+00, -7.5232e-02, -1.7704e+00],
        [ 4.7457e-02,  1.4127e+00, -1.1064e-01, -2.0854e+00],
        [ 7.5712e-02,  1.2189e+00, -1.5235e-01, -1.8289e+00],
        [ 1.0009e-01,  1.0257e+00, -1.8893e-01, -1.5872e+00],
        [        nan,         nan,         nan,         nan],
        [-3.9405e-02, -1.7766e-01, -1.0403e-02,  3.0626e-01],
        [-4.2959e-02, -3.7263e-01, -4.2775e-03,  5.9564e-01],
        [-5.0411e-02, -5.6769e-01,  7.6354e-03,  8.8698e-01],
        [-6.1765e-02, -7.6292e-01,  2.5375e-02,  1.1820e+00],
        [-7.7023e-02, -9.5836e-01,  4.9016e-02,  1.4826e+00],
        [-9.6191e-02, -7.6387e-01,  7.8667e-02,  1.2056e+00],
        [-1.1147e-01, -9.5991e-01,  1.0278e-01,  1.5219e+00],
        [-1.3067e-01, -7.6617e-01,  1.3322e-01,  1.2629e+00],
        [-1.4599e-01, -5.7298e-01,  1.5848e-01,  1.0148e+00],
        [-1.5745e-01, -7.6982e-01,  1.7877e-01,  1.3527e+00],
        [-1.7285e-01, -9.6668e-01,  2.0583e-01,  1.6956e+00],
        [        nan,         nan,         nan,         nan],
        [-4.3962e-02,  1.9845e-01, -4.5015e-02, -2.5903e-01],
        [-3.9993e-02,  3.9418e-01, -5.0196e-02, -5.6557e-01],
        [-3.2109e-02,  5.8997e-01, -6.1507e-02, -8.7363e-01],
        [-2.0310e-02,  3.9574e-01, -7.8980e-02, -6.0090e-01]])

動態規範

通常透過建立用於在程序之間傳遞資訊的記憶體緩衝區來並行執行環境。在某些情況下，可能無法預測環境在 rollout 期間是否會有一致的輸入或輸出，因為它們的形狀可能是可變的。我們將此稱為動態規範。

TorchRL 能夠處理動態規範，但需要讓批處理環境和收集器瞭解此特性。請注意，在實踐中，這是自動檢測的。

要指示張量沿某個維度具有可變大小，可以將所需維度的尺寸值設定為 -1。由於資料無法連續堆疊，呼叫 env.rollout 時需要使用 return_contiguous=False 引數。這裡是一個工作示例

>>> from torchrl.envs import EnvBase
>>> from torchrl.data import Unbounded, Composite, Bounded, Binary
>>> import torch
>>> from tensordict import TensorDict, TensorDictBase
>>>
>>> class EnvWithDynamicSpec(EnvBase):
...     def __init__(self, max_count=5):
...         super().__init__(batch_size=())
...         self.observation_spec = Composite(
...             observation=Unbounded(shape=(3, -1, 2)),
...         )
...         self.action_spec = Bounded(low=-1, high=1, shape=(2,))
...         self.full_done_spec = Composite(
...             done=Binary(1, shape=(1,), dtype=torch.bool),
...             terminated=Binary(1, shape=(1,), dtype=torch.bool),
...             truncated=Binary(1, shape=(1,), dtype=torch.bool),
...         )
...         self.reward_spec = Unbounded((1,), dtype=torch.float)
...         self.count = 0
...         self.max_count = max_count
...
...     def _reset(self, tensordict=None):
...         self.count = 0
...         data = TensorDict(
...             {
...                 "observation": torch.full(
...                     (3, self.count + 1, 2),
...                     self.count,
...                     dtype=self.observation_spec["observation"].dtype,
...                 )
...             }
...         )
...         data.update(self.done_spec.zero())
...         return data
...
...     def _step(
...         self,
...         tensordict: TensorDictBase,
...     ) -> TensorDictBase:
...         self.count += 1
...         done = self.count >= self.max_count
...         observation = TensorDict(
...             {
...                 "observation": torch.full(
...                     (3, self.count + 1, 2),
...                     self.count,
...                     dtype=self.observation_spec["observation"].dtype,
...                 )
...             }
...         )
...         done = self.full_done_spec.zero() | done
...         reward = self.full_reward_spec.zero()
...         return observation.update(done).update(reward)
...
...     def _set_seed(self, seed: Optional[int]):
...         self.manual_seed = seed
...         return seed
>>> env = EnvWithDynamicSpec()
>>> print(env.rollout(5, return_contiguous=False))
LazyStackedTensorDict(
    fields={
        action: Tensor(shape=torch.Size([5, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: LazyStackedTensorDict(
            fields={
                done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([5, 3, -1, 2]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            exclusive_fields={
            },
            batch_size=torch.Size([5]),
            device=None,
            is_shared=False,
            stack_dim=0),
        observation: Tensor(shape=torch.Size([5, 3, -1, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([5]),
    device=None,
    is_shared=False,
    stack_dim=0)

警告

在 ParallelEnv 和資料收集器中缺少記憶體緩衝區會極大地影響這些類的效能。任何此類用法都應與在單個程序上的普通執行仔細進行基準測試，因為序列化和反序列化大量張量可能非常昂貴。

目前，check_env_specs() 對於沿某些維度形狀變化的動態規範會透過，但對於在某個步驟中存在而在其他步驟中缺失的鍵，或維度數量變化的情況，則不會透過。

變換

在大多數情況下，環境的原始輸出在傳遞給其他物件（如策略或值運算子）之前必須經過處理。為此，TorchRL 提供了一系列變換，旨在復現 torch.distributions.Transform 和 torchvision.transforms 的變換邏輯。我們的環境 教程 提供了關於如何設計自定義變換的更多資訊。

變換環境使用 TransformedEnv 基本單元構建。複合變換使用 Compose 類構建

變換環境

     >>> base_env = GymEnv("Pendulum-v1", from_pixels=True, device="cuda:0")
     >>> transform = Compose(ToTensorImage(in_keys=["pixels"]), Resize(64, 64, in_keys=["pixels"]))
     >>> env = TransformedEnv(base_env, transform)

變換通常是 Transform 的子類，但也可以是任何 Callable[[TensorDictBase], TensorDictBase]。

預設情況下，變換環境將繼承傳遞給它的 base_env 的裝置。然後變換將在該裝置上執行。現在很明顯，根據要計算的操作型別，這可以帶來顯著的加速。

環境包裝器的一個巨大優勢是，可以檢視直到該包裝器的環境。使用 TorchRL 變換環境也可以實現同樣的效果：parent 屬性將返回一個新的 TransformedEnv，其中包含直到目標變換的所有變換。重用上面的示例

變換父級

     >>> resize_parent = env.transform[-1].parent  # returns the same as TransformedEnv(base_env, transform[:-1])

變換環境可以與向量化環境一起使用。由於每個變換都使用 "in_keys"/"out_keys" 關鍵字引數集，因此也很容易將變換圖連線到觀測資料的每個元件（例如畫素或狀態等）。

正向和逆向變換

變換還具有一個 inv() 方法，該方法在動作以逆序應用於複合變換鏈之前被呼叫。這允許在環境中執行動作之前對環境中的資料應用變換。要包含在此逆向變換中的鍵透過 “in_keys_inv” 關鍵字引數傳遞，並且在大多數情況下，out-keys 預設使用這些值

逆向變換

     >>> env.append_transform(DoubleToFloat(in_keys_inv=["action"]))  # will map the action from float32 to float64 before calling the base_env.step

以下段落詳細闡述瞭如何考慮哪些應被視為 in_ 或 out_ 特性。

理解變換鍵

在變換中，in_keys 和 out_keys 定義了基礎環境與外部世界（例如你的策略）之間的互動

• in_keys 指的是基礎環境的視角（內部 = TransformedEnv 的 base_env）。

• out_keys 指的是外部世界（外部 = policy、agent 等）。

例如，使用 in_keys=[“obs”] 和 out_keys=[“obs_standardized”] 時，策略將“看到”一個標準化觀測，而基礎環境輸出的是常規觀測。

類似地，對於逆向鍵

• in_keys_inv 指基礎環境看到的條目。

• out_keys_inv 指策略看到或產生的條目。

下圖說明了 RenameTransform 類的這個概念：step 函式的輸入 TensorDict 必須包含 out_keys_inv，因為它們是外部世界的一部分。變換使用 in_keys_inv 將這些名稱更改為與內部基礎環境的名稱匹配。逆向過程使用輸出 tensordict 執行，其中 in_keys 對映到對應的 out_keys。

Rename 變換邏輯

變換 Tensor 和規範

變換實際張量（來自策略）時，該過程示意性地表示為

>>> for t in reversed(self.transform):
...     td = t.inv(td)

這從最外層變換開始到最內層變換結束，確保暴露給策略的動作值被正確變換。

對於變換動作規範，過程應從最內層到最外層（類似於觀測規範）

>>> def transform_action_spec(self, action_spec):
...     for t in self.transform:
...         action_spec = t.transform_action_spec(action_spec)
...     return action_spec

單個 transform_action_spec 的虛擬碼可以是

>>> def transform_action_spec(self, action_spec):
...    return spec_from_random_values(self._apply_transform(action_spec.rand()))

這種方法確保“外部”規範是從“內部”規範推斷出來的。請注意，我們有意呼叫的是 _apply_transform 而不是 _inv_apply_transform！

向外部世界暴露規範

TransformedEnv 將暴露對應於動作和狀態的 out_keys_inv 的規範。例如，使用 ActionDiscretizer 時，環境的動作（例如 “action”）是一個浮點值張量，在使用變換環境呼叫 rand_action() 時不應生成該張量。相反，應該生成 “action_discrete”，並透過變換獲取其連續對應項。因此，使用者應看到 “action_discrete” 條目被暴露，而不是 “action”。

克隆變換

由於附加到環境的變換透過 transform.parent 屬性“註冊”到該環境，因此在操作變換時，我們應記住父級可能根據對變換進行的操作出現和消失。這裡有一些示例：如果我們從一個 Compose 物件獲取單個變換，這個變換將保留其父級

>>> third_transform = env.transform[2]
>>> assert third_transform.parent is not None

這意味著禁止將此變換用於另一個環境，因為另一個環境會替換父級，這可能導致意外行為。幸運的是，Transform 類帶有一個 clone() 方法，該方法將擦除父級，同時保留所有註冊緩衝區的標識。

>>> TransformedEnv(base_env, third_transform)  # raises an Exception as third_transform already has a parent
>>> TransformedEnv(base_env, third_transform.clone())  # works

在單個程序上或如果緩衝區放置在共享記憶體中，這將導致所有克隆的變換保持相同的行為，即使緩衝區被原地修改（例如，CatFrames 變換就是如此）。在分散式設定中，這可能不適用，在此上下文中應小心處理克隆變換的預期行為。最後，請注意，從 Compose 變換中索引多個變換也可能導致這些變換丟失父級：原因是索引 Compose 變換會產生另一個 Compose 變換，該變換沒有父環境。因此，我們必須克隆子變換才能建立這個其他組合

>>> env = TransformedEnv(base_env, Compose(transform1, transform2, transform3))
>>> last_two = env.transform[-2:]
>>> assert isinstance(last_two, Compose)
>>> assert last_two.parent is None
>>> assert last_two[0] is not transform2
>>> assert isinstance(last_two[0], type(transform2))  # and the buffers will match
>>> assert last_two[1] is not transform3
>>> assert isinstance(last_two[1], type(transform3))  # and the buffers will match

Transform([in_keys, out_keys, in_keys_inv, ...])	環境變換父類。
TransformedEnv(*args, **kwargs)	一個 transformed_in 環境。
ActionDiscretizer(num_intervals[, ...])	將連續動作空間離散化的變換。
ActionMask([action_key, mask_key])	自適應動作掩碼器。
AutoResetEnv(*args, **kwargs)	用於自動重置環境的子類。
AutoResetTransform(*[, replace, fill_float, ...])	用於自動重置環境的變換。
BatchSizeTransform(*[, batch_size, ...])	修改環境批大小的變換。
BinarizeReward([in_keys, out_keys])	如果獎勵為 null 或非 null，則將獎勵對映到二進位制值（分別為 0 或 1）。
BurnInTransform(modules, burn_in[, out_keys])	用於部分燒入資料序列的變換。
CatFrames(N, dim[, in_keys, out_keys, ...])	將連續觀測幀連線到單個張量中。
CatTensors([in_keys, out_key, dim, ...])	將多個鍵連線到單個張量中。
CenterCrop(w[, h, in_keys, out_keys])	裁剪影像中心。
ClipTransform([in_keys, out_keys, ...])	用於裁剪輸入（狀態、動作）或輸出（觀測、獎勵）值的變換。
Compose(*transforms)	組合一系列變換。
Crop(w[, h, top, left, in_keys, out_keys])	在指定位置和輸出尺寸裁剪輸入影像。
DTypeCastTransform(dtype_in, dtype_out[, ...])	將選定鍵的一種資料型別轉換為另一種。
DeviceCastTransform(device[, orig_device, ...])	將資料從一個裝置移動到另一個裝置。
DiscreteActionProjection(...[, action_key, ...])	將離散動作從高維空間投射到低維空間。
DoubleToFloat([in_keys, out_keys, ...])	將選定鍵的一種資料型別轉換為另一種。
EndOfLifeTransform([eol_key, lives_key, ...])	使用 lives 方法註冊來自具有 lives 方法的 Gym 環境的生命週期結束訊號。
ExcludeTransform(*excluded_keys[, inverse])	從資料中排除鍵。
FiniteTensorDictCheck()	此變換將檢查 tensordict 的所有項是否為有限值，如果不是則引發異常。
FlattenObservation(first_dim, last_dim[, ...])	展平張量的相鄰維度。
FrameSkipTransform([frame_skip])	幀跳躍變換。
GrayScale([in_keys, out_keys])	將畫素觀測轉換為灰度影像。
Hash(in_keys, out_keys, *[, hash_fn, seed, ...])	向 tensordict 新增雜湊值。
InitTracker	重置跟蹤器。
KLRewardTransform(actor[, coef, in_keys, ...])	向獎勵新增 KL[pi_current||pi_0] 校正項的變換。
LineariseRewards(in_keys[, out_keys, weights])	透過加權求和將多目標獎勵訊號轉換為單目標訊號。
NoopResetEnv([noops, random])	環境重置時執行一系列隨機動作。
ObservationNorm([loc, scale, in_keys, ...])	觀測仿射變換層。
ObservationTransform([in_keys, out_keys, ...])	觀測變換的抽象類。
PermuteTransform(dims[, in_keys, out_keys, ...])	置換變換。
PinMemoryTransform()	對 tensordict 呼叫 pin_memory 以方便在 CUDA 裝置上寫入。
R3MTransform(*args, **kwargs)	R3M 變換類。
RandomCropTensorDict(sub_seq_len[, ...])	ReplayBuffer 和模組的軌跡子取樣器。
RemoveEmptySpecs([in_keys, out_keys, ...])	從環境中移除空的規範和內容。
RenameTransform(in_keys, out_keys[, ...])	用於重新命名輸出 tensordict（或透過逆向鍵重新命名輸入 tensordict）中條目的變換。
Resize(w[, h, interpolation, in_keys, out_keys])	調整畫素觀測的大小。
Reward2GoTransform([gamma, in_keys, ...])	根據回合獎勵和折扣因子計算剩餘獎勵 (reward to go)。
RewardClipping([clamp_min, clamp_max, ...])	將獎勵裁剪在 clamp_min 和 clamp_max 之間。
RewardScaling(loc, scale[, in_keys, ...])	獎勵的仿射變換。
RewardSum([in_keys, out_keys, reset_keys, ...])	跟蹤回合累積獎勵。
SelectTransform(*selected_keys[, ...])	從輸入 tensordict 中選擇鍵。
SignTransform([in_keys, out_keys, ...])	計算 TensorDict 值符號的變換。
SqueezeTransform(*args, **kwargs)	移除指定位置上大小為 1 的維度。
Stack(in_keys, out_key[, in_key_inv, ...])	堆疊張量和 tensordict。
StepCounter([max_steps, truncated_key, ...])	計算從重置以來的步數，並在達到一定步數後可選地將截斷狀態設定為 True。
TargetReturn(target_return[, mode, in_keys, ...])	為智慧體設定一個在環境中達成的目標回報。
TensorDictPrimer([primers, random, ...])	在重置時用於 TensorDict 初始化的引導器。
TimeMaxPool([in_keys, out_keys, T, reset_key])	在過去 T 個觀測中，取每個位置的最大值。
Tokenizer(in_keys, out_keys[, in_keys_inv, ...])	對指定輸入應用分詞操作。
ToTensorImage([from_int, unsqueeze, dtype, ...])	將類似 numpy 的影像 (W x H x C) 轉換為 pytorch 影像 (C x W x H)。
TrajCounter([out_key])	全域性軌跡計數變換。
UnaryTransform(in_keys, out_keys[, ...])	對指定輸入應用一元操作。
UnsqueezeTransform(*args, **kwargs)	在指定位置插入一個大小為 1 的維度。
VC1Transform(in_keys, out_keys, model_name)	VC1 變換類。
VIPRewardTransform(*args, **kwargs)	基於嵌入相似性計算獎勵的 VIP 變換。
VIPTransform(*args, **kwargs)	VIP 變換類。
VecGymEnvTransform([final_name])	用於 GymWrapper 子類的變換，以一致的方式處理自動重置。
VecNorm([in_keys, out_keys, shared_td, ...])	torchrl 環境的移動平均歸一化層。
gSDENoise([state_dim, action_dim, shape])	一個 gSDE 噪聲初始化器。

帶有遮罩動作（masked actions）的環境

在一些具有離散動作的環境中，智慧體可用的動作可能會在執行過程中發生變化。在這種情況下，環境會輸出一個動作遮罩（預設鍵為 "action_mask"）。需要使用此遮罩來過濾掉該步驟中不可用的動作。

如果您正在使用自定義策略，可以將此遮罩傳遞給您的機率分佈，如下所示

帶有動作遮罩的類別策略

     >>> from tensordict.nn import TensorDictModule, ProbabilisticTensorDictModule, TensorDictSequential
     >>> import torch.nn as nn
     >>> from torchrl.modules import MaskedCategorical
     >>> module = TensorDictModule(
     >>>     nn.Linear(in_feats, out_feats),
     >>>     in_keys=["observation"],
     >>>     out_keys=["logits"],
     >>> )
     >>> dist = ProbabilisticTensorDictModule(
     >>>     in_keys={"logits": "logits", "mask": "action_mask"},
     >>>     out_keys=["action"],
     >>>     distribution_class=MaskedCategorical,
     >>> )
     >>> actor = TensorDictSequential(module, dist)

如果您想使用預設策略，則需要使用 ActionMask 轉換來包裝您的環境。此轉換可以負責更新動作規範中的動作遮罩，以便預設策略始終知道最新的可用動作。您可以按如下方式進行操作

如何使用動作遮罩轉換

     >>> from tensordict.nn import TensorDictModule, ProbabilisticTensorDictModule, TensorDictSequential
     >>> import torch.nn as nn
     >>> from torchrl.envs.transforms import TransformedEnv, ActionMask
     >>> env = TransformedEnv(
     >>>     your_base_env
     >>>     ActionMask(action_key="action", mask_key="action_mask"),
     >>> )

注意

如果您使用並行環境，將轉換新增到並行環境本身而不是其子環境非常重要。

記錄器

在環境 rollout 執行期間記錄資料對於關注演算法效能以及訓練後報告結果至關重要。

TorchRL 提供了幾個工具來與環境輸出互動：首先，可呼叫的回撥函式 callback 可以傳遞給 rollout() 方法。此函式將在每次 rollout 迭代時對收集到的 tensordict 進行呼叫（如果某些迭代需要跳過，應在 callback 內部新增一個內部變數來跟蹤呼叫次數）。

要將收集到的 tensordicts 儲存到磁碟，可以使用 TensorDictRecorder。

錄製影片

幾個後端提供了錄製環境渲染影像的可能性。如果畫素已經是環境輸出的一部分（例如 Atari 或其他遊戲模擬器），可以將 VideoRecorder 新增到環境中。此環境轉換將能夠錄製影片的日誌記錄器（例如 CSVLogger、WandbLogger 或 TensorBoardLogger）以及指示影片應儲存位置的標籤作為輸入。例如，要將 mp4 影片儲存到磁碟，可以使用帶有 video_format=”mp4” 引數的 CSVLogger。

VideoRecorder 轉換可以處理批次影像，並自動檢測 numpy 或 PyTorch 格式的影像（WHC 或 CWH）。

>>> logger = CSVLogger("dummy-exp", video_format="mp4")
>>> env = GymEnv("ALE/Pong-v5")
>>> env = env.append_transform(VideoRecorder(logger, tag="rendered", in_keys=["pixels"]))
>>> env.rollout(10)
>>> env.transform.dump()  # Save the video and clear cache

請注意，轉換的快取會一直增長，直到呼叫 dump。使用者有責任根據需要在必要時呼叫 dump 以避免 OOM（記憶體不足）問題。

在某些情況下，建立一個可以收集影像的測試環境既繁瑣又昂貴，或者根本不可能（某些庫只允許每個工作區一個環境例項）。在這種情況下，假設環境中有 render 方法可用，可以使用 PixelRenderTransform 在父環境上呼叫 render 並將影像儲存在 rollout 資料流中。此類別同樣適用於單個環境和批次環境。

>>> from torchrl.envs import GymEnv, check_env_specs, ParallelEnv, EnvCreator
>>> from torchrl.record.loggers import CSVLogger
>>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
>>>
>>> def make_env():
>>>     env = GymEnv("CartPole-v1", render_mode="rgb_array")
>>>     # Uncomment this line to execute per-env
>>>     # env = env.append_transform(PixelRenderTransform())
>>>     return env
>>>
>>> if __name__ == "__main__":
...     logger = CSVLogger("dummy", video_format="mp4")
...
...     env = ParallelEnv(16, EnvCreator(make_env))
...     env.start()
...     # Comment this line to execute per-env
...     env = env.append_transform(PixelRenderTransform())
...
...     env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
...     env.rollout(3)
...
...     check_env_specs(env)
...
...     r = env.rollout(30)
...     env.transform.dump()
...     env.close()

記錄器是註冊接收到的資料以便進行日誌記錄的轉換。

TensorDictRecorder(out_file_base[, ...])	TensorDict 記錄器。
VideoRecorder(logger, tag[, in_keys, skip, ...])	影片記錄器轉換。
PixelRenderTransform([out_keys, preproc, ...])	用於在父環境上呼叫 render 並在 tensordict 中註冊畫素觀察結果的轉換。

助手/工具函式

RandomPolicy(action_spec[, action_key])	用於資料收集器的隨機策略。
check_env_specs(env[, return_contiguous, ...])	根據短期 rollout 的結果測試環境規範。
exploration_type()	返回當前的取樣型別。
get_available_libraries()	返回所有支援的庫。
make_composite_from_td(data[, ...])	從 tensordict 建立一個 Composite 例項，假設所有值都是無界的。
set_exploration_type	`set_interaction_type` 的別名
step_mdp(tensordict[, next_tensordict, ...])	建立一個反映輸入 tensordict 時間步長的新 tensordict。
terminated_or_truncated(data[, ...])	讀取 tensordict 中的 done / terminated / truncated 鍵，並寫入一個新的 tensor，其中聚合了兩個訊號的值。

特定領域

ModelBasedEnvBase(*args, **kwargs)	用於基於模型的強化學習 sota 實現的基本環境。
model_based.dreamer.DreamerEnv(*args, **kwargs)	Dreamer 模擬環境。
model_based.dreamer.DreamerDecoder([...])	用於記錄 Dreamer 中解碼觀察結果的轉換。

庫

TorchRL 的使命是讓控制和決策演算法的訓練儘可能簡單，無論使用何種模擬器（如果使用）。針對 DMControl、Habitat、Jumanji 當然還有 Gym 提供了多個包裝器。

後一個庫在強化學習社群中具有特殊地位，是編寫模擬器最常用的框架。其成功的 API 奠定了基礎並啟發了許多其他框架，其中包括 TorchRL。然而，Gym 經歷了多次設計變更，作為外部採用庫有時難以適應這些變化：使用者通常有他們“偏好”的庫版本。此外，Gym 現在由另一個團隊以“gymnasium”的名義維護，這不利於程式碼相容性。實際上，我們必須考慮到使用者可能在同一個虛擬環境中安裝了 Gym 和 Gymnasium 的版本，並且必須允許兩者同時工作。幸運的是，TorchRL 為此問題提供瞭解決方案：特殊的裝飾器 set_gym_backend 允許控制相關函式中使用哪個庫。

>>> from torchrl.envs.libs.gym import GymEnv, set_gym_backend, gym_backend
>>> import gymnasium, gym
>>> with set_gym_backend(gymnasium):
...     print(gym_backend())
...     env1 = GymEnv("Pendulum-v1")
<module 'gymnasium' from '/path/to/venv/python3.9/site-packages/gymnasium/__init__.py'>
>>> with set_gym_backend(gym):
...     print(gym_backend())
...     env2 = GymEnv("Pendulum-v1")
<module 'gym' from '/path/to/venv/python3.9/site-packages/gym/__init__.py'>
>>> print(env1._env.env.env)
<gymnasium.envs.classic_control.pendulum.PendulumEnv at 0x15147e190>
>>> print(env2._env.env.env)
<gym.envs.classic_control.pendulum.PendulumEnv at 0x1629916a0>

我們可以看到這兩個庫修改了 gym_backend() 返回的值，該值可以進一步用於指示當前計算需要使用哪個庫。set_gym_backend 也是一個裝飾器：我們可以用它來告訴特定函式在其執行期間需要使用哪個 Gym 後端。torchrl.envs.libs.gym.gym_backend() 函式允許您獲取當前的 Gym 後端或其任何模組。

>>> import mo_gymnasium
>>> with set_gym_backend("gym"):
...     wrappers = gym_backend('wrappers')
...     print(wrappers)
<module 'gym.wrappers' from '/path/to/venv/python3.9/site-packages/gym/wrappers/__init__.py'>
>>> with set_gym_backend("gymnasium"):
...     wrappers = gym_backend('wrappers')
...     print(wrappers)
<module 'gymnasium.wrappers' from '/path/to/venv/python3.9/site-packages/gymnasium/wrappers/__init__.py'>

另一個與 Gym 和其他外部依賴項相關的有用工具是 torchrl._utils.implement_for 類。使用 @implement_for 裝飾函式會告訴 torchrl，根據指定的版本，預期會有特定的行為。這使得我們能夠輕鬆支援多個版本的 Gym，而無需使用者付出任何努力。例如，假設我們的虛擬環境安裝了 v0.26.2，那麼查詢時以下函式將返回 1。

>>> from torchrl._utils import implement_for
>>> @implement_for("gym", None, "0.26.0")
... def fun():
...     return 0
>>> @implement_for("gym", "0.26.0", None)
... def fun():
...     return 1
>>> fun()
1

BraxEnv(*args, **kwargs)	Google Brax 環境包裝器，使用環境名稱構建。
BraxWrapper(*args, **kwargs)	Google Brax 環境包裝器。
DMControlEnv(*args, **kwargs)	DeepMind Control lab 環境包裝器。
DMControlWrapper(*args, **kwargs)	DeepMind Control lab 環境包裝器。
GymEnv(*args, **kwargs)	OpenAI Gym 環境包裝器，直接由環境 ID 構建。
GymWrapper(*args, **kwargs)	OpenAI Gym 環境包裝器。
HabitatEnv(*args, **kwargs)	Habitat 環境的包裝器。
IsaacGymEnv(*args, **kwargs)	IsaacGym 環境的 TorchRL 環境介面。
IsaacGymWrapper(*args, **kwargs)	IsaacGymEnvs 環境的包裝器。
JumanjiEnv(*args, **kwargs)	Jumanji 環境包裝器，使用環境名稱構建。
JumanjiWrapper(*args, **kwargs)	Jumanji 的環境包裝器。
MeltingpotEnv(*args, **kwargs)	Meltingpot 環境包裝器。
MeltingpotWrapper(*args, **kwargs)	Meltingpot 環境包裝器。
MOGymEnv(*args, **kwargs)	FARAMA MO-Gymnasium 環境包裝器。
MOGymWrapper(*args, **kwargs)	FARAMA MO-Gymnasium 環境包裝器。
MultiThreadedEnv(*args, **kwargs)	基於 EnvPool 的環境的多執行緒執行。
MultiThreadedEnvWrapper(*args, **kwargs)	基於 envpool 的多執行緒環境的包裝器。
OpenMLEnv(*args, **kwargs)	用於 bandit 上下文中使用的 OpenML 資料的環境介面。
OpenSpielWrapper(*args, **kwargs)	Google DeepMind OpenSpiel 環境包裝器。
OpenSpielEnv(*args, **kwargs)	Google DeepMind OpenSpiel 環境包裝器，使用遊戲字串構建。
PettingZooEnv(*args, **kwargs)	PettingZoo 環境。
PettingZooWrapper(*args, **kwargs)	PettingZoo 環境包裝器。
RoboHiveEnv(*args, **kwargs)	用於 RoboHive gym 環境的包裝器。
SMACv2Env(*args, **kwargs)	SMACv2 (StarCraft Multi-Agent Challenge v2) 環境包裝器。
SMACv2Wrapper(*args, **kwargs)	SMACv2 (StarCraft Multi-Agent Challenge v2) 環境包裝器。
UnityMLAgentsEnv(*args, **kwargs)	Unity ML-Agents 環境包裝器。
UnityMLAgentsWrapper(*args, **kwargs)	Unity ML-Agents 環境包裝器。
VmasEnv(*args, **kwargs)	Vmas 環境包裝器。
VmasWrapper(*args, **kwargs)	Vmas 環境包裝器。
gym_backend([submodule])	返回 gym 後端或其子模組。
set_gym_backend(backend)	將 gym 後端設定為某個值。
register_gym_spec_conversion(spec_type)	註冊特定 spec 型別轉換函式的裝飾器。