LLMEnv¶

class torchrl.envs.llm.LLMEnv(*args, **kwargs)[原始碼]¶

一個用於語言模型的文字生成環境。

此環境旨在與語言模型配合使用，其中觀測值是表示令牌序列的字串或整數張量。動作也是一個字串或整數張量，它被連線到之前的觀測值以形成新的觀測值。

預設情況下，此環境用於跟蹤提示的歷史記錄。使用者可以新增轉換來根據其用例進行定製，例如思維鏈 (CoT) 推理或其他自定義處理。

使用者必須新增一個轉換來設定“done”條件，這將觸發載入下一個提示。當環境透過 from_dataloader() 建立時，可以在環境 reset 時載入語言模型的提示。

注意

LLMEnv 類的預設引數設定為使其易於使用 vllm 後端 (vLLMWrapper) 執行此環境。

關鍵字引數:

token_key (NestedKey, 可選) – 當 from_text=False 時，儲存令牌的 tensordict 中的鍵。預設為 "tokens"。
str_key (NestedKey, 可選) – 當 from_text=True 時，儲存字串輸入的 tensordict 中的鍵。預設為 "text"。
attention_key (NestedKey, 可選) – 儲存注意力掩碼的 tensordict 中的鍵。預設為 "attention_mask"。
action_key (NestedKey, 可選) – 儲存動作的 tensordict 中的鍵。預設為 "tokens_response" 或 "text_response"。
reward_key (NestedKey, 可選) – 如果 assign_reward=True，則儲存獎勵的 tensordict 中的鍵。預設為 "reward"。
from_text (bool, 可選) – 環境是否應期望字串作為輸入和輸出。預設為 True。
device (torch.device | None, 可選) – 環境應執行的裝置。預設為 None。
vocab_size (int | None, 可選) – 詞彙表的大小。如果為 None，則環境將假定一個無界詞彙表。預設為 None。
has_attention (bool, 可選) – 如果為 True，則使用由 attention_key 指示的鍵下的注意力掩碼。預設為 True。
assign_reward (bool, 可選) – 如果為 True，則在呼叫 step() 時將寫入形狀與動作形狀相同的零值獎勵。預設為 False。
assign_done (bool, 可選) –
如果為 True，則在呼叫 step() 時將寫入形狀與動作形狀相同的零值 done 和 terminated 狀態。預設為 False。 .. note:: 無論 assign_done 的值如何，done 狀態都會寫入根目錄

因為這是所有 TorchRL 環境的要求。
batch_size (int 或 torch.Size, 可選) –
環境的批處理大小。如果留空，則假定批處理大小為空。批處理大小可以為 null (torch.Size([])) 或一維。不支援無批處理的環境。

注意

使用 DataLoadingPrimer 轉換時，環境和轉換的批處理大小應匹配。
eos_token_id (int, 可選) – 序列結束的 token ID。如果傳遞，則在檢測到時將 done 狀態設定為 True。預設為 None。

另請參閱

DataLoadingPrimer 示例。

from_dataloader()[原始碼]¶: 從資料載入器建立 LLMEnv 例項。

property action_key: NestedKey¶

環境的 action 鍵。

預設情況下，這通常是 “action”。

如果環境中存在多個 action 鍵，此函式將引發異常。

property action_keys: list[tensordict._nestedkey.NestedKey]¶

環境的 action 鍵。

預設情況下，只有一個名為 “action” 的鍵。

鍵按資料樹的深度排序。

property action_spec: TensorSpec¶

action spec。

的 action_spec 始終儲存為複合 spec。

如果 action spec 作為簡單 spec 提供，則將返回該 spec。

>>> env.action_spec = Unbounded(1)
>>> env.action_spec
UnboundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

如果 action spec 作為複合 spec 提供且僅包含一個葉子，則此函式將僅返回該葉子。

>>> env.action_spec = Composite({"nested": {"action": Unbounded(1)}})
>>> env.action_spec
UnboundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

如果 action spec 作為複合 spec 提供且包含多個葉子，則此函式將返回整個 spec。

>>> env.action_spec = Composite({"nested": {"action": Unbounded(1), "another_action": Categorical(1)}})
>>> env.action_spec
Composite(
    nested: Composite(
        action: UnboundedContinuous(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        another_action: Categorical(
            shape=torch.Size([]),
            space=DiscreteBox(n=1),
            device=cpu,
            dtype=torch.int64,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

要檢索傳遞的完整 spec，請使用

>>> env.input_spec["full_action_spec"]

此屬性是可變的。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.action_spec
BoundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

property action_spec_unbatched: TensorSpec¶: 返回環境的 action spec，就好像它沒有批次維度一樣。

add_module(name: str, module: Optional[Module]) → None¶

將子模組新增到當前模組。

可以使用給定的名稱作為屬性訪問該模組。

引數:

name (str) – 子模組的名稱。子模組可以透過給定名稱從此模組訪問
module (Module) – 要新增到模組中的子模組。

add_truncated_keys() → EnvBase¶: 將截斷鍵新增到環境中。

all_actions(tensordict: TensorDictBase | None = None) → TensorDictBase¶

從 action spec 生成所有可能的 action。

這僅適用於具有完全離散 action 的環境。

引數:: tensordict (TensorDictBase, 可選) – 如果提供，將使用此 tensordict 呼叫 reset()。
返回:: 一個 tensordict 物件，其中 “action” 條目已更新為批次的所有可能 action。Action 被堆疊在主維度上。

any_done(tensordict: TensorDictBase) → bool¶

檢查 tensordict 是否處於“結束”狀態（或批次中的某個元素是否處於）。

結果將寫入 “_reset” 條目。

返回: 一個布林值，指示 tensordict 中是否有標記: 為結束的元素。

注意

傳入的 tensordict 應該是 “next” tensordict 或等價物——即，它不應包含 “next” 值。

append_transform(transform: Transform | Callable[[TensorDictBase], TensorDictBase]) → torchrl.envs.TransformedEnv¶

返回一個轉換後的環境，其中應用了傳入的可呼叫函式/轉換。

引數:: transform (Transform 或 Callable[[TensorDictBase], TensorDictBase]) – 要應用於環境的轉換。

示例

>>> from torchrl.envs import GymEnv
>>> import torch
>>> env = GymEnv("CartPole-v1")
>>> loc = 0.5
>>> scale = 1.0
>>> transform = lambda data: data.set("observation", (data.get("observation") - loc)/scale)
>>> env = env.append_transform(transform=transform)
>>> print(env)
TransformedEnv(
    env=GymEnv(env=CartPole-v1, batch_size=torch.Size([]), device=cpu),
    transform=_CallableTransform(keys=[]))

apply(fn: Callable[[Module], None]) → Self¶

將 fn 遞迴應用於每個子模組（由 .children() 返回）以及自身。

典型用法包括初始化模型引數（另請參閱 torch.nn.init）。

引數:: fn (Module -> None) – 要應用於每個子模組的函式
返回:: self
返回型別:: 模組

示例

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

auto_specs_(policy: Callable[[TensorDictBase], TensorDictBase], *, tensordict: TensorDictBase | None = None, action_key: NestedKey | list[NestedKey] = 'action', done_key: NestedKey | list[NestedKey] | None = None, observation_key: NestedKey | list[NestedKey] = 'observation', reward_key: NestedKey | list[NestedKey] = 'reward')¶

根據使用給定策略的隨機 rollout 自動設定環境的規範（specs）。

此方法使用提供的策略執行 rollout，以推斷環境的輸入和輸出規範。它會根據 rollout 期間收集的資料更新環境的操作、觀察、獎勵和完成訊號的規範。

引數:

policy (Callable[[TensorDictBase], TensorDictBase]) – 一個可呼叫的策略，接受 TensorDictBase 作為輸入並返回 TensorDictBase 作為輸出。此策略用於執行 rollout 並確定規範。

關鍵字引數:

tensordict (TensorDictBase, optional) – 可選的 TensorDictBase 例項，用作 rollout 的初始狀態。如果未提供，將呼叫環境的 reset 方法來獲取初始狀態。
action_key (NestedKey 或 List[NestedKey], optional) – 在 TensorDictBase 中用於標識 action 的鍵。預設為 “action”。
done_key (NestedKey 或 List[NestedKey], optional) – 用於在 TensorDictBase 中標識完成訊號的鍵。預設為 None，它將嘗試使用 [“done”, “terminated”, “truncated”] 作為潛在鍵。
observation_key (NestedKey 或 List[NestedKey], optional) – 在 TensorDictBase 中用於標識 observation 的鍵。預設為 “observation”。
reward_key (NestedKey 或 List[NestedKey], optional) – 在 TensorDictBase 中用於標識 reward 的鍵。預設為 “reward”。

返回:

已更新 spec 的環境例項。

返回型別:

EnvBase

丟擲:

RuntimeError – 如果輸出 spec 中存在未被提供的鍵所覆蓋的鍵。

property batch_dims: int¶: 環境的批次維度數。

property batch_locked: bool¶

環境是否可以用於與初始化時不同的批次大小。

如果為 True，則需要在與環境相同批次大小的 tensordict 上使用該環境。batch_locked 是一個不可變屬性。

property batch_size: Size¶

此環境例項中批次化環境的數量，組織為 torch.Size() 物件。

環境可能相似或不同，但假定它們之間幾乎沒有（如果有的話）互動（例如，多工或並行批處理執行）。

bfloat16() → Self¶

將所有浮點引數和緩衝區轉換為 bfloat16 資料型別。

注意

此方法就地修改模組。

返回:: self
返回型別:: 模組

buffers(recurse: bool = True) → Iterator[Tensor]¶

返回模組緩衝區的迭代器。

引數:: recurse (bool) – 如果為 True，則會產生此模組及其所有子模組的 buffer。否則，僅會產生此模組的直接成員 buffer。
產生:: torch.Tensor – 模組緩衝區

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

cardinality(tensordict: TensorDictBase | None = None) → int¶

動作空間的基數。

預設情況下，這只是 env.action_space.cardinality 的一個包裝器。

此類在動作規範可變時很有用

動作數量可能未定義，例如 Categorical(n=-1)；
動作基數可能取決於動作掩碼；
形狀可以是動態的，如 Unbound(shape=(-1))。

在這些情況下，應覆蓋 cardinality()，

引數:: tensordict (TensorDictBase, optional) – 包含計算基數所需資料的 tensordict。

check_env_specs(*args, **kwargs)¶

使用簡短的 rollout 來測試環境規範。

此測試函式應作為 torchrl 的 EnvBase 子類包裝的環境的健全性檢查：預期的資料與收集到的資料之間的任何差異都應引發斷言錯誤。

損壞的環境規範很可能會使並行環境無法使用。

引數:

env (EnvBase) – 要檢查其規格與資料是否匹配的環境。
return_contiguous (bool, optional) – 如果 True，則會使用 return_contiguous=True 呼叫隨機回放。這在某些情況下會失敗（例如，輸入/輸出的異構形狀）。預設為 None（由動態 spec 的存在決定）。
check_dtype (bool, optional) – 如果為 False，則會跳過 dtype 檢查。預設為 True。
seed (int, optional) – 為了可復現性，可以設定一個種子。種子會臨時設定在 pytorch 中，然後 RNG 狀態會恢復到之前的狀態。對於 env，我們設定了種子，但由於將 RNG 狀態恢復到之前的狀態不是大多數環境的功能，因此我們將其留給使用者自行完成。預設為 None。
tensordict (TensorDict, optional) – 用於重置的可選 tensordict 例項。
break_when_any_done (bool 或 str, optional) – break_when_any_done 在 rollout() 中的值。如果為 "both"，則測試會在 True 和 False 上執行。

注意：此函式會重置環境種子。它應該“離線”使用，以檢查環境是否已充分構建，但它可能會影響實驗的播種，因此應將其排除在訓練指令碼之外。

children() → Iterator[Module]¶

返回直接子模組的迭代器。

產生:: Module – 子模組

property collector: DataCollectorBase | None¶: 返回與容器關聯的收集器（如果存在）。

compile(*args, **kwargs)¶

使用 torch.compile() 編譯此 Module 的前向傳播。

此 Module 的 __call__ 方法將被編譯，並且所有引數將按原樣傳遞給 torch.compile()。

有關此函式的引數的詳細資訊，請參閱 torch.compile()。

cpu() → Self¶

將所有模型引數和緩衝區移動到 CPU。

注意

此方法就地修改模組。

返回:: self
返回型別:: 模組

cuda(device: Optional[Union[device, int]] = None) → Self¶

將所有模型引數和緩衝區移動到 GPU。

這也會使相關的引數和緩衝區成為不同的物件。因此，如果模組在最佳化時將駐留在 GPU 上，則應在構建最佳化器之前呼叫此函式。

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數將複製到該裝置
返回:: self
返回型別:: 模組

property done_key¶

環境的 done 鍵。

預設為“done”。

如果環境中存在多個 done 鍵，此函式將引發異常。

property done_keys: list[tensordict._nestedkey.NestedKey]¶

環境的 done 鍵。

預設情況下，只有一個名為“done”的鍵。

鍵按資料樹的深度排序。

property done_keys_groups¶

done 鍵的列表，按重置鍵分組。

這是一個列表的列表。外層列表的長度等於重置鍵的數量，內層列表包含 done 鍵（例如，done 和 truncated），這些鍵可以在 absence 時讀取以確定重置。

property done_spec: TensorSpec¶

done 規範。

done_spec 始終儲存為複合規範。

如果 done 規範作為簡單規範提供，則將返回該規範。

>>> env.done_spec = Categorical(2, dtype=torch.bool)
>>> env.done_spec
Categorical(
    shape=torch.Size([]),
    space=DiscreteBox(n=2),
    device=cpu,
    dtype=torch.bool,
    domain=discrete)

如果 done 規範作為複合規範提供且僅包含一個葉子，則此函式將僅返回該葉子。

>>> env.done_spec = Composite({"nested": {"done": Categorical(2, dtype=torch.bool)}})
>>> env.done_spec
Categorical(
    shape=torch.Size([]),
    space=DiscreteBox(n=2),
    device=cpu,
    dtype=torch.bool,
    domain=discrete)

如果 done 規範作為複合規範提供且具有多個葉子，則此函式將返回整個規範。

>>> env.done_spec = Composite({"nested": {"done": Categorical(2, dtype=torch.bool), "another_done": Categorical(2, dtype=torch.bool)}})
>>> env.done_spec
Composite(
    nested: Composite(
        done: Categorical(
            shape=torch.Size([]),
            space=DiscreteBox(n=2),
            device=cpu,
            dtype=torch.bool,
            domain=discrete),
        another_done: Categorical(
            shape=torch.Size([]),
            space=DiscreteBox(n=2),
            device=cpu,
            dtype=torch.bool,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

要始終檢索傳入的完整規範，請使用

>>> env.output_spec["full_done_spec"]

此屬性是可變的。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.done_spec
Categorical(
    shape=torch.Size([1]),
    space=DiscreteBox(n=2),
    device=cpu,
    dtype=torch.bool,
    domain=discrete)

property done_spec_unbatched: TensorSpec¶: 返回環境的 done 規範，就好像它沒有批處理維度一樣。

double() → Self¶

將所有浮點引數和緩衝區轉換為 double 資料型別。

注意

此方法就地修改模組。

返回:: self
返回型別:: 模組

empty_cache()¶

清除所有快取值。

對於常規環境，鍵列表（獎勵、完成等）會被快取，但在某些情況下，它們可能在程式碼執行期間發生更改（例如，新增轉換時）。

eval() → Self¶

將模組設定為評估模式。

這僅對某些模組有影響。有關模組在訓練/評估模式下的行為，例如它們是否受影響（如 Dropout、BatchNorm 等），請參閱具體模組的文件。

這等同於 self.train(False)。

有關 .eval() 和幾種可能與之混淆的類似機制之間的比較，請參閱區域性停用梯度計算。

返回:: self
返回型別:: 模組

extra_repr() → str¶

返回模組的額外表示。

要列印自定義額外資訊，您應該在自己的模組中重新實現此方法。單行和多行字串均可接受。

fake_tensordict() → TensorDictBase¶: 返回一個假的 tensordict，其鍵值對在形狀、裝置和 dtype 上與環境 rollout 期間預期的一致。

float() → Self¶

將所有浮點引數和緩衝區轉換為 float 資料型別。

注意

此方法就地修改模組。

返回:: self
返回型別:: 模組

forward(*args, **kwargs)¶

定義每次呼叫時執行的計算。

所有子類都應重寫此方法。

注意

儘管前向傳播的實現需要在此函式中定義，但您應該在之後呼叫 Module 例項而不是此函式，因為前者會處理註冊的鉤子，而後者則會靜默忽略它們。

classmethod from_dataloader(dataloader: DataLoader, *, tokenizer: transformers.PretrainedTokenizerBase | None = None, token_key: NestedKey | None = None, str_key: NestedKey | None = None, attention_key: NestedKey | None = None, action_key: NestedKey | None = None, reward_key: NestedKey = 'reward', from_text: bool = True, device: torch.device | None = None, vocab_size: int | None = None, batch_size: int | torch.Size | None = None, has_attention: bool = True, assign_reward: bool = False, assign_done: bool = False, primers: Composite | None = None, example_data: Any = None, stack_method: Callable[[Any], Any] | Literal['as_nested_tensor', 'as_padded_tensor'] = None, repeats: int | None = None, group_repeats: bool = True, eos_token_id: int | None = None) → LLMEnv[原始碼]¶

從資料載入器建立 LLMEnv 例項。

此方法建立一個 LLMEnv 例項併為其新增 DataLoadingPrimer，該 primer 在重置環境時使用來自提供的資料載入器的資料填充 data_keys（預設為 observation_key）。

引數:

dataloader (DataLoader) – 要從中載入資料的資料載入器。

關鍵字引數:

tokenizer (transformers.PretrainedTokenizerBase 或 str, 可選) –
要使用的分詞器。如果為 None，則預設使用“bert-base-uncased”。如果提供字串，則應為預訓練分詞器的名稱。

注意

使用 tokenizer 將會向環境新增一個 Tokenizer 轉換。如果 from_text 設定為 True，則分詞器將在每次迭代中呼叫，並且 rollout 將同時包含 tokens 和 text 資料。如果 from_text 設定為 False，則分詞器僅在重置時呼叫，並且 rollout 中唯一的文字資料將是從資料集中取樣的文字。
token_key (NestedKey, 可選) – 當 from_text=False 時，儲存令牌的 tensordict 中的鍵。預設為 ("tokens_in", "input_ids")。
str_key (NestedKey, 可選) – 當 from_text=True 時，儲存字串輸入的 tensordict 中的鍵。預設為 "test"。
attention_key (NestedKey, 可選) – 儲存注意力掩碼的 tensordict 中的鍵。預設為 ("tokens_in", "input_ids")
action_key (NestedKey, 可選) – 儲存動作的 tensordict 中的鍵。預設為 ("tokens_out", "sequences")。
reward_key (NestedKey, 可選) – 如果 assign_reward=True，則儲存獎勵的 tensordict 中的鍵。預設為 "reward"。
from_text (bool, 可選) – 環境是否應期望字串作為輸入和輸出。預設為 True。
device (torch.device | None, 可選) – 環境應執行的裝置。預設為 None。
vocab_size (int | None, 可選) – 詞彙表的大小。如果為 None，則環境將假定一個無界詞彙表。預設為 None。
has_attention (bool, 可選) – 如果為 True，則使用由 attention_key 指示的鍵下的注意力掩碼。預設為 True。
assign_reward (bool, 可選) – 如果為 True，則在呼叫 step() 時將寫入形狀與動作形狀相同的零值獎勵。預設為 False。
assign_done (bool, 可選) –
如果為 True，則在呼叫 step() 時將寫入形狀與動作形狀相同的零值 done 和 terminated 狀態。預設為 False。

注意

無論 assign_done 的值如何，done 狀態都會寫入根目錄，因為這是所有 TorchRL 環境的要求。
batch_size (int 或 torch.Size, 可選) –
環境的批處理大小。如果留空，則從 dataloader.batch_size 推斷批處理大小（如果該屬性存在），否則設定為 ()。批處理大小可以為 null (torch.Size([])) 或一維。不支援無批處理的環境。

注意

使用 DataLoadingPrimer 轉換時，環境和轉換的批處理大小應匹配。
primers (Composite | None, 可選) – 資料載入器中每個鍵使用的 primer。預設為 None（從第一個資料批次自動推斷）。
stack_method (Callable[[Any], Any] | Literal["as_nested_tensor", "as_padded_tensor"], 可選) – 用於堆疊資料的方法。預設為 None。
repeats (int, 可選) – 同一樣本需要連續出現多少次。這在 GRPO 等情況下很有用，其中單個提示會多次用於通過蒙特卡洛樣本估計優勢（而不是優勢模組）。
group_repeats (bool, 可選) – 如果為 True，則批處理大小乘以重複次數，以便所有重複項都分組在一個從緩衝區收集的批次中。預設為 True。
eos_token_id (int, 可選) – 序列結束的 token ID。如果傳遞，則在檢測到時將 done 狀態設定為 True。預設為 None。

返回:

建立的 LLMEnv 例項。

返回型別:

LLMEnv

property full_action_spec: Composite¶

完整的動作規範。

full_action_spec 是一個 Composite` 例項，其中包含所有動作條目。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.full_action_spec
Composite(
    action: BoundedContinuous(
        shape=torch.Size([8]),
        space=ContinuousBox(
            low=Tensor(shape=torch.Size([8]), device=cpu, dtype=torch.float32, contiguous=True),
            high=Tensor(shape=torch.Size([8]), device=cpu, dtype=torch.float32, contiguous=True)),
        device=cpu,
        dtype=torch.float32,
        domain=continuous), device=cpu, shape=torch.Size([]))

property full_action_spec_unbatched: Composite¶: 返回環境的 action spec，就好像它沒有批次維度一樣。

property full_done_spec: Composite¶

完整的 done 規範。

full_done_spec 是一個 Composite` 例項，其中包含所有完成條目。它可以用於生成結構類似於執行時獲得的假資料。

示例

>>> import gymnasium
>>> from torchrl.envs import GymWrapper
>>> env = GymWrapper(gymnasium.make("Pendulum-v1"))
>>> env.full_done_spec
Composite(
    done: Categorical(
        shape=torch.Size([1]),
        space=DiscreteBox(n=2),
        device=cpu,
        dtype=torch.bool,
        domain=discrete),
    truncated: Categorical(
        shape=torch.Size([1]),
        space=DiscreteBox(n=2),
        device=cpu,
        dtype=torch.bool,
        domain=discrete), device=cpu, shape=torch.Size([]))

property full_done_spec_unbatched: Composite¶: 返回環境的 done 規範，就好像它沒有批處理維度一樣。

property full_observation_spec_unbatched: Composite¶: 返回環境的 observation 規範，就好像它沒有批處理維度一樣。

property full_reward_spec: Composite¶

完整的 reward 規範。

full_reward_spec 是一個 Composite` 例項，其中包含所有 reward 條目。

示例

>>> import gymnasium
>>> from torchrl.envs import GymWrapper, TransformedEnv, RenameTransform
>>> base_env = GymWrapper(gymnasium.make("Pendulum-v1"))
>>> env = TransformedEnv(base_env, RenameTransform("reward", ("nested", "reward")))
>>> env.full_reward_spec
Composite(
    nested: Composite(
        reward: UnboundedContinuous(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=None, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

property full_reward_spec_unbatched: Composite¶: 返回環境的 reward 規範，就好像它沒有批處理維度一樣。

property full_state_spec: Composite¶

完整的 state 規範。

full_state_spec 是一個 Composite` 例項，其中包含所有狀態條目（即，非操作的輸入資料）。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.full_state_spec
Composite(
    state: Composite(
        pipeline_state: Composite(
            q: UnboundedContinuous(
                shape=torch.Size([15]),
                space=None,
                device=cpu,
                dtype=torch.float32,
                domain=continuous),
    [...], device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

property full_state_spec_unbatched: Composite¶: 返回環境的 state 規範，就好像它沒有批處理維度一樣。

get_buffer(target: str) → Tensor¶

返回由 target 給定的緩衝區（如果存在），否則丟擲錯誤。

有關此方法功能的更詳細解釋以及如何正確指定 target，請參閱 get_submodule 的文件字串。

引數:: target – 要查詢的 buffer 的完全限定字串名稱。（要指定完全限定字串，請參閱 get_submodule。）
返回:: 由 target 引用的緩衝區
返回型別:: torch.Tensor
丟擲:: AttributeError – 如果目標字串引用了無效路徑或解析為非 buffer 物件。

get_extra_state() → Any¶

返回要包含在模組 state_dict 中的任何額外狀態。

如果您的模組需要儲存額外狀態，請實現此方法和相應的 set_extra_state()。在構建模組的 state_dict() 時會呼叫此函式。

注意，為了保證 state_dict 的序列化工作正常，額外狀態應該是可被 pickle 的。我們僅為 Tensors 的序列化提供向後相容性保證；其他物件的序列化形式若發生變化，可能導致向後相容性中斷。

返回:: 要儲存在模組 state_dict 中的任何額外狀態
返回型別:: 物件

get_parameter(target: str) → Parameter¶

如果存在，返回由 target 給定的引數，否則丟擲錯誤。

有關此方法功能的更詳細解釋以及如何正確指定 target，請參閱 get_submodule 的文件字串。

引數:: target – 要查詢的 Parameter 的完全限定字串名稱。（要指定完全限定字串，請參閱 get_submodule。）
返回:: 由 target 引用的引數
返回型別:: torch.nn.Parameter
丟擲:: AttributeError – 如果目標字串引用了無效路徑或解析為非 nn.Parameter 的物件。

get_submodule(target: str) → Module¶

如果存在，返回由 target 給定的子模組，否則丟擲錯誤。

例如，假設您有一個 nn.Module A，它看起來像這樣

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(圖示了一個 nn.Module A。 A 包含一個巢狀子模組 net_b，該子模組本身有兩個子模組 net_c 和 linear。 net_c 隨後又有一個子模組 conv。)

要檢查是否存在 linear 子模組，可以呼叫 get_submodule("net_b.linear")。要檢查是否存在 conv 子模組，可以呼叫 get_submodule("net_b.net_c.conv")。

get_submodule 的執行時複雜度受 target 中模組巢狀深度的限制。與 named_modules 的查詢相比，後者的複雜度是按傳遞模組數量計算的 O(N)。因此，對於簡單地檢查某個子模組是否存在，應始終使用 get_submodule。

引數:: target – 要查詢的子模組的完全限定字串名稱。（要指定完全限定字串，請參閱上面的示例。）
返回:: 由 target 引用的子模組
返回型別:: torch.nn.Module
丟擲:: AttributeError – 如果在目標字串解析的任何路徑中，子路徑解析為不存在的屬性名或不是 nn.Module 例項的物件。

half() → Self¶

將所有浮點引數和緩衝區轉換為 half 資料型別。

注意

此方法就地修改模組。

返回:: self
返回型別:: 模組

property input_spec: TensorSpec¶

輸入規範。

包含輸入到環境的所有規範的複合規範。

它包含

“full_action_spec”: 輸入動作的規範
“full_state_spec”: 所有其他環境輸入的規範

此屬性是鎖定的，應該是隻讀的。相反，要設定其中包含的規範，請使用相應的屬性。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.input_spec
Composite(
    full_state_spec: None,
    full_action_spec: Composite(
        action: BoundedContinuous(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

property input_spec_unbatched: Composite¶: 返回環境的 input 規範，就好像它沒有批處理維度一樣。

ipu(device: Optional[Union[device, int]] = None) → Self¶

將所有模型引數和緩衝區移動到 IPU。

這也會使關聯的引數和緩衝區成為不同的物件。因此，如果模組在最佳化時將駐留在 IPU 上，則應在構建最佳化器之前呼叫它。

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數將複製到該裝置
返回:: self
返回型別:: 模組

property is_spec_locked¶

獲取環境的規範是否已鎖定。

此屬性可以直接修改。

返回:: 如果規範已鎖定，則為 True，否則為 False。
返回型別:: 布林值

另請參閱

鎖定環境 spec.

load_state_dict(state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False)¶

將 state_dict 中的引數和緩衝區複製到此模組及其子模組中。

如果 strict 為 True，則 state_dict 的鍵必須與此模組的 state_dict() 函式返回的鍵完全匹配。

警告

如果 assign 為 True，則必須在呼叫 load_state_dict 後建立最佳化器，除非 get_swap_module_params_on_conversion() 為 True。

引數:

state_dict (dict) – 包含引數和持久 buffer 的字典。
strict (bool, 可選) – 是否嚴格強制 state_dict 中的鍵與此模組的 state_dict() 函式返回的鍵匹配。預設為 True
assign (bool, optional) – 當設定為 False 時，將保留當前模組中張量的屬性；當設定為 True 時，將保留 state_dict 中張量的屬性。唯一的例外是 Parameter 的 requires_grad 欄位，此時將保留模組的值。預設值：False

返回:

missing_keys 是一個包含此模組期望但
在提供的 state_dict 中缺失的任何鍵的字串列表。
unexpected_keys 是一個字串列表，包含此模組
不期望但在提供的 state_dict 中存在的鍵。

返回型別:

NamedTuple，包含 missing_keys 和 unexpected_keys 欄位。

注意

如果引數或緩衝區註冊為 None 且其對應的鍵存在於 state_dict 中，load_state_dict() 將引發 RuntimeError。

maybe_reset(tensordict: TensorDictBase) → TensorDictBase¶

檢查輸入 tensordict 的 done 鍵，如果需要，則重置已完成的環境。

引數:: tensordict (TensorDictBase) – 來自 step_mdp() 輸出的 tensordict。
返回:: 一個與輸入相同的 tensordict，其中環境未被重置，並且在環境被重置的地方包含新的重置資料。

modules() → Iterator[Module]¶

返回網路中所有模組的迭代器。

產生:: Module – 網路中的一個模組

注意

重複的模組只返回一次。在以下示例中，l 只返回一次。

示例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device: Optional[Union[device, int]] = None) → Self¶

將所有模型引數和緩衝區移動到 MTIA。

這也會使關聯的引數和緩衝區成為不同的物件。因此，如果模組在最佳化時將駐留在 MTIA 上，則應在構建最佳化器之前呼叫它。

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數將複製到該裝置
返回:: self
返回型別:: 模組

named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[tuple[str, torch.Tensor]]¶

返回模組緩衝區上的迭代器，同時生成緩衝區的名稱和緩衝區本身。

引數:

prefix (str) – 為所有 buffer 名稱新增字首。
recurse (bool, optional) – 如果為 True，則會生成此模組及其所有子模組的 buffers。否則，僅生成此模組直接成員的 buffers。預設為 True。
remove_duplicate (bool, optional) – 是否在結果中刪除重複的 buffers。預設為 True。

產生:

(str, torch.Tensor) – 包含名稱和緩衝區的元組

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children() → Iterator[tuple[str, 'Module']]¶

返回對直接子模組的迭代器，生成模組的名稱和模組本身。

產生:: (str, Module) – 包含名稱和子模組的元組

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo: Optional[set['Module']] = None, prefix: str = '', remove_duplicate: bool = True)¶

返回網路中所有模組的迭代器，同時生成模組的名稱和模組本身。

引數:

memo – 用於儲存已新增到結果中的模組集合的 memo
prefix – 將新增到模組名稱的名稱字首
remove_duplicate – 是否從結果中刪除重複的模組例項

產生:

(str, Module) – 名稱和模組的元組

注意

重複的模組只返回一次。在以下示例中，l 只返回一次。

示例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[tuple[str, torch.nn.parameter.Parameter]]¶

返回模組引數的迭代器，同時生成引數的名稱和引數本身。

引數:

prefix (str) – 為所有引數名稱新增字首。
recurse (bool) – 如果為 True，則會生成此模組及其所有子模組的引數。否則，僅生成此模組直接成員的引數。
remove_duplicate (bool, optional) – 是否在結果中刪除重複的引數。預設為 True。

產生:

(str, Parameter) – 包含名稱和引數的元組

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

property observation_keys: list[tensordict._nestedkey.NestedKey]¶

環境的 observation keys。

預設情況下，只有一個名為“observation”的 key。

鍵按資料樹的深度排序。

property observation_spec: Composite¶

Observation spec。

必須是 torchrl.data.Composite 例項。spec 中列出的鍵在重置和步進後可直接訪問。

在 TorchRL 中，即使它們不嚴格來說是“observation”，所有 info、state、transforms 的結果等環境輸出都儲存在 observation_spec 中。

因此，"observation_spec" 應被視為環境輸出（非 done 或 reward 資料）的通用資料容器。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.observation_spec
Composite(
    observation: BoundedContinuous(
        shape=torch.Size([3]),
        space=ContinuousBox(
            low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
            high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
        device=cpu,
        dtype=torch.float32,
        domain=continuous), device=cpu, shape=torch.Size([]))

property observation_spec_unbatched: Composite¶: 返回環境的 observation 規範，就好像它沒有批處理維度一樣。

property output_spec: TensorSpec¶

Output spec。

包含環境所有資料輸出 spec 的複合 spec。

它包含

“full_reward_spec”: reward 的 spec
“full_done_spec”: done 的 spec
“full_observation_spec”: 所有其他環境輸出的 spec

此屬性是鎖定的，應該是隻讀的。相反，要設定其中包含的規範，請使用相應的屬性。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.output_spec
Composite(
    full_reward_spec: Composite(
        reward: UnboundedContinuous(
            shape=torch.Size([1]),
            space=None,
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])),
    full_observation_spec: Composite(
        observation: BoundedContinuous(
            shape=torch.Size([3]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])),
    full_done_spec: Composite(
        done: Categorical(
            shape=torch.Size([1]),
            space=DiscreteBox(n=2),
            device=cpu,
            dtype=torch.bool,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

property output_spec_unbatched: Composite¶: 返回環境的 output spec，就好像它沒有 batch 維度一樣。

parameters(recurse: bool = True) → Iterator[Parameter]¶

返回模組引數的迭代器。

這通常傳遞給最佳化器。

引數:: recurse (bool) – 如果為 True，則會生成此模組及其所有子模組的引數。否則，僅生成此模組直接成員的引數。
產生:: Parameter – 模組引數

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

rand_action(tensordict: TensorDictBase | None = None)¶

根據 action_spec 屬性執行隨機動作。

引數:: tensordict (TensorDictBase, optional) – 要將生成的動作寫入的 tensordict。
返回:: 一個 tensordict 物件，其“action”條目已用從 action-spec 中隨機抽取的樣本更新。

rand_step(tensordict: TensorDictBase | None = None) → TensorDictBase¶

根據 action_spec 屬性在環境中執行隨機步長。

引數:: tensordict (TensorDictBase, optional) – 要將生成的 info 寫入的 tensordict。
返回:: 一個 tensordict 物件，其中包含在環境中隨機步長後的新 observation。動作將以“action”鍵儲存。

register_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]]) → RemovableHandle¶

在模組上註冊一個反向傳播鉤子。

此函式已棄用，建議使用 register_full_backward_hook()，並且此函式在未來版本中的行為將發生變化。

返回:: 一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子
返回型別:: torch.utils.hooks.RemovableHandle

register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) → None¶

向模組新增一個緩衝區。

這通常用於註冊不應被視為模型引數的緩衝區。例如，BatchNorm 的 running_mean 不是引數，但屬於模組的狀態。預設情況下，緩衝區是持久的，將與引數一起儲存。可以透過將 persistent 設定為 False 來更改此行為。持久緩衝區和非持久緩衝區之間的唯一區別是後者不會成為此模組 state_dict 的一部分。

可以使用給定名稱作為屬性訪問緩衝區。

引數:

name (str) – buffer 的名稱。可以使用給定的名稱從此模組訪問 buffer
tensor (Tensor 或 None) – 要註冊的緩衝區。如果為 None，則在緩衝區上執行的操作（如 cuda）將被忽略。如果為 None，則緩衝區 **不** 包含在模組的 state_dict 中。
persistent (bool) – 緩衝區是否是此模組 state_dict 的一部分。

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_collector(collector: DataCollectorBase)¶

將 collector 註冊到 environment。

引數:: collector (DataCollectorBase) – 要註冊的 collector。

register_forward_hook(hook: Union[Callable[[T, tuple[Any, ...], Any], Optional[Any]], Callable[[T, tuple[Any, ...], dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False) → RemovableHandle¶

在模組上註冊一個前向鉤子。

在每次呼叫 forward() 計算輸出後，都會呼叫 hook。

如果 with_kwargs 為 False 或未指定，則輸入僅包含傳遞給模組的位置引數。關鍵字引數不會傳遞給 hook，只傳遞給 forward。hook 可以修改輸出。它可以就地修改輸入，但這不會影響 forward，因為它是在 forward() 呼叫之後呼叫的。hook 的簽名應如下所示

hook(module, args, output) -> None or modified output

如果 with_kwargs 為 True，則前向鉤子將接收傳遞給 forward 函式的 kwargs，並需要返回可能已修改的輸出。鉤子應該具有以下簽名

hook(module, args, kwargs, output) -> None or modified output

引數:

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 True，則提供的 hook 將在當前 torch.nn.Module 的所有現有 forward hook 之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 forward hook 之後觸發。請注意，使用 register_module_forward_hook() 註冊的全域性 forward hook 將在由此方法註冊的所有 hook 之前觸發。預設為 False
with_kwargs (bool) – 如果為 True，則 hook 將接收傳遞給 forward 函式的 kwargs。預設為 False。
always_call (bool) – 如果為 True，則無論在呼叫 Module 時是否引發異常，都會執行 hook。預設為 False。

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook: Union[Callable[[T, tuple[Any, ...]], Optional[Any]], Callable[[T, tuple[Any, ...], dict[str, Any]], Optional[tuple[Any, dict[str, Any]]]]], *, prepend: bool = False, with_kwargs: bool = False) → RemovableHandle¶

在模組上註冊一個前向預鉤子。

每次呼叫 forward() 之前都會呼叫此鉤子。

如果 with_kwargs 為 false 或未指定，則輸入僅包含傳遞給模組的位置引數。關鍵字引數不會傳遞給鉤子，而只會傳遞給 forward。鉤子可以修改輸入。使用者可以返回一個元組或單個修改後的值。我們將把值包裝成一個元組，如果返回的是單個值（除非該值本身就是元組）。鉤子應該具有以下簽名

hook(module, args) -> None or modified input

如果 with_kwargs 為 true，則前向預鉤子將接收傳遞給 forward 函式的 kwargs。如果鉤子修改了輸入，則應該返回 args 和 kwargs。鉤子應該具有以下簽名

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

引數:

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 true，則提供的 hook 將在當前 torch.nn.Module 的所有現有 forward_pre 鉤子之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 forward_pre 鉤子之後觸發。請注意，使用 register_module_forward_pre_hook() 註冊的全域性 forward_pre 鉤子將在使用此方法註冊的所有鉤子之前觸發。預設為 False。
with_kwargs (bool) – 如果為 True，則 hook 將接收傳遞給 forward 函式的 kwargs。預設為 False。

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor]], Union[tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模組上註冊一個反向傳播鉤子。

每次計算相對於模組的梯度時，將呼叫此鉤子，其觸發規則如下：

通常，鉤子在計算相對於模組輸入的梯度時觸發。

如果模組輸入都不需要梯度，則在計算相對於模組輸出的梯度時觸發鉤子。

如果模組輸出都不需要梯度，則鉤子將不觸發。

鉤子應具有以下簽名

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

grad_input 和 grad_output 是包含關於輸入和輸出的梯度的元組。鉤子不應修改其引數，但可以選擇性地返回一個關於輸入的新的梯度，該梯度將用於替換後續計算中的 grad_input。 grad_input 將僅對應於作為位置引數給出的輸入，並且所有關鍵字引數都將被忽略。對於所有非 Tensor 引數，grad_input 和 grad_output 中的條目將為 None。

由於技術原因，當此鉤子應用於模組時，其前向函式將接收傳遞給模組的每個張量的檢視。類似地，呼叫者將接收模組前向函式返回的每個張量的檢視。

警告

使用反向傳播鉤子時不允許就地修改輸入或輸出，否則將引發錯誤。

引數:

hook (Callable) – 要註冊的使用者定義鉤子。
prepend (bool) – 如果為 true，則提供的 hook 將在當前 torch.nn.Module 的所有現有 backward 鉤子之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 backward 鉤子之後觸發。請注意，使用 register_module_full_backward_hook() 註冊的全域性 backward 鉤子將在使用此方法註冊的所有鉤子之前觸發。

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模組上註冊一個反向預鉤子。

每次計算模組的梯度時，將呼叫此鉤子。鉤子應具有以下簽名

hook(module, grad_output) -> tuple[Tensor] or None

grad_output 是一個元組。鉤子不應修改其引數，但可以選擇返回一個新的輸出梯度，該梯度將取代 grad_output 用於後續計算。對於所有非 Tensor 引數，grad_output 中的條目將為 None。

由於技術原因，當此鉤子應用於模組時，其前向函式將接收傳遞給模組的每個張量的檢視。類似地，呼叫者將接收模組前向函式返回的每個張量的檢視。

警告

使用反向傳播鉤子時不允許就地修改輸入，否則將引發錯誤。

引數:

hook (Callable) – 要註冊的使用者定義鉤子。
prepend (bool) – 如果為 true，則提供的 hook 將在當前 torch.nn.Module 的所有現有 backward_pre 鉤子之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 backward_pre 鉤子之後觸發。請注意，使用 register_module_full_backward_pre_hook() 註冊的全域性 backward_pre 鉤子將在使用此方法註冊的所有鉤子之前觸發。

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

classmethod register_gym(id: str, *, entry_point: Callable | None = None, transform: Transform | None = None, info_keys: list[NestedKey] | None = None, backend: str = None, to_numpy: bool = False, reward_threshold: float | None] = None, nondeterministic: bool = False, max_episode_steps: int | None = None, order_enforce: bool = True, autoreset: bool | None = None, disable_env_checker: bool = False, apply_api_compatibility: bool = False, **kwargs)¶

註冊一個 gym(nasium) 環境。

此方法的設計考慮了以下範圍：

將 TorchRL-first 環境納入使用 Gym 的框架；
將其他環境（例如 DeepMind Control、Brax、Jumanji 等）納入使用 Gym 的框架。

引數:

id (str) – 環境的名稱。應遵循 gym 命名約定。

關鍵字引數:

entry_point (callable, optional) –

用於構建環境的入口點。如果未傳入，則父類將用作入口點。通常，這用於註冊不一定繼承自正在使用的基類的環境。

>>> from torchrl.envs import DMControlEnv
>>> DMControlEnv.register_gym("DMC-cheetah-v0", env_name="cheetah", task="run")
>>> # equivalently
>>> EnvBase.register_gym("DMC-cheetah-v0", entry_point=DMControlEnv, env_name="cheetah", task="run")

transform (torchrl.envs.Transform) – 要與 env 一起使用的變換（或 torchrl.envs.Compose 例項內的變換列表）。此引數可以在呼叫 make() 時傳遞（請參閱下面的示例）。
info_keys (List[NestedKey], optional) –
如果提供了 info_keys，它們將用於構建 info 字典，並將從 observation keys 中排除。此引數可以在呼叫 make() 時傳遞（請參閱下面的示例）。

警告

使用 info_keys 可能會導致 spec 為空，因為內容已被移至 info 字典。Gym 不喜歡 spec 中為空的 Dict，因此應使用 RemoveEmptySpecs 移除此空內容。
backend (str, optional) – 後端。可以是 “gym” 或 “gymnasium” 或與 set_gym_backend 相容的任何其他後端。
to_numpy (bool, optional) – 如果為 True，則呼叫 step 和 reset 的結果將對映到 numpy 陣列。預設為 False（結果為 tensor）。此引數可以在呼叫 make() 時傳遞（請參閱下面的示例）。
reward_threshold (float, optional) – [Gym kwarg] considered learned an environment.
nondeterministic (bool, optional) – [Gym kwarg] 如果環境是不確定的（即使已知初始種子和所有操作）。預設為 False。
max_episode_steps (int, optional) – [Gym kwarg] 截斷之前的最大單集步數。由 Time Limit 包裝器使用。
order_enforce (bool, optional) – [Gym >= 0.14] 是否應用 order enforcer wrapper 以確保使用者按正確順序執行函式。預設為 True。
autoreset (bool, optional) – [Gym >= 0.14 and <1.0.0] 是否新增 autoreset wrapper，以便不需要呼叫 reset。預設為 False。
disable_env_checker – [Gym >= 0.14] 是否停用環境的環境檢查器。預設為 False。
apply_api_compatibility – [Gym >= 0.26 and <1.0.0] 是否應用 StepAPICompatibility 包裝器。預設為 False。
**kwargs – 傳遞給環境建構函式的任意關鍵字引數。

注意

TorchRL 的環境沒有 "info" 字典的概念，因為 TensorDict 提供了大多數訓練設定所需的所有儲存要求。儘管如此，您可以使用 info_keys 引數來精細控制哪些內容被視為 observation，哪些被視為 info。

示例

>>> # Register the "cheetah" env from DMControl with the "run" task
>>> from torchrl.envs import DMControlEnv
>>> import torch
>>> DMControlEnv.register_gym("DMC-cheetah-v0", to_numpy=False, backend="gym", env_name="cheetah", task_name="run")
>>> import gym
>>> envgym = gym.make("DMC-cheetah-v0")
>>> envgym.seed(0)
>>> torch.manual_seed(0)
>>> envgym.reset()
({'position': tensor([-0.0855,  0.0215, -0.0881, -0.0412, -0.1101,  0.0080,  0.0254,  0.0424],
       dtype=torch.float64), 'velocity': tensor([ 1.9609e-02, -1.9776e-04, -1.6347e-03,  3.3842e-02,  2.5338e-02,
         3.3064e-02,  1.0381e-04,  7.6656e-05,  1.0204e-02],
       dtype=torch.float64)}, {})
>>> envgym.step(envgym.action_space.sample())
({'position': tensor([-0.0833,  0.0275, -0.0612, -0.0770, -0.1256,  0.0082,  0.0186,  0.0476],
       dtype=torch.float64), 'velocity': tensor([ 0.2221,  0.2256,  0.5930,  2.6937, -3.5865, -1.5479,  0.0187, -0.6825,
         0.5224], dtype=torch.float64)}, tensor([0.0018], dtype=torch.float64), tensor([False]), tensor([False]), {})
>>> # same environment with observation stacked
>>> from torchrl.envs import CatTensors
>>> envgym = gym.make("DMC-cheetah-v0", transform=CatTensors(in_keys=["position", "velocity"], out_key="observation"))
>>> envgym.reset()
({'observation': tensor([-0.1005,  0.0335, -0.0268,  0.0133, -0.0627,  0.0074, -0.0488, -0.0353,
        -0.0075, -0.0069,  0.0098, -0.0058,  0.0033, -0.0157, -0.0004, -0.0381,
        -0.0452], dtype=torch.float64)}, {})
>>> # same environment with numpy observations
>>> envgym = gym.make("DMC-cheetah-v0", transform=CatTensors(in_keys=["position", "velocity"], out_key="observation"), to_numpy=True)
>>> envgym.reset()
({'observation': array([-0.11355747,  0.04257728,  0.00408397,  0.04155852, -0.0389733 ,
       -0.01409826, -0.0978704 , -0.08808327,  0.03970837,  0.00535434,
       -0.02353762,  0.05116226,  0.02788907,  0.06848346,  0.05154399,
        0.0371798 ,  0.05128025])}, {})
>>> # If gymnasium is installed, we can register the environment there too.
>>> DMControlEnv.register_gym("DMC-cheetah-v0", to_numpy=False, backend="gymnasium", env_name="cheetah", task_name="run")
>>> import gymnasium
>>> envgym = gymnasium.make("DMC-cheetah-v0")
>>> envgym.seed(0)
>>> torch.manual_seed(0)
>>> envgym.reset()
({'position': tensor([-0.0855,  0.0215, -0.0881, -0.0412, -0.1101,  0.0080,  0.0254,  0.0424],
       dtype=torch.float64), 'velocity': tensor([ 1.9609e-02, -1.9776e-04, -1.6347e-03,  3.3842e-02,  2.5338e-02,
         3.3064e-02,  1.0381e-04,  7.6656e-05,  1.0204e-02],
       dtype=torch.float64)}, {})

注意

此功能也適用於無狀態環境（例如 BraxEnv）。

>>> import gymnasium
>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.envs import BraxEnv, SelectTransform
>>>
>>> # get action for dydactic purposes
>>> env = BraxEnv("ant", batch_size=[2])
>>> env.set_seed(0)
>>> torch.manual_seed(0)
>>> td = env.rollout(10)
>>>
>>> actions = td.get("action")
>>>
>>> # register env
>>> env.register_gym("Brax-Ant-v0", env_name="ant", batch_size=[2], info_keys=["state"])
>>> gym_env = gymnasium.make("Brax-Ant-v0")
>>> gym_env.seed(0)
>>> torch.manual_seed(0)
>>>
>>> gym_env.reset()
>>> obs = []
>>> for i in range(10):
...     obs, reward, terminated, truncated, info = gym_env.step(td[..., i].get("action"))

register_load_state_dict_post_hook(hook)¶

註冊一個後鉤子，用於在模組的 load_state_dict() 被呼叫後執行。

它應該具有以下簽名：: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

如果需要，可以就地修改給定的 incompatible_keys。

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

返回:: 一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子
返回型別:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)¶

註冊一個預鉤子，用於在模組的 load_state_dict() 被呼叫之前執行。

它應該具有以下簽名：: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

引數:: hook (Callable) – 在載入狀態字典之前將呼叫的可呼叫鉤子。

register_module(name: str, module: Optional[Module]) → None¶: Alias for add_module().

register_parameter(name: str, param: Optional[Parameter]) → None¶

向模組新增一個引數。

可以使用給定名稱作為屬性訪問該引數。

引數:

name (str) – 引數的名稱。可以透過給定名稱從該模組訪問該引數。
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

register_state_dict_post_hook(hook)¶

註冊 state_dict() 方法的後置鉤子。

它應該具有以下簽名：: hook(module, state_dict, prefix, local_metadata) -> None

註冊的鉤子可以就地修改 state_dict。

register_state_dict_pre_hook(hook)¶

註冊 state_dict() 方法的前置鉤子。

它應該具有以下簽名：: hook(module, prefix, keep_vars) -> None

註冊的鉤子可用於在進行 state_dict 呼叫之前執行預處理。

requires_grad_(requires_grad: bool = True) → Self¶

更改自動梯度是否應記錄此模組中引數的操作。

此方法就地設定引數的 requires_grad 屬性。

此方法有助於凍結模組的一部分以進行微調或單獨訓練模型的一部分（例如，GAN 訓練）。

請參閱本地停用梯度計算以比較 .requires_grad_() 和幾種可能與之混淆的類似機制。

引數:: requires_grad (bool) – 自動求導是否應記錄此模組上的引數操作。預設為 True。
返回:: self
返回型別:: 模組

reset(tensordict: TensorDictBase | None = None, *, **kwargs) → TensorDictBase¶

重置環境。

與 step 和 _step 一樣，只有私有方法 _reset 應該被 EnvBase 子類覆蓋。

引數:

tensordict (TensorDictBase, optional) – 用於包含新 observation 的 tensordict。在某些情況下，此輸入還可用於向 reset 函式傳遞引數。
kwargs (optional) – 傳遞給原生 reset 函式的其他引數。

返回:

一個 tensordict（或任何輸入的 tensordict），原地修改以包含相應的 observation。

注意

reset 不應被 EnvBase 子類覆蓋。應修改的方法是 _reset()。

property reset_keys: list[tensordict._nestedkey.NestedKey]¶

返回重置鍵列表。

Reset keys 是指示部分重置的鍵，用於批次、多工或多代理設定。它們的結構是 (*prefix, "_reset")，其中 prefix 是一個（可能為空的）字串元組，指向 tensordict 中的一個位置，在該位置可以找到 done 狀態。

鍵按資料樹的深度排序。

property reward_key¶

環境的獎勵鍵。

預設情況下，這將是“reward”。

如果環境中存在多個獎勵鍵，此函式將引發異常。

property reward_keys: list[tensordict._nestedkey.NestedKey]¶

環境的獎勵鍵。

預設情況下，只有一個鍵，名為“reward”。

鍵按資料樹的深度排序。

property reward_spec: TensorSpec¶

reward spec。

reward_spec 始終儲存為複合 spec。

如果 reward spec 作為簡單 spec 提供，則返回該 spec。

>>> env.reward_spec = Unbounded(1)
>>> env.reward_spec
UnboundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

如果 reward spec 作為複合 spec 提供並且僅包含一個葉子，則此函式將僅返回該葉子。

>>> env.reward_spec = Composite({"nested": {"reward": Unbounded(1)}})
>>> env.reward_spec
UnboundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

如果 reward spec 作為複合 spec 提供並且具有多個葉子，則此函式將返回整個 spec。

>>> env.reward_spec = Composite({"nested": {"reward": Unbounded(1), "another_reward": Categorical(1)}})
>>> env.reward_spec
Composite(
    nested: Composite(
        reward: UnboundedContinuous(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        another_reward: Categorical(
            shape=torch.Size([]),
            space=DiscreteBox(n=1),
            device=cpu,
            dtype=torch.int64,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

要檢索傳遞的完整 spec，請使用

>>> env.output_spec["full_reward_spec"]

此屬性是可變的。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.reward_spec
UnboundedContinuous(
    shape=torch.Size([1]),
    space=None,
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

property reward_spec_unbatched: TensorSpec¶: 返回環境的 reward 規範，就好像它沒有批處理維度一樣。

rollout(max_steps: int, policy: Callable[[TensorDictBase], TensorDictBase] | None = None, callback: Callable[[TensorDictBase, ...], Any] | None = None, *, auto_reset: bool = True, auto_cast_to_device: bool = False, break_when_any_done: bool | None = None, break_when_all_done: bool | None = None, return_contiguous: bool | None = False, tensordict: TensorDictBase | None = None, set_truncated: bool = False, out=None, trust_policy: bool = False) → TensorDictBase¶

在環境中執行 rollout。

只要其中任何一個環境達到任何一個完成狀態，該函式就會返回。

引數:

max_steps (int) – 要執行的最大步數。實際步數可能較少，如果環境在 max_steps 執行完畢之前達到完成狀態。
policy (callable, optional) – 用於計算所需動作的可呼叫函式。如果未提供策略，將使用 env.rand_step() 呼叫動作。該策略可以是任何讀取 tensordict 或所有觀察條目的序列的可呼叫函式，這些條目 **按** env.observation_spec.keys() 的順序排序。預設為 None。
callback (Callable[[TensorDict], Any], optional) – 在每次迭代時使用給定的 TensorDict 呼叫該函式。預設為 None。callback 的輸出不會被收集，使用者有責任在 callback 呼叫中儲存任何資料，以便資料能在呼叫 rollout 之外保留。

關鍵字引數:

auto_reset (bool, optional) – 如果為 True，則將在開始 rollout 之前重置包含的環境。如果為 False，則 rollout 將從先前狀態繼續，這需要 tensordict 引數與先前 rollout 一起傳遞。預設為 True。
auto_cast_to_device (bool, optional) – 如果為 True，則 tensordict 的裝置將在使用策略之前自動轉換為策略的裝置。預設為 False。
break_when_any_done (bool) –
如果為 True，則當任何包含的環境達到任何 done 狀態時停止。如果為 False，則 done 的環境會自動重置。預設為 True。

另請參閱

文件的部分重置提供了更多關於部分重置的資訊。
break_when_all_done (bool, optional) –
如果為 True，則在所有包含的環境都達到任何 done 狀態時停止。如果為 False，則在至少一個環境達到任何 done 狀態時停止。預設為 False。

另請參閱

文件的部分步進提供了更多關於部分重置的資訊。
return_contiguous (bool) – If False，將返回一個 LazyStackedTensorDict。如果環境沒有動態規範，則預設為 True，否則為 False。
tensordict (TensorDict, optional) – 如果 auto_reset 為 False，則必須提供初始 tensordict。Rollout 將檢查此 tensordict 是否具有 done 標誌，並在必要時重置該維度中的環境。如果 tensordict 是 reset 的輸出，這通常不應該發生，但如果 tensordict 是前一個 rollout 的最後一步，則可能會發生。當 auto_reset=True 時，也可以提供 tensordict，以便將元資料傳遞給 reset 方法，例如批次大小或無狀態環境的裝置。
set_truncated (bool, optional) – 如果為 True，則在 rollout 完成後，將 "truncated" 和 "done" 鍵設定為 True。如果在 done_spec 中未找到 "truncated"，則會引發異常。截斷鍵可以透過 env.add_truncated_keys 設定。預設為 False。
trust_policy (bool, optional) – 如果為 True，則將信任非 TensorDictModule 策略，並假定其與收集器相容。對於 CudaGraphModules，這預設為 True，否則為 False。

返回:

包含結果軌跡的 TensorDict 物件。

返回的資料將在 tensordict 的最後一個維度（在 env.ndim 索引處）用“time”維度名稱進行標記。

rollout 對於顯示環境的資料結構非常有幫助。

示例

>>> # Using rollout without a policy
>>> from torchrl.envs.libs.gym import GymEnv
>>> from torchrl.envs.transforms import TransformedEnv, StepCounter
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), StepCounter(max_steps=20))
>>> rollout = env.rollout(max_steps=1000)
>>> print(rollout)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                truncated: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([20]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        truncated: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([20]),
    device=cpu,
    is_shared=False)
>>> print(rollout.names)
['time']
>>> # with envs that contain more dimensions
>>> from torchrl.envs import SerialEnv
>>> env = SerialEnv(3, lambda: TransformedEnv(GymEnv("Pendulum-v1"), StepCounter(max_steps=20)))
>>> rollout = env.rollout(max_steps=1000)
>>> print(rollout)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                truncated: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([3, 20]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        truncated: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([3, 20]),
    device=cpu,
    is_shared=False)
>>> print(rollout.names)
[None, 'time']

使用策略（普通的 Module 或 TensorDictModule）也很容易。

示例

>>> from torch import nn
>>> env = GymEnv("CartPole-v1", categorical_action_encoding=True)
>>> class ArgMaxModule(nn.Module):
...     def forward(self, values):
...         return values.argmax(-1)
>>> n_obs = env.observation_spec["observation"].shape[-1]
>>> n_act = env.action_spec.n
>>> # A deterministic policy
>>> policy = nn.Sequential(
...     nn.Linear(n_obs, n_act),
...     ArgMaxModule())
>>> env.rollout(max_steps=10, policy=policy)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)
>>> # Under the hood, rollout will wrap the policy in a TensorDictModule
>>> # To speed things up we can do that ourselves
>>> from tensordict.nn import TensorDictModule
>>> policy = TensorDictModule(policy, in_keys=list(env.observation_spec.keys()), out_keys=["action"])
>>> env.rollout(max_steps=10, policy=policy)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)

在某些情況下，無法獲得連續的 tensordict，因為它們無法堆疊。當每步返回的資料可能具有不同的形狀，或者當不同環境一起執行時，可能會發生這種情況。在這種情況下，return_contiguous=False 將導致返回的 tensordict 成為 tensordict 的惰性堆疊。

非連續 rollout 的示例

>>> rollout = env.rollout(4, return_contiguous=False)
>>> print(rollout)
LazyStackedTensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: LazyStackedTensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                truncated: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        truncated: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=cpu,
    is_shared=False)
    >>> print(rollout.names)
    [None, 'time']

Rollout 可以在迴圈中使用以模擬資料收集。要做到這一點，您需要將上一個 rollout 的最後一個 tensordict 作為輸入，並在其上呼叫 step_mdp()。

資料收集 rollout 的示例

>>> from torchrl.envs import GymEnv, step_mdp
>>> env = GymEnv("CartPole-v1")
>>> epochs = 10
>>> input_td = env.reset()
>>> for i in range(epochs):
...     rollout_td = env.rollout(
...         max_steps=100,
...         break_when_any_done=False,
...         auto_reset=False,
...         tensordict=input_td,
...     )
...     input_td = step_mdp(
...         rollout_td[..., -1],
...     )

set_extra_state(state: Any) → None¶

設定載入的 state_dict 中包含的額外狀態。

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

引數:: state (dict) – 來自 state_dict 的額外狀態

set_seed(seed: int | None = None, static_seed: bool = False) → int | None¶

設定環境的種子，並返回要使用的下一個種子（如果存在單個環境，則為輸入種子）。

引數:

seed (int) – 要設定的種子。種子僅在環境中本地設定。要處理全域性種子，請參閱 manual_seed()。
static_seed (bool, optional) – 如果 True，種子不會遞增。預設為 False

返回:

即，如果同時建立此環境，則應為另一個環境使用的種子。

返回型別:

代表“下一個種子”的整數

set_spec_lock_(mode: bool = True) → EnvBase¶

鎖定或解鎖環境的規範。

引數:: mode (bool) – 是否鎖定（True）或解鎖（False）規範。預設為 True。
返回:: 環境例項本身。
返回型別:: EnvBase

另請參閱

鎖定環境 spec.

set_submodule(target: str, module: Module, strict: bool = False) → None¶

如果存在，設定由 target 給定的子模組，否則丟擲錯誤。

注意

如果 strict 設定為 False（預設），該方法將替換現有子模組或在父模組存在的情況下建立新子模組。如果 strict 設定為 True，該方法將僅嘗試替換現有子模組，並在子模組不存在時引發錯誤。

例如，假設您有一個 nn.Module A，它看起來像這樣

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(圖示了一個 nn.Module A。 A 包含一個巢狀子模組 net_b，該子模組本身有兩個子模組 net_c 和 linear。 net_c 隨後又有一個子模組 conv。)

要用一個新的 Linear 子模組覆蓋 Conv2d，可以呼叫 set_submodule("net_b.net_c.conv", nn.Linear(1, 1))，其中 strict 可以是 True 或 False。

要將一個新的 Conv2d 子模組新增到現有的 net_b 模組中，可以呼叫 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1))。

在上面，如果設定 strict=True 並呼叫 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True)，則會引發 AttributeError，因為 net_b 中不存在名為 conv 的子模組。

引數:

target – 要查詢的子模組的完全限定字串名稱。（要指定完全限定字串，請參閱上面的示例。）
module – 要設定子模組的物件。
strict – 如果為 False，該方法將替換現有子模組或建立新子模組（如果父模組存在）。如果為 True，則該方法只會嘗試替換現有子模組，如果子模組不存在則丟擲錯誤。

丟擲:

ValueError – 如果 target 字串為空或 module 不是 nn.Module 的例項。
AttributeError – 如果 target 字串路徑中的任何一點解析為一個不存在的屬性名或不是 nn.Module 例項的物件。

property shape¶: Equivalent to batch_size.

share_memory() → Self¶: 請參閱 torch.Tensor.share_memory_()。

property specs: Composite¶

返回一個 Composite 容器，其中包含所有環境。

此功能允許建立環境，在一個數據容器中檢索所有規範，然後從工作區中刪除該環境。

state_dict(*args, destination=None, prefix='', keep_vars=False)¶

返回一個字典，其中包含對模組整個狀態的引用。

引數和持久緩衝區（例如，執行平均值）都包含在內。鍵是相應的引數和緩衝區名稱。設定為 None 的引數和緩衝區不包含在內。

注意

返回的物件是淺複製。它包含對模組引數和緩衝區的引用。

警告

當前 state_dict() 還接受 destination、prefix 和 keep_vars 的位置引數，順序為。但是，這正在被棄用，並且在未來的版本中將強制使用關鍵字引數。

警告

請避免使用引數 destination，因為它不是為終端使用者設計的。

引數:

destination (dict, optional) – 如果提供，模組的狀態將更新到 dict 中，並返回相同的物件。否則，將建立一個 OrderedDict 並返回。預設為 None。
prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''。
keep_vars (bool, optional) – 預設情況下，state dict 中返回的 Tensors 會從 autograd 中分離。如果設定為 True，則不會執行分離。預設為 False。

返回:

包含模組整體狀態的字典

返回型別:

dict

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

property state_keys: list[tensordict._nestedkey.NestedKey]¶

環境的狀態鍵。

預設情況下，只有一個名為“state”的鍵。

鍵按資料樹的深度排序。

property state_spec: Composite¶

狀態規範。

必須是 torchrl.data.Composite 例項。此處列出的鍵應與操作一起作為輸入提供給環境。

在 TorchRL 中，即使它們不是嚴格意義上的“狀態”，所有不是動作的環境輸入都儲存在 state_spec 中。

因此，"state_spec" 應被視為非動作資料的通用環境輸入資料容器。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.state_spec
Composite(
    state: Composite(
        pipeline_state: Composite(
            q: UnboundedContinuous(
                shape=torch.Size([15]),
                space=None,
                device=cpu,
                dtype=torch.float32,
                domain=continuous),
    [...], device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

property state_spec_unbatched: TensorSpec¶: 返回環境的 state 規範，就好像它沒有批處理維度一樣。

step(tensordict: TensorDictBase) → TensorDictBase¶

在環境中執行一步。

Step 接受一個引數 tensordict，它通常包含一個‘action’鍵，表示要執行的操作。Step 將呼叫一個非原地私有方法 _step，這是 EnvBase 子類需要重寫的方法。

引數:: tensordict (TensorDictBase) – 包含要執行的操作的 tensordict。如果輸入 tensordict 包含 "next" 條目，則其中包含的值將優先於新計算的值。這提供了一種覆蓋底層計算的機制。
返回:: 輸入 tensordict，原地修改，包含結果觀察、完成狀態和獎勵（+其他如有需要）。

step_and_maybe_reset(tensordict: TensorDictBase) → tuple[tensordict.base.TensorDictBase, tensordict.base.TensorDictBase]¶

在環境中執行一步，並在需要時（部分）重置它。

引數:: tensordict (TensorDictBase) – an input data structure for the step() method.

此方法允許輕鬆編寫非停止的 rollout 函式。

示例

>>> from torchrl.envs import ParallelEnv, GymEnv
>>> def rollout(env, n):
...     data_ = env.reset()
...     result = []
...     for i in range(n):
...         data, data_ = env.step_and_maybe_reset(data_)
...         result.append(data)
...     return torch.stack(result)
>>> env = ParallelEnv(2, lambda: GymEnv("CartPole-v1"))
>>> print(rollout(env, 2))
TensorDict(
    fields={
        done: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([2, 2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([2, 2]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([2, 2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([2, 2]),
    device=cpu,
    is_shared=False)

step_mdp(next_tensordict: TensorDictBase) → TensorDictBase¶

使用提供的 next_tensordict 將環境狀態推進一步。

此方法透過從當前狀態轉換到 next_tensordict 定義的下一個狀態來更新環境狀態。生成的 tensordict 包含更新後的觀測以及任何其他相關狀態資訊，鍵的管理遵迴圈境的規範。

內部，此方法利用預先計算的 _StepMDP 例項來高效處理狀態、觀測、動作、獎勵和完成（done）鍵的轉換。_StepMDP 類透過預先計算要包含和排除的鍵來最佳化過程，從而減少重複呼叫期間的執行時開銷。_StepMDP 例項使用 exclude_action=False 建立，這意味著動作鍵將保留在根 tensordict 中。

引數:: next_tensordict (TensorDictBase) – A tensordict containing the state of the environment at the next time step. This tensordict should include keys for observations, actions, rewards, and done flags, as defined by the environment’s specifications. – 包含下一時間步環境狀態的 tensordict。此 tensordict 應包含根據環境規範定義的觀測、動作、獎勵和完成（done）標誌的鍵。
返回:: 一個代表環境狀態推進一步後的新 tensordict。
返回型別:: TensorDictBase

注意

該方法確保環境的鍵規範已針對提供的 next_tensordict 進行驗證，如果發現不匹配，則會發出警告。

注意

此方法旨在與具有一致鍵規範的環境高效工作，並利用 _StepMDP 類來最小化開銷。

示例

>>> from torchrl.envs import GymEnv
>>> env = GymEnv("Pendulum-1")
>>> data = env.reset()
>>> for i in range(10):
...     # compute action
...     env.rand_action(data)
...     # Perform action
...     next_data = env.step(reset_data)
...     data = env.step_mdp(next_data)

to(device: Union[device, str, int]) → EnvBase¶

移動和/或轉換引數和緩衝區。

這可以這樣呼叫

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

有關示例，請參閱下文。

注意

此方法就地修改模組。

引數:

device (torch.device) – the desired device of the parameters and buffers in this module – 此模組中引數和緩衝區的目標裝置。
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module – 此模組中引數和緩衝區的目標浮點數或複數 dtype。
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module – 其 dtype 和 device 是此模組中所有引數和緩衝區的目標 dtype 和 device 的 Tensor。
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument) – 此模組中 4D 引數和緩衝區的目標記憶體格式（僅關鍵字引數）。

返回:

self

返回型別:

模組

示例

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device: Optional[Union[int, str, device]], recurse: bool = True) → Self¶

將引數和緩衝區移動到指定裝置，而不復制儲存。

引數:

device (torch.device) – The desired device of the parameters and buffers in this module. – 此模組中引數和緩衝區的目標裝置。
recurse (bool) – 是否遞迴地將子模組的引數和緩衝區移動到指定裝置。

返回:

self

返回型別:

模組

train(mode: bool = True) → Self¶

將模組設定為訓練模式。

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc. – 這隻對某些模組有影響。有關其在訓練/評估模式下的行為的詳細資訊，例如它們是否受影響，請參閱特定模組的文件，例如 Dropout、BatchNorm 等。

引數:: mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True. – 設定訓練模式（True）或評估模式（False）。預設值：True。
返回:: self
返回型別:: 模組

type(dst_type: Union[dtype, str]) → Self¶

將所有引數和緩衝區轉換為 dst_type。

注意

此方法就地修改模組。

引數:: dst_type (type or string) – 目標型別
返回:: self
返回型別:: 模組

xpu(device: Optional[Union[device, int]] = None) → Self¶

將所有模型引數和緩衝區移動到 XPU。

這也會使關聯的引數和緩衝區成為不同的物件。因此，如果模組在最佳化時將駐留在 XPU 上，則應在構建最佳化器之前呼叫它。

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數將複製到該裝置
返回:: self
返回型別:: 模組

zero_grad(set_to_none: bool = True) → None¶

重置所有模型引數的梯度。

See similar function under torch.optim.Optimizer for more context. – 有關更多背景資訊，請參閱 torch.optim.Optimizer 下的類似函式。

引數:: set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details. – 與其設定為零，不如將 grad 設定為 None。有關詳細資訊，請參閱 torch.optim.Optimizer.zero_grad()。

LLMEnv¶

文件

教程

資源