ThreadingAsyncEnvPool¶

class torchrl.envs.ThreadingAsyncEnvPool(*args, **kwargs)[源代码]¶

使用线程进行环境并行执行的 AsyncEnvPool 实现。

此类管理着一个环境池，每个环境都在自己的线程中运行，并提供使用线程池执行器异步执行环境的步进和重置的方法。

注意

当嵌套在 TransformedEnv 和批处理环境中使用时，此类及其子类应该可以正常工作，但用户目前无法在基础环境嵌套在这些类中使用其异步功能。用户应优先选择将转换后的环境嵌套在 AsyncEnvPool 中。如果无法做到这一点，请提交一个 issue。

_setup()[源代码]¶: 初始化线程池和环境实例。

async_step_send(tensordict)[源代码]¶: 将步进命令发送到环境。

async_step_recv(min_get)[源代码]¶: 接收步进命令的结果。

async_reset_send(tensordict)[源代码]¶: 将重置命令发送到环境。

async_reset_recv(min_get)[源代码]¶: 接收重置命令的结果。

shutdown()[源代码]¶: 关闭线程池。

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, optional) – 如果提供，将使用此 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]) → T¶

将 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）。

此方法使用提供的策略执行滚动，以推断环境的输入和输出规范。它会根据滚动过程中收集的数据更新环境的动作、观察、奖励和完成信号的规范。

参数:

policy (Callable[[TensorDictBase], TensorDictBase]) – 一个可调用的策略，它接受一个 TensorDictBase 作为输入并返回一个 TensorDictBase 作为输出。此策略用于执行滚动并确定规范。

关键字参数:

tensordict (TensorDictBase, optional) – 一个可选的 TensorDictBase 实例，用作滚动的初始状态。如果未提供，将调用环境的 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() → T¶

将所有浮点参数和缓冲区转换为 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(return_contiguous: bool | None = None, check_dtype=True, seed: int | None = None, tensordict: TensorDictBase | None = None, break_when_any_done: bool | Literal['both'] = None)¶

使用简短的 rollout 来测试环境规范。

此测试函数应作为 torchrl 的 EnvBase 子类包装的环境的健全性检查：预期的数据与收集到的数据之间的任何差异都应引发断言错误。

损坏的环境规范很可能会使并行环境无法使用。

参数:

env (EnvBase) – 要检查其规格与数据是否匹配的环境。
return_contiguous (bool, optional) – 如果为 True，将使用 return_contiguous=True 调用随机滚动。这在某些情况下可能会失败（例如，输入的/输出的形状不一致）。默认为 None（由动态规范的存在决定）。
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() → T¶

将所有模型参数和缓冲区移动到 CPU。

注意

此方法就地修改模块。

返回:: self
返回类型:: 模块

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

将所有模型参数和缓冲区移动到 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() → T¶

将所有浮点参数和缓冲区转换为 double 数据类型。

注意

此方法就地修改模块。

返回:: self
返回类型:: 模块

empty_cache()¶

清除所有缓存值。

对于常规环境，键列表（奖励、完成等）会被缓存，但在某些情况下，它们可能在代码执行期间发生更改（例如，添加转换时）。

property env_batch_sizes: list[torch.Size]¶: 返回每个环境的批次大小。

eval() → T¶

将模块设置为评估模式。

这只对某些模块有影响。有关其在训练/评估模式下的行为的更多详细信息，例如它们是否受到影响，请参阅特定模块的文档，例如 Dropout、BatchNorm 等。

这等效于 self.train(False)。

有关 .eval() 和几种可能与之混淆的类似机制之间的比较，请参阅局部禁用梯度计算。

返回:: self
返回类型:: 模块

extra_repr() → str¶

返回模块的额外表示。

要打印自定义额外信息，您应该在自己的模块中重新实现此方法。单行和多行字符串均可接受。

fake_tensordict() → TensorDictBase¶: 返回一个假的 tensordict，其键值对在形状、设备和 dtype 上与环境 rollout 期间预期的一致。

float() → T¶

将所有浮点参数和缓冲区转换为 float 数据类型。

注意

此方法就地修改模块。

返回:: self
返回类型:: 模块

forward(*args, **kwargs)¶

定义每次调用时执行的计算。

所有子类都应重写此方法。

注意

虽然前向传播的配方需要在该函数内定义，但之后应调用 Module 实例而不是这个，因为前者会处理运行已注册的钩子，而后者会默默地忽略它们。

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 的序列化正常工作。我们仅为序列化 Tensor 提供向后兼容性保证；其他对象的序列化后的 pickled 形式如果发生更改，可能会破坏向后兼容性。

返回:: 要存储在模块 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 – 如果在 target 字符串解析出的路径中的任何一点，(子)路径解析为一个不存在的属性名称或一个不是 nn.Module 实例的对象。

half() → T¶

将所有浮点参数和缓冲区转换为 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[int, device]] = None) → T¶

将所有模型参数和缓冲区移动到 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, optional) – 是否严格执行 state_dict 中的键与此模块的 state_dict() 函数返回的键是否匹配。默认为：True
assign (bool, optional) – 当设置为 False 时，将保留当前模块中张量的属性，而设置为 True 时，将保留 state dict 中张量的属性。唯一的例外是 requires_grad 字段 Default: ``False`

返回:

missing_keys 是一个包含任何预期键的 str 列表。
在提供的 state_dict 中缺失的任何键的字符串列表。
unexpected_keys 是一个包含不匹配的键的 str 列表。
不期望但在提供的 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[int, device]] = None) → T¶

将所有模型参数和缓冲区移动到 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 实例。规范中列出的键在重置和步进后可直接访问。

在 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]], Union[None, 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¶

向模块添加一个缓冲区。

这通常用于注册不应被视为模型参数的 buffer。例如，BatchNorm 的 running_mean 不是参数，而是模块状态的一部分。默认情况下，buffers 是持久的，并将与参数一起保存。此行为可以通过将 persistent 设置为 False 来更改。持久 buffer 和非持久 buffer 之间的唯一区别是后者不会是此模块 state_dict 的一部分。

可以使用给定名称作为属性访问缓冲区。

参数:

name (str) – buffer 的名称。可以使用给定的名称从此模块访问 buffer
tensor (Tensor 或 None) – 要注册的 buffer。如果为 None，则运行在 buffers 上的操作，例如 cuda，将被忽略。如果为 None，则该 buffer **不会** 包含在模块的 state_dict 中。
persistent (bool) – buffer 是否是此模块 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 调用后调用的，因此不会影响前向传播。hook 应具有以下签名

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

如果 with_kwargs 为 True，则会将传递给 forward 函数的 kwargs 传递给 forward hook，并期望它返回可能已修改的输出。hook 应具有以下签名

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() 之前都会调用该 hook。

如果 with_kwargs 为 false 或未指定，则输入仅包含传递给模块的位置参数。关键字参数不会传递给 hook，只会传递给 forward。hook 可以修改输入。用户可以返回一个元组或 hook 中的单个修改值。我们将返回的值包装成一个元组（如果返回的是单个值）（除非该值已经是元组）。hook 应具有以下签名

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

如果 with_kwargs 为 true，则会将传递给 forward 函数的 kwargs 传递给 forward pre-hook。如果 hook 修改了输入，则应返回 args 和 kwargs。hook 应具有以下签名

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 之前触发。否则，提供的 hook 将在当前 torch.nn.Module 的所有现有 forward_pre hook 之后触发。请注意，使用 register_module_forward_pre_hook() 注册的全局 forward_pre hook 将在通过此方法注册的所有 hook 之前触发。默认值：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]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模块上注册一个反向传播钩子。

hook 将在计算模块的梯度时被调用，即只有在计算模块输出的梯度时 hook 才会执行。hook 应具有以下签名

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

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入或输出，否则将引发错误。

参数:

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

返回:

一个句柄，可用于通过调用 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

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入，否则将引发错误。

参数:

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

返回:

一个句柄，可用于通过调用 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 字典，并将从 observation 键中排除。此参数可以在调用 make() 时传递（参见下面的示例）。

警告

使用 info_keys 可能会导致 spec 为空，因为内容已移至 info 字典。Gym 不喜欢 specs 中空的 Dict，因此应使用 RemoveEmptySpecs 移除此空内容。
backend (str, optional) – 后端。可以是 “gym” 或 “gymnasium” 或任何其他与 set_gym_backend 兼容的后端。
to_numpy (bool, optional) – 如果为 True，则调用 step 和 reset 的结果将被映射到 numpy 数组。默认为 False（结果是 tensors）。此参数可以在调用 make() 时传递（参见下面的示例）。
reward_threshold (float, optional) – [Gym kwarg] 学习环境的奖励阈值。
nondeterministic (bool, optional) – [Gym kwarg] 如果环境是非确定性的（即使知道初始种子和所有动作）。默认为 False。
max_episode_steps (int, optional) – [Gym kwarg] 截断之前的最大单集步数。由 Time Limit 包装器使用。
order_enforce (bool, optional) – [Gym >= 0.14] 是否应用顺序强制器包装器以确保用户以正确的顺序运行函数。默认为 True。
autoreset (bool, optional) – [Gym >= 0.14 and <1.0.0] 是否添加 autores集包装器，以便不需要调用 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¶: add_module() 的别名。

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

向模块添加一个参数。

可以使用给定名称作为属性访问该参数。

参数:

name (str) – 参数的名称。可以通过给定名称从该模块访问该参数。
param (Parameter 或 None) – 要添加到模块的参数。如果为 None，则运行在参数上的操作，例如 cuda，将被忽略。如果为 None，则该参数 **不会** 包含在模块的 state_dict 中。

register_state_dict_post_hook(hook)¶

为 state_dict() 方法注册一个后置 hook。

它应该具有以下签名：: hook(module, state_dict, prefix, local_metadata) -> None

注册的钩子可以就地修改 state_dict。

register_state_dict_pre_hook(hook)¶

为 state_dict() 方法注册一个前置 hook。

它应该具有以下签名：: hook(module, prefix, keep_vars) -> None

注册的钩子可用于在进行 state_dict 调用之前执行预处理。

requires_grad_(requires_grad: bool = True) → T¶

更改自动梯度是否应记录此模块中参数的操作。

此方法就地设置参数的 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 位置的（可能为空的）字符串元组，可以在其中找到完成状态。

键按数据树的深度排序。

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 或整个 observation 条目序列的 callable，这些序列按 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，则在任何包含的环境达到任何完成状态时中断。如果为 False，则已完成的环境会自动重置。默认值为 True。

另请参阅

文档的部分重置提供了更多关于部分重置的信息。
break_when_all_done (bool, optional) –
如果为 True，则在所有包含的环境都达到任何完成状态时中断。如果为 False，则在至少一个环境达到任何完成状态时中断。默认值为 False。

另请参阅

文档的部分步进提供了更多关于部分重置的信息。
return_contiguous (bool) – If False，将返回一个 LazyStackedTensorDict。如果环境没有动态规范，则默认为 True，否则为 False。
tensordict (TensorDict, optional) – 如果 auto_reset 为 False，则必须提供初始 tensordict。Rollout 将检查此 tensordict 是否具有完成标志，并在那些维度上重置环境（如果需要）。如果 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 策略将被信任并假定与 collector 兼容。对于 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 成为 tensordicts 的懒惰堆栈。

非连续 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']

Rollouts 可以在循环中使用以模拟数据收集。为此，您需要在调用 step_mdp() 后，将上一个 rollout 的最后一个 tensordict 作为输入。

数据收集 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 中包含的额外状态。

此函数从 load_state_dict() 调用，以处理 state_dict 中找到的任何额外状态。如果需要将额外状态存储在模块的 state_dict 中，请为您自己的模块实现此函数和相应的 get_extra_state()。

参数:: 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¶: 等同于 batch_size。

share_memory() → T¶: 参见 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) – 如果提供，模块的状态将更新到该字典中，并返回相同的对象。否则，将创建一个 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) – 默认情况下，状态字典中返回的 Tensor 会从自动微分中分离。如果设置为 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) – step() 方法的输入数据结构。

此方法允许轻松编写非停止的 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 实例来高效处理状态、观测、动作、奖励和完成键的过渡。_StepMDP 类通过预先计算要包含和排除的键来优化此过程，从而在重复调用期间减少运行时开销。_StepMDP 实例以 exclude_action=False 创建，这意味着动作键会保留在根 tensordict 中。

参数:: next_tensordict (TensorDictBase) – 包含下一时间步环境状态的 tensordict。此 tensordict 应包含观测、动作、奖励和完成标志的键，具体取决于环境的规范。
返回:: 一个代表环境状态推进一步后的新 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)

其签名与 torch.Tensor.to() 类似，但只接受浮点数或复数 dtype。此外，此方法只会将浮点数或复数参数和缓冲区转换为 dtype（如果给定）。整数参数和缓冲区将被移动到 device（如果给定），但 dtypes 不变。当设置 non_blocking 时，如果可能，它会尝试异步（相对于主机）进行转换/移动，例如，将具有固定内存的 CPU Tensor 移动到 CUDA 设备。

有关示例，请参阅下文。

注意

此方法就地修改模块。

参数:

device (torch.device) – 该模块中参数和缓冲区的期望设备
dtype (torch.dtype) – 该模块中参数和缓冲区的期望浮点数或复数 dtype
tensor (torch.Tensor) – 其 dtype 和 device 是该模块中所有参数和缓冲区的期望 dtype 和 device 的 Tensor
memory_format (torch.memory_format) – 该模块中 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) → T¶

将参数和缓冲区移动到指定设备，而不复制存储。

参数:

device (torch.device) – 该模块中参数和缓冲区的期望设备。
recurse (bool) – 是否递归地将子模块的参数和缓冲区移动到指定设备。

返回:

self

返回类型:

模块

train(mode: bool = True) → T¶

将模块设置为训练模式。

这仅对某些模块有效。请参阅特定模块的文档，了解它们在训练/评估模式下的行为，即它们是否受影响，例如 Dropout、BatchNorm 等。

参数:: mode (bool) – 设置训练模式（True）或评估模式（False）。默认为 True。
返回:: self
返回类型:: 模块

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

将所有参数和缓冲区转换为 dst_type。

注意

此方法就地修改模块。

参数:: dst_type (type or string) – 目标类型
返回:: self
返回类型:: 模块

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

将所有模型参数和缓冲区移动到 XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 XPU 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:: device (int, optional) – 如果指定，所有参数将复制到该设备
返回:: self
返回类型:: 模块

zero_grad(set_to_none: bool = True) → None¶

重置所有模型参数的梯度。

请参见 torch.optim.Optimizer 下类似的函数以获取更多上下文。

参数:: set_to_none (bool) – 在设置为零而不是设置为零的情况下，将梯度设置为 None。有关详细信息，请参阅 torch.optim.Optimizer.zero_grad()。

ThreadingAsyncEnvPool¶

文档

教程

资源