torchrl.envs 封包¶
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 預期存在的裝置。環境裝置並不意味著實際的步進操作將在裝置上計算(這是後端的責任,TorchRL 對此無能為力)。環境的裝置僅表示輸入到環境或從環境中檢索資料時預期存在的裝置。TorchRL 負責將資料對應到所需的裝置。這對於轉換尤其有用(請參閱下文)。對於參數化環境(例如,基於模型的環境),裝置確實代表將用於計算操作的硬體。env.observation_spec
:一個Composite
物件,包含所有觀察鍵-規範配對。env.state_spec
:一個Composite
物件,包含所有輸入鍵-規範配對(除了動作)。對於大多數有狀態環境,此容器將為空。env.action_spec
: 一個TensorSpec
物件,代表動作規格 (action spec)。env.reward_spec
: 一個TensorSpec
物件,代表獎勵規格 (reward spec)。env.done_spec
: 一個TensorSpec
物件,代表完成標記規格 (done-flag 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) 資料,可以使用 NonTensorSpec
實例。
重要的是,環境規格的形狀 (shapes) 應包含批次大小 (batch size),例如,一個 env.batch_size == torch.Size([4])
的環境應該有一個形狀為 torch.Size([4, action_size])
的 env.action_spec
。這在預先分配張量、檢查形狀一致性等方面很有幫助。
基於這些,實作了以下方法:
env.reset()
: 一個重置 (reset) 方法,可能會 (但不一定需要) 接收一個tensordict.TensorDict
輸入。它會回傳 rollout 的第一個 tensordict,通常包含一個"done"
狀態和一組觀察 (observations)。如果不存在,將會實例化一個 “reward” 鍵,其值為 0 且具有適當的形狀。env.step()
: 一個步進 (step) 方法,接收一個tensordict.TensorDict
輸入,其中包含輸入動作 (action) 以及其他輸入 (例如,針對基於模型的或無狀態的環境)。env.step_and_maybe_reset()
: 執行一個步進,並且在需要時 (部分) 重置環境。它會回傳更新後的輸入,其中包含一個"next"
鍵,該鍵包含下一個步進的資料,以及一個包含下一個步進的輸入資料的 tensordict (即,重置或結果或step_mdp()
)。這是透過讀取done_keys
並為每個完成狀態分配一個"_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=model
)。該策略應該使用tensordict.nn.TensorDictModule
(或任何其他與tensordict.TensorDict
相容的模組) 進行編碼。產生的tensordict.TensorDict
實例將標記一個尾隨的"time"
命名維度,其他模組可以使用該維度將此批次維度視為應有的樣子。
下圖總結了在 torchrl 中如何執行 rollout。

使用 TensorDict 的 TorchRL rollouts。¶
簡而言之,TensorDict 由 reset()
方法建立,然後由策略填入一個動作,之後傳遞給 step()
方法,該方法在 "next"
條目下寫入觀察、完成標記和獎勵。此呼叫的結果會被儲存以進行傳遞,而 "next"
條目則由 step_mdp()
函數收集。
注意
一般來說,所有 TorchRL 環境的輸出 tensordict 中都會有 "done"
和 "terminated"
條目。如果因為設計而沒有這些條目,EnvBase
元類別會確保每個 done 或 terminated 都會與其對應物配對。在 TorchRL 中,"done"
嚴格來說是指所有軌跡結束訊號的聯集,應被解釋為「軌跡的最後一步」或等效地「表示需要重置的訊號」。如果環境有提供(例如,Gymnasium),則截斷條目也會寫入 EnvBase.step()
輸出中的 "truncated"
條目下。如果環境僅帶有一個值,預設情況下它會被解釋為 "terminated"
訊號。預設情況下,TorchRL 的收集器和 rollout 方法會尋找 "done"
條目來評估是否應該重置環境。
注意
torchrl.collectors.utils.split_trajectories 函數可用於分割相鄰的軌跡。它依賴於輸入 tensordict 中的 "traj_ids"
條目,或者如果缺少 "traj_ids"
,則依賴於 "done"
和 "truncated"
鍵的結合。
注意
在某些情況下,標記軌跡的第一步可能很有用。TorchRL 通過 InitTracker
轉換來提供此功能。
我們的環境 教學 提供了更多關於如何從頭設計自訂環境的資訊。
|
抽象環境父類別。 |
|
類似 Gym 的環境是一種環境。 |
|
用於在多進程設定中儲存和傳遞環境元資料的類別。 |
向量化環境¶
向量化(或更好的是:平行)環境是強化學習中常見的一個特性,在其中執行環境步驟可能需要大量的 CPU 資源。一些函式庫,例如 gym3 或 EnvPool,提供了同時執行多批環境的介面。雖然它們通常提供非常有競爭力的計算優勢,但它們不一定能擴展到 TorchRL 支援的各種環境函式庫。因此,TorchRL 提供了自己的通用 ParallelEnv
類別來平行執行多個環境。由於此類別繼承自 SerialEnv
,因此它與其他環境享有完全相同的 API。當然,ParallelEnv
將具有與其環境數量相對應的批次大小。
注意
鑑於該函式庫的許多可選依賴項(例如,Gym、Gymnasium 和許多其他函式庫),警告在多進程/分散式設定中可能會很快變得非常煩人。預設情況下,TorchRL 會過濾子進程中的這些警告。如果仍然希望看到這些警告,可以通過設定 torchrl.filter_warnings_subprocess=False
來顯示它們。
重要的是,您的環境規格與其傳送和接收的輸入和輸出相符,因為 ParallelEnv
將從這些規格建立緩衝區,以便與 spawn 進程進行通訊。檢查 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"
鍵來指示環境應該重置哪個元件(子環境或代理)。這允許重置某些但不是所有的元件。
"_reset"
鍵有兩個不同的功能:1. 在調用 _reset()
期間,"_reset"
鍵可能存在於輸入 tensordict 中,也可能不存在。
TorchRL 的慣例是在給定的
"done"
層級缺少"_reset"
鍵表示該層級的完全重置(除非在較高層級找到"_reset"
鍵,請參閱下面的詳細資訊)。如果它存在,則預期那些條目和只有那些"_reset"
條目為True
的元件(沿著鍵和形狀維度)將被重置。環境在其
_reset()
方法中處理"_reset"
鍵的方式是其類別特有的。設計一個根據"_reset"
輸入運作的環境是開發人員的責任,因為 TorchRL 無法控制_reset()
的內部邏輯。儘管如此,在設計該方法時應牢記以下幾點。
在呼叫
_reset()
之後,輸出將會使用"_reset"
條目進行遮罩,並且先前step()
的輸出將會寫入"_reset"
為False
的位置。實際上,這表示如果"_reset"
修改了未公開的資料,則此修改將會遺失。在此遮罩操作之後,"_reset"
條目將會從reset()
的輸出中移除。
必須指出的是,"_reset"
是一個私有鍵,它應該只在編寫面向內部的特定環境功能時使用。換句話說,它 *不應該* 在函式庫外部使用,並且開發人員將保留修改透過 "_reset"
設定進行部分重置邏輯的權利,而無需事先保證,只要它們不影響 TorchRL 內部測試。
最後,設計重置功能時,應牢記以下假設:
每個
"_reset"
都與一個"done"
條目配對(+"terminated"
以及可能的"truncated"
)。這表示不允許以下結構:TensorDict({"done": done, "nested": {"_reset": reset}}, [])
,因為"_reset"
與"done"
位於不同的巢狀層級。一個層級的重置不會排除較低層級存在
"_reset"
的可能性,但它會消除其影響。原因很簡單,因為根層級的"_reset"
對應於all()
、any()
還是對巢狀"done"
條目的自定義呼叫,這是無法預先知道的,並且明確假設根層級的"_reset"
被放置在那裡是為了取代巢狀值(例如,可以查看PettingZooWrapper
實作,其中每個群組都有一個或多個與之關聯的"done"
條目,這些條目在根層級上使用any
或all
邏輯進行聚合,具體取決於任務)。當使用部分
"_reset"
條目呼叫env.reset(tensordict)()
時,該條目將重置一些但並非全部完成的子環境,輸入資料應包含 *未* 重置的子環境的資料。 此限制的原因在於,只能預測重置條目的env._reset(data)
的輸出。 對於其他條目,TorchRL 無法預先知道它們是否具有意義。 例如,可以完美地僅填充未重置元件的值,在這種情況下,未重置的資料將沒有意義,應將其丟棄。
下面,我們给出一些 "_reset"
鍵在重置後返回零的環境中預期產生的影響的範例
>>> # 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 second 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])
>>> 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
(和其他導入)所需的時間,啟動平行環境可能是一個瓶頸。 這就是為什麼例如 TorchRL 測試如此緩慢的原因。 一旦啟動環境,就應該觀察到很大的加速。
注意
TorchRL 需要精確的規格:另一個需要考慮的是 ParallelEnv
(以及資料收集器)將根據環境規格建立資料緩衝區,以便將資料從一個進程傳遞到另一個進程。 這表示規格錯誤(輸入、觀察或獎勵)將在運行時導致中斷,因為資料無法寫入預先分配的緩衝區中。 一般來說,在使用 ParallelEnv
之前,應該使用 check_env_specs()
測試函式測試環境。 只要預先分配的緩衝區和收集的資料不匹配,此函式就會引發斷言錯誤。
我們還提供 SerialEnv
類別,該類別享有完全相同的 API,但會依序執行。 這主要用於測試目的,當人們想要評估 ParallelEnv
的行為,而無需啟動子進程時。
除了提供基於進程並行的 ParallelEnv
之外,我們還提供了一種使用 MultiThreadedEnv
建立多執行緒環境的方法。這個類別底層使用了 EnvPool 函式庫,它可以提供更高的效能,但同時也限制了彈性 - 只能建立在 EnvPool
中實作的環境。這涵蓋了許多流行的 RL 環境類型(Atari、Classic Control 等),但不能像使用 ParallelEnv
那樣,使用任意的 TorchRL 環境。執行 benchmarks/benchmark_batched_envs.py 來比較平行化批次環境的不同方法的效能。
|
在同一個進程中建立一系列環境。批次環境允許使用者查詢遠端執行環境的任意方法/屬性。 |
|
為每個進程建立一個環境。 |
|
環境建立器類別。 |
自訂原生 TorchRL 環境¶
TorchRL 提供了一系列自訂的內建環境。
|
一個無狀態的 Pendulum 環境。 |
|
一個井字遊戲的實作。 |
多代理人環境¶
TorchRL 支援開箱即用的多代理人學習。在單代理人學習管道中使用的相同類別可以無縫地用於多代理人環境中,無需任何修改或專用的多代理人基礎設施。
在這個觀點中,環境在多代理人方面扮演核心角色。在多代理人環境中,許多決策代理人在一個共享的世界中行動。代理人可以觀察不同的事物,以不同的方式行動,並且獲得不同的獎勵。因此,存在許多典範來建模多代理人環境(DecPODPs、Markov Games)。這些典範之間的一些主要差異包括:
observation 可以是每個代理人的,也可以有一些共享的組件
reward 可以是每個代理人的,也可以是共享的
done (和
"truncated"
或"terminated"
) 可以是每個代理人的,也可以是共享的。
得益於它的 tensordict.TensorDict
資料載體,TorchRL 可以容納所有這些可能的典範。特別是,在多代理人環境中,每個代理人的鍵將被攜帶在巢狀的 “agents” TensorDict 中。這個 TensorDict 將具有額外的代理人維度,因此可以對每個代理人不同的資料進行分組。另一方面,共享的鍵將保持在第一層,就像在單代理人情況下一樣。
讓我們看一個範例來更好地理解這一點。在這個範例中,我們將使用 VMAS,一個也基於 PyTorch 的多機器人任務模擬器,它在設備上運行並行的批次模擬。
我們可以建立一個 VMAS 環境,並查看隨機步驟的輸出是什麼樣子
>>> 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,存在於具有批次大小 (num_envs,) 的根 tensordict 中,它代表模擬的環境數量。
另一方面,代理人之間不同的鍵,例如 action、reward、observation 和 info,存在於具有批次大小 (num_envs, n_agents) 的巢狀 “agents” tensordict 中,它代表額外的代理人維度。
多代理人張量規格將遵循與 tensordict 相同的樣式。與代理人之間變化的值相關的規格需要嵌套在 “agents” 條目中。
這是一個範例,說明如何在一個僅共享 done 標誌的多代理人環境中建立規格(如 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”]。對於 action 和 done 規格也是如此。請注意,env.reward_spec == env.output_spec[“full_reward_spec”][env.reward_key]。
|
Marl 群組映射類型。 |
|
檢查 MARL 群組映射。 |
自動重置環境¶
自動重置環境是指當環境在 rollout 期間達到 "done"
狀態時,不需要呼叫 reset()
的環境,因為重置會自動發生。 通常,在這種情況下,伴隨 done 和 reward 一起傳遞的觀察(實際上是環境中執行動作的結果)實際上是新 episode 的第一次觀察,而不是當前 episode 的最後一次觀察。
為了處理這些情況,torchrl 提供了一個 AutoResetTransform
,它會將 step 呼叫產生的觀察複製到下一個 reset,並在 rollout 期間跳過對 reset 的呼叫(在 rollout()
和 SyncDataCollector
迭代中)。 此轉換類別還提供了對無效觀察採用行為的細粒度控制,可以使用 “nan” 或任何其他值進行遮罩,或者根本不遮罩。
要告訴 torchrl 環境是自動重置的,只需在建構期間提供一個 auto_reset
參數。 如果提供,auto_reset_replace
參數還可以控制是否應該用一些佔位符替換 episode 的最後一次觀察的值。
>>> 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
。 由於數據無法連續堆疊,因此需要使用 return_contiguous=False
參數呼叫 env.rollout
。 這是一個可運作的範例
>>> 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()
將通過沿某些維度變化形狀的動態規格,但當一個 key 在一個 step 中存在而在其他 step 中不存在時,或者當維度數量變化時,將不會通過。
轉換¶
在大多數情況下,環境的原始輸出必須在傳遞給另一個物件(例如策略或價值運算符)之前進行處理。 為此,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"
關鍵字參數傳遞。
>>> env.append_transform(DoubleToFloat(in_keys_inv=["action"])) # will map the action from float32 to float64 before calling the base_env.step
in_keys
與 in_keys_inv
之間的關係可以透過將基礎環境視為轉換的「內部」部分來理解。 相反地,使用者輸入和輸出到轉換以及從轉換輸出,應被視為外部世界。 下圖顯示了這對於 RenameTransform
類別在實踐中的意義:step
函數的輸入 TensorDict
必須在其條目中列出 out_keys_inv
,因為它們是外部世界的一部分。 轉換會更改這些名稱,使它們與使用 in_keys_inv
的內部基礎環境的名稱相符。 反向流程會使用輸出 tensordict 執行,其中 in_keys
會被映射到對應的 out_keys
。

重新命名轉換邏輯¶
複製轉換¶
由於附加到環境的轉換會透過 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
|
環境轉換父類別。 |
|
一個經過轉換的環境。 |
|
一個離散化連續動作空間的轉換。 |
|
一個自適應動作遮罩器。 |
|
用於自動重設 env 的子類別。 |
|
用於自動重設環境的轉換。 |
|
一個用於修改環境批次大小的轉換。 |
|
如果獎勵為空或非空,則分別將獎勵映射到二進制值(0 或 1)。 |
|
轉換為部分燒錄資料序列。 |
|
將連續的觀察影格連接到單一張量中。 |
|
將多個鍵連接到單一張量中。 |
|
裁剪圖片的中心。 |
|
一個用於裁剪輸入(狀態、動作)或輸出(觀察、獎勵)值的轉換。 |
|
組合一系列轉換。 |
|
在指定位置和輸出大小裁剪輸入圖像。 |
|
將所選鍵的資料類型從一種轉換為另一種。 |
|
將資料從一個裝置移動到另一個裝置。 |
|
將離散動作從高維空間投影到低維空間。 |
|
將所選鍵的資料類型從一種轉換為另一種。 |
|
使用 lives 方法從 Gym 環境註冊生命週期結束訊號。 |
|
從資料中排除鍵。 |
此轉換將檢查 tensordict 的所有項目是否為有限值,如果不是,則引發例外。 |
|
|
展平張量的相鄰維度。 |
|
一個影格跳躍轉換。 |
|
將像素觀測值轉換為灰階。 |
|
重設追蹤器。 |
|
一個轉換,用於將 KL[pi_current||pi_0] 校正項新增到獎勵中。 |
|
當環境重設時,執行一系列隨機動作。 |
|
觀測值仿射轉換層。 |
|
觀測值轉換的抽象類別。 |
|
排列轉換。 |
在 tensordict 上呼叫 pin_memory 以方便在 CUDA 裝置上寫入。 |
|
|
R3M 轉換類別。 |
|
用於 ReplayBuffer 和模組的軌跡子採樣器。 |
|
從環境中移除空規格和內容。 |
|
一個轉換,用於重新命名輸出 tensordict 中的條目(或透過反向鍵重新命名輸入 tensordict)。 |
|
調整像素觀測值的大小。 |
|
根據 episode 獎勵和折扣因子計算 reward to go。 |
|
將獎勵限制在 clamp_min 和 clamp_max 之間。 |
|
獎勵的仿射轉換。 |
|
追蹤 episode 的累積獎勵。 |
|
從輸入 tensordict 中選擇鍵。 |
|
一個轉換,用於計算 TensorDict 值的符號。 |
|
移除指定位置大小為一的維度。 |
|
計算從重置開始的步數,並可選擇在達到一定步數後將截斷狀態設定為 |
|
設定 agent 在環境中要達成的目標回報。 |
|
一個用於在重置時初始化 TensorDict 的引導器。 |
|
取得過去 T 個觀測值中,每個位置的最大值。 |
|
將類似 numpy 的影像 (W x H x C) 轉換為 pytorch 影像 (C x W x H)。 |
|
在指定位置插入大小為 1 的維度。 |
|
VC1 轉換類別。 |
|
一個基於嵌入相似性計算獎勵的 VIP 轉換。 |
|
VIP 轉換類別。 |
|
GymWrapper 子類別的轉換,以一致的方式處理自動重置。 |
|
用於 torchrl 環境的移動平均正規化層。 |
|
一個 gSDE 雜訊初始化器。 |
具有遮罩動作的環境¶
在一些具有離散動作的環境中,agent 可用的動作可能會在執行過程中發生變化。在這種情況下,環境將輸出一個動作遮罩(預設情況下在 "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
中的呼叫計數)。
若要將收集到的 tensordict 儲存到磁碟上,可以使用 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 為止。使用者有責任在需要時呼叫 dumpy 以避免 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()
記錄器是在資料進入時註冊資料的轉換,用於記錄目的。
|
TensorDict 記錄器。 |
|
影片錄製轉換。 |
|
一個轉換,用於在父環境上呼叫 render,並將像素觀測結果註冊到 tensordict 中。 |
輔助函數¶
|
用於資料收集器的隨機策略。 |
|
針對短期 rollout 的結果測試環境規格。 |
傳回目前的取樣類型。 |
|
傳回所有支援的函式庫。 |
|
|
從 tensordict 建立一個 Composite 實例,假設所有值都是無界的。 |
|
|
|
建立一個新的 tensordict,反映輸入 tensordict 的一個時間步。 |
|
讀取 tensordict 內的 done / terminated / truncated 鍵,並寫入一個新的張量,其中聚合了兩個訊號的值。 |
特定領域¶
|
用於 Model Based RL sota 實現的基本環境。 |
|
Dreamer 模擬環境。 |
一個用於在 Dreamer 中記錄解碼觀測結果的轉換。 |
函式庫¶
TorchRL 的使命是盡可能簡化控制和決策演算法的訓練,而無需考慮正在使用的模擬器(如果有的話)。DMControl、Habitat、Jumanji 以及 Gym(自然也包括)都有多個封裝器可用。
最後一個函式庫在 RL 社群中具有特殊的地位,因為它是用於編碼模擬器最常用的框架。它成功的 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
|
使用環境名稱建構的 Google Brax 環境封裝器。 |
|
Google Brax 環境封裝器。 |
|
DeepMind Control lab 環境封裝器。 |
|
DeepMind Control lab 環境封裝器。 |
|
直接由環境 ID 建構的 OpenAI Gym 環境封裝器。 |
|
OpenAI Gym 環境包裝器。 |
|
Habitat 環境的包裝器。 |
|
用於 IsaacGym 環境的 TorchRL Env 介面。 |
|
用於 IsaacGymEnvs 環境的包裝器。 |
|
使用環境名稱構建的 Jumanji 環境包裝器。 |
|
Jumanji 環境包裝器。 |
|
Meltingpot 環境包裝器。 |
|
Meltingpot 環境包裝器。 |
|
FARAMA MO-Gymnasium 環境包裝器。 |
|
FARAMA MO-Gymnasium 環境包裝器。 |
|
基於 EnvPool 的環境多線程執行。 |
|
用於基於 envpool 的多線程環境的包裝器。 |
|
用於 OpenML 資料的環境介面,用於 bandits 情境。 |
|
Google DeepMind OpenSpiel 環境包裝器。 |
|
使用遊戲字串構建的 Google DeepMind OpenSpiel 環境包裝器。 |
|
PettingZoo 環境。 |
|
PettingZoo 環境包裝器。 |
|
RoboHive gym 環境的包裝器。 |
|
SMACv2 (StarCraft Multi-Agent Challenge v2) 環境包裝器。 |
|
SMACv2 (StarCraft Multi-Agent Challenge v2) 環境包裝器。 |
|
Unity ML-Agents 環境包裝器。 |
|
Unity ML-Agents 環境包裝器。 |
|
Vmas 環境包裝器。 |
|
Vmas 環境包裝器。 |
|
返回 gym 後端,或其子模組。 |
|
將 gym-backend 設置為特定值。 |