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¶

文档

教程

资源