SerialEnv¶

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

在同一个进程中创建一系列环境。批处理环境允许用户查询远程运行的环境的任意方法/属性。

这些查询将返回一个长度等于工作进程数量的列表，其中包含这些查询的结果值。

>>> env = ParallelEnv(3, my_env_fun)
>>> custom_attribute_list = env.custom_attribute
>>> custom_method_list = env.custom_method(*args)

参数：

num_workers – 要同时部署的工作进程（即环境实例）的数量；
create_env_fn (callable 或 可调用列表) – 用于创建环境的函数（或函数列表）。如果使用单一任务，则应使用可调用对象，而不是使用相同可调用对象的列表：如果提供可调用对象列表，则环境将被执行，就好像需要多个不同的任务一样，这会带来轻微的计算开销；

关键字参数：

create_env_kwargs (dict 或 dicts 列表, 可选) – 将用于创建的环境的 kwargs；
share_individual_td (bool, 可选) – 如果为 True，则为每个进程/工作进程创建不同的 tensordict，并返回一个延迟堆栈。默认值 = None（如果为单一任务，则为 False）；
shared_memory (bool) – 返回的 tensordict 是否将放在共享内存中；
memmap (bool) – 返回的 tensordict 是否将放在内存映射中。
policy_proof (callable, 可选) – 如果提供，它将用于获取要通过 step() 和 reset() 方法返回的张量列表，例如 "hidden" 等。
device (str, int, torch.device) – 批处理环境的设备可以传递。如果没有，它将从 env 推断出来。在这种情况下，假设所有环境的设备都匹配。如果提供，它可能与子环境设备不同。在这种情况下，数据将在收集期间自动转换为适当的设备。这可用于加快收集速度，以防转换为设备会带来开销（例如，基于 numpy 的环境等）：通过使用批处理环境的 "cuda" 设备，但使用嵌套环境的 "cpu" 设备，可以将开销降至最低。
num_threads (int, 可选) – 此进程的线程数。应等于每个子进程中启动的进程数加 1（如果只启动单个进程，则为 1）。默认值为工作进程数加 1。此参数对 SerialEnv 类无效。
num_sub_threads (int, 可选) – 子进程的线程数。默认值为 1，以确保安全：如果没有指定，启动多个工作进程可能会过度占用 CPU 负载，并损害性能。此参数对 SerialEnv 类无效。
serial_for_single (bool, 可选) – 如果为 True，则使用单个工作进程创建并行环境将返回一个 SerialEnv。此选项对 SerialEnv 无效。默认值为 False。
non_blocking (bool, 可选) – 如果为 True，则设备移动将使用 non_blocking=True 选项执行。默认值为 True。
mp_start_method (str, 可选) – 多处理启动方法。如果没有指定，则使用默认启动方法（如果在第一次导入之前没有以其他方式启动，则在 TorchRL 中默认使用 ‘spawn’）。仅与 ParallelEnv 子类一起使用。
use_buffers (bool, 可选) –
工作进程之间的通信是否应通过循环预分配内存缓冲区进行。默认值为 True，除非环境之一具有动态规格。

注意

详细了解动态规格和环境此处.

注意

可以使用以下技术将关键字参数传递给每个子环境：reset() 中的每个关键字参数都将传递给每个环境，除了 list_of_kwargs 参数，如果存在，它应包含一个与工作进程数量相同长度的列表，其中 worker 特定的关键字参数存储在一个字典中。如果查询部分重置，则将忽略与未重置的子环境相对应的 list_of_kwargs 元素。

示例

>>> from torchrl.envs import GymEnv, ParallelEnv, SerialEnv, EnvCreator
>>> make_env = EnvCreator(lambda: GymEnv("Pendulum-v1")) # EnvCreator ensures that the env is sharable. Optional in most cases.
>>> env = SerialEnv(2, make_env)  # Makes 2 identical copies of the Pendulum env, runs them on the same process serially
>>> env = ParallelEnv(2, make_env)  # Makes 2 identical copies of the Pendulum env, runs them on dedicated processes
>>> from torchrl.envs import DMControlEnv
>>> env = ParallelEnv(2, [
...     lambda: DMControlEnv("humanoid", "stand"),
...     lambda: DMControlEnv("humanoid", "walk")])  # Creates two independent copies of Humanoid, one that walks one that stands
>>> rollout = env.rollout(10)  # executes 10 random steps in the environment
>>> rollout[0]  # data for Humanoid stand
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
        com_velocity: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        extremities: Tensor(shape=torch.Size([10, 12]), device=cpu, dtype=torch.float64, is_shared=False),
        head_height: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
        joint_angles: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
        next: TensorDict(
            fields={
                com_velocity: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                extremities: Tensor(shape=torch.Size([10, 12]), device=cpu, dtype=torch.float64, is_shared=False),
                head_height: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                joint_angles: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                torso_vertical: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                velocity: Tensor(shape=torch.Size([10, 27]), device=cpu, dtype=torch.float64, is_shared=False)},
            batch_size=torch.Size([10]),
            device=cpu,
            is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        torso_vertical: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        velocity: Tensor(shape=torch.Size([10, 27]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)
>>> rollout[1]  # data for Humanoid walk
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
        com_velocity: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        extremities: Tensor(shape=torch.Size([10, 12]), device=cpu, dtype=torch.float64, is_shared=False),
        head_height: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
        joint_angles: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
        next: TensorDict(
            fields={
                com_velocity: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                extremities: Tensor(shape=torch.Size([10, 12]), device=cpu, dtype=torch.float64, is_shared=False),
                head_height: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                joint_angles: Tensor(shape=torch.Size([10, 21]), device=cpu, dtype=torch.float64, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                torso_vertical: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                velocity: Tensor(shape=torch.Size([10, 27]), device=cpu, dtype=torch.float64, is_shared=False)},
            batch_size=torch.Size([10]),
            device=cpu,
            is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        torso_vertical: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float64, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        velocity: Tensor(shape=torch.Size([10, 27]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)
>>> # serial_for_single to avoid creating parallel envs if not necessary
>>> env = ParallelEnv(1, make_env, serial_for_single=True)
>>> assert isinstance(env, SerialEnv)  # serial_for_single allows you to avoid creating parallel envs when not necessary

property action_key: NestedKey¶

环境的动作键。

默认情况下，这将为 “action”。

如果环境中存在多个动作键，则此函数将引发异常。

property action_keys: List[NestedKey]¶

环境的动作键。

默认情况下，将只有一个名为 “action” 的键。

键按数据树中的深度排序。

property action_spec: TensorSpec¶

The action spec.

The action_spec 始终存储为复合规范。

如果动作规范以简单规范的形式提供，则将返回此规范。

>>> env.action_spec = UnboundedContinuousTensorSpec(1)
>>> env.action_spec
UnboundedContinuousTensorSpec(
    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)

如果动作规范以复合规范的形式提供，并且仅包含一个叶子节点，则此函数将只返回该叶子节点。

>>> env.action_spec = CompositeSpec({"nested": {"action": UnboundedContinuousTensorSpec(1)}})
>>> env.action_spec
UnboundedContinuousTensorSpec(
    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)

如果动作规范以复合规范的形式提供，并且包含多个叶子节点，则此函数将返回整个规范。

>>> env.action_spec = CompositeSpec({"nested": {"action": UnboundedContinuousTensorSpec(1), "another_action": DiscreteTensorSpec(1)}})
>>> env.action_spec
CompositeSpec(
    nested: CompositeSpec(
        action: UnboundedContinuousTensorSpec(
            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: DiscreteTensorSpec(
            shape=torch.Size([]),
            space=DiscreteBox(n=1),
            device=cpu,
            dtype=torch.int64,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

要检索传递的完整规范，请使用

>>> env.input_spec["full_action_spec"]

此属性是可变的。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.action_spec
BoundedTensorSpec(
    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)

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

将子模块添加到当前模块。

可以使用给定的名称以属性的形式访问该模块。

参数：

name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块。
module (Module) – 要添加到模块的子模块。

add_truncated_keys()¶: 在环境中添加截断键。

append_transform(transform: 'Transform' | Callable[[TensorDictBase], TensorDictBase]) → None¶

返回一个经过转换的环境，其中应用了传递的可调用函数/转换。

参数：: transform (Transform or 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
返回类型:: Module

示例

>>> @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)
)

property batch_locked: bool¶

环境是否可以使用与初始化时不同的批次大小。

如果为 True，则需要使用具有与环境相同的批次大小的张量字典来使用 env。batch_locked 是一个不可变属性。

property batch_size: Size¶

在此环境实例中批处理的环境数量，组织为 torch.Size() 对象。

环境可能相似或不同，但假定它们之间几乎没有或根本没有交互（例如，多任务或并行批处理执行）。

bfloat16() → T¶

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

注意

此方法在原地修改模块。

返回值:: self
返回类型:: Module

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

返回模块缓冲区的迭代器。

参数：: recurse (bool) – 如果为 True，则会生成此模块和所有子模块的缓冲区。否则，只生成此模块的直接成员的缓冲区。
返回值:: 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)

children() → Iterator[
返回值:: Module – a child module

compile(*args, **kwargs)¶

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

cpu() → T¶

Move all model parameters and buffers to the CPU.

注意

此方法在原地修改模块。

返回值:: self
返回类型:: Module

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

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

注意

此方法在原地修改模块。

参数：: device (int, optional) – if specified, all parameters will be copied to that device
返回值:: self
返回类型:: Module

property done_key¶

The done key of an environment.

By default, this will be “done”.

If there is more than one done key in the environment, this function will raise an exception.

property done_keys: List[NestedKey]¶

The done keys of an environment.

By default, there will only be one key named “done”.

键按数据树中的深度排序。

property done_keys_groups¶

A list of done keys, grouped as the reset keys.

This is a list of lists. The outer list has the length of reset keys, the inner lists contain the done keys (eg, done and truncated) that can be read to determine a reset when it is absent.

property done_spec: TensorSpec¶

The done spec.

The done_spec is always stored as a composite spec.

If the done spec is provided as a simple spec, this will be returned.

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

If the done spec is provided as a composite spec and contains only one leaf, this function will return just the leaf.

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

If the done spec is provided as a composite spec and has more than one leaf, this function will return the whole spec.

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

To always retrieve the full spec passed, use

>>> env.output_spec["full_done_spec"]

此属性是可变的。

示例

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

double() → T¶

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

注意

此方法在原地修改模块。

返回值:: self
返回类型:: Module

empty_cache()¶

清除所有缓存的值。

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

eval() → T¶

将模块设置为评估模式。

这仅对某些模块有效。有关特定模块在训练/评估模式下的行为细节，如果它们受到影响，例如 Dropout、BatchNorm 等，请参阅相关文档。

这等效于 self.train(False)。

请参阅本地禁用梯度计算以比较 .eval() 和一些可能与之混淆的类似机制。

返回值:: self
返回类型:: Module

extra_repr() → str¶

设置模块的额外表示。

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

fake_tensordict() → TensorDictBase¶: 返回一个假的张量字典，其键值对在形状、设备和数据类型方面与环境回滚期间预期的一致。

float() → T¶

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

注意

此方法在原地修改模块。

返回值:: self
返回类型:: Module

forward(tensordict: TensorDictBase) → TensorDictBase¶

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

应由所有子类覆盖。

注意

虽然正向传递的配方需要在此函数内定义，但应该随后调用 Module 实例，而不是它，因为前者负责运行注册的钩子，而后者则默默地忽略它们。

property full_action_spec: CompositeSpec¶

完整的动作规格。

full_action_spec 是一个 CompositeSpec` 实例，包含所有动作条目。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.full_action_spec

CompositeSpec(

action: BoundedTensorSpec(

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_done_spec: CompositeSpec¶

完整的完成规格。

full_done_spec 是一个 CompositeSpec` 实例，包含所有完成条目。它可以用来生成与运行时获得的结构相似的假数据。

示例

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

property full_reward_spec: CompositeSpec¶

完整的奖励规格。

full_reward_spec 是一个 CompositeSpec` 实例，包含所有奖励条目。

示例

>>> 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
CompositeSpec(
    nested: CompositeSpec(
        reward: UnboundedContinuousTensorSpec(
            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_state_spec: CompositeSpec¶

完整的状态规格。

full_state_spec 是一个 CompositeSpec` 实例，包含所有状态条目（即，不是动作的输入数据）。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.full_state_spec
CompositeSpec(
    state: CompositeSpec(
        pipeline_state: CompositeSpec(
            q: UnboundedContinuousTensorSpec(
                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([]))

get_buffer(target: str) → Tensor¶

如果存在，则返回由 target 给出的缓冲区，否则抛出错误。

有关此方法功能的更详细解释以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数：: target – 要查找的缓冲区的完全限定字符串名称。 (有关如何指定完全限定字符串，请参阅 get_submodule)。
返回值:: 由 target 引用的缓冲区
返回类型:: torch.Tensor
引发:: AttributeError – 如果目标字符串引用无效路径或解析为不是缓冲区的内容

get_extra_state() → Any¶

返回要包含在模块的 state_dict 中的任何额外状态。

如果需要存储额外状态，请实现此方法和相应的 set_extra_state() 。此函数在构建模块的 state_dict() 时被调用。

请注意，额外状态应该是可腌制的，以确保 state_dict 的序列化正常工作。我们只提供对序列化张量进行向后兼容性的保证；其他对象如果它们的序列化腌制形式发生变化，可能会破坏向后兼容性。

返回值:: 要存储在模块的 state_dict 中的任何额外状态
返回类型:: object

get_parameter(target: str) → Parameter¶

如果存在，则返回由 target 给出的参数，否则抛出错误。

有关此方法功能的更详细解释以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数：: target – 要查找的参数的完全限定字符串名称。 (有关如何指定完全限定字符串，请参阅 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() → T¶

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

注意

此方法在原地修改模块。

返回值:: self
返回类型:: Module

property input_spec: TensorSpec¶

输入规范。

包含所有环境数据输入规范的组合规范。

它包含

“full_action_spec”: 输入动作的规范
“full_state_spec”: 所有其他环境输入的规范

此属性已锁定，应为只读。相反，要设置其中包含的规范，请使用相应的属性。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.input_spec
CompositeSpec(
    full_state_spec: None,
    full_action_spec: CompositeSpec(
        action: BoundedTensorSpec(
            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([]))

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

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

这也使关联的参数和缓冲区成为不同的对象。因此，如果模块将在 IPU 上运行并进行优化，则应在构造优化器之前调用它。

注意

此方法在原地修改模块。

参数：: device (int, optional) – if specified, all parameters will be copied to that device
返回值:: self
返回类型:: Module

load_state_dict(*args, **kwargs)[source]¶

将参数和缓冲区从state_dict复制到此模块及其后代。

如果strict为True，则state_dict中的键必须与该模块的state_dict()函数返回的键完全匹配。

警告

如果assign为True，则优化器必须在调用load_state_dict之后创建，除非get_swap_module_params_on_conversion()为True。

参数：

state_dict (dict) – 包含参数和持久缓冲区的字典。
strict (bool, optional) – 是否严格强制state_dict中的键与该模块的state_dict()函数返回的键匹配。默认值：True
assign (bool, optional) – 当False时，保留当前模块中张量的属性，而当True时，保留状态字典中张量的属性。唯一例外是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) – if specified, all parameters will be copied to that device
返回值:: self
返回类型:: Module

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

返回模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数：

prefix (str) – 要添加到所有缓冲区名称之前的首缀。
recurse (bool, optional) – 如果为 True，则生成此模块和所有子模块的缓冲区。否则，仅生成作为此模块的直接成员的缓冲区。默认值为 True。
remove_duplicate (bool, optional) – 是否从结果中删除重复的缓冲区。默认值为 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 – 用于存储已添加到结果中的模块集的备忘录
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, 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_spec: CompositeSpec¶

观测规范。

必须是 torchrl.data.CompositeSpec 实例。规范中列出的键在重置和步骤后可以直接访问。

在 TorchRL 中，即使它们不是严格意义上的“观测”，环境的所有信息、状态、变换结果等输出都存储在 observation_spec 中。

因此，"observation_spec" 应该被认为是环境输出的通用数据容器，这些输出不是 done 或 reward 数据。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.observation_spec
CompositeSpec(
    observation: BoundedTensorSpec(
        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 output_spec: TensorSpec¶

输出规范。

包含环境输出数据所有规范的复合规范。

它包含

“full_reward_spec”: 奖励的规范
“full_done_spec”: done 的规范
“full_observation_spec”: 所有其他环境输出的规范

此属性已锁定，应为只读。相反，要设置其中包含的规范，请使用相应的属性。

示例

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = GymEnv("Pendulum-v1")
>>> env.output_spec
CompositeSpec(
    full_reward_spec: CompositeSpec(
        reward: UnboundedContinuousTensorSpec(
            shape=torch.Size([1]),
            space=None,
            device=cpu,
            dtype=torch.float32,
            domain=continuous), device=cpu, shape=torch.Size([])),
    full_observation_spec: CompositeSpec(
        observation: BoundedTensorSpec(
            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: CompositeSpec(
        done: DiscreteTensorSpec(
            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([]))

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: Optional[TensorDictBase] = None)¶

根据 action_spec 属性执行随机动作。

参数：: tensordict (TensorDictBase, optional) – 应写入结果动作的 tensordict。
返回值:: 一个 tensordict 对象，其中“action”条目更新为来自 action-spec 的随机样本。

rand_step(tensordict: Optional[TensorDictBase] = None) → TensorDictBase¶

根据 action_spec 属性在环境中执行随机步骤。

参数：: tensordict (TensorDictBase, optional) – 应写入结果信息的 tensordict。
返回值:: 一个 tensordict 对象，其中包含在环境中执行随机步骤后的新观测。动作将使用“action”键存储。

register_backward_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor], Union[Tuple[Tensor, ...], Tensor]], Union[None, Tuple[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) – 缓冲区的名称。可以使用给定的名称从此模块访问缓冲区
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_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() 计算出输出后，都会调用该钩子。

如果 with_kwargs 为 False 或未指定，则输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给 forward。钩子可以修改输出。它可以就地修改输入，但它不会影响前向，因为这是在调用 forward() 之后调用的。钩子应具有以下签名

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

如果 with_kwargs 为 True，则前向钩子将传递给前向函数的 kwargs，并期望返回可能被修改的输出。钩子应具有以下签名

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

参数：

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – 如果为 True，则提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward 钩子之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward 钩子将在本方法注册的所有钩子之前触发。默认值：False
with_kwargs (bool) – 如果为 True，则 hook 将传递给前向函数的 kwargs。默认值：False
always_call (bool) – 如果为 True，则无论在调用模块时是否抛出异常，都会运行 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，则前向预钩子将传递给前向函数的 kwargs。如果钩子修改了输入，则应返回 args 和 kwargs。钩子应具有以下签名

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

参数：

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – 如果为 true，则提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 forward_pre 钩子之后触发。请注意，使用 register_module_forward_pre_hook() 注册的全局 forward_pre 钩子将在本方法注册的所有钩子之前触发。默认值：False
with_kwargs (bool) – 如果为 true，则 hook 将传递给前向函数的 kwargs。默认值：False

返回值:

一个句柄，可用于通过调用 handle.remove() 来删除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

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

在模块上注册反向钩子。

该钩子将在每次计算模块的梯度时调用，即当且仅当计算模块输出的梯度时，该钩子才会执行。该钩子应具有以下签名

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

The grad_input 和 grad_output 是包含关于输入和输出的梯度的元组。钩子不应该修改其参数，但它可以选择返回一个新的关于输入的梯度，该梯度将在后续计算中代替 grad_input 使用。 grad_input 将仅对应于作为位置参数给出的输入，并且所有 kwarg 参数都被忽略。 grad_input 和 grad_output 中的条目将对于所有非张量参数为 None。

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

警告

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

参数：

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – 如果为真，则提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward 钩子之后触发。请注意，使用 register_module_full_backward_hook() 注册的全局 backward 钩子将在使用此方法注册的所有钩子之前触发。

返回值:

一个句柄，可用于通过调用 handle.remove() 来删除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

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

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

该钩子将在每次计算模块的梯度时调用。该钩子应具有以下签名

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

The grad_output 是一个元组。钩子不应该修改其参数，但它可以选择返回一个新的关于输出的梯度，该梯度将在后续计算中代替 grad_output 使用。 grad_output 中的条目将对于所有非张量参数为 None。

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

警告

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

参数：

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – 如果为真，则提供的 hook 将在该 torch.nn.modules.Module 上的所有现有 backward_pre 钩子之前触发。否则，提供的 hook 将在该 torch.nn.modules.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 = False, disable_env_checker: bool = False, apply_api_compatibility: bool = False, **kwargs)¶

在 gym(nasium) 中注册一个环境。

此方法的设计考虑了以下范围

将 TorchRL 首要环境纳入使用 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) – 将与环境一起使用的转换（或 torchrl.envs.Compose 实例中的转换列表）。此参数可以在调用 make() 时传递（见下例）。
info_keys (List[NestedKey], optional) –
如果提供，这些键将用于构建 info 字典，并将从观察键中排除。此参数可以在调用 make() 时传递（见下例）。

警告

使用 info_keys 可能会使规范为空，因为内容已移动到 info 字典中。Gym 不喜欢规范中的空 Dict，因此应使用 RemoveEmptySpecs 删除此空内容。
backend (str, optional) – 后端。可以是 “gym” 或 “gymnasium”，或任何其他与 set_gym_backend 兼容的后端。
to_numpy (bool, optional) – 如果为 True，则对 step 和 reset 的调用的结果将映射到 numpy 数组。默认为 False（结果为张量）。此参数可以在调用 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] 是否应添加自动重置包装器，以便无需调用重置。默认为 False。
disable_env_checker – [Gym >= 0.14] 是否应禁用环境检查器以进行环境检查。默认为 False。
apply_api_compatibility – [Gym >= 0.26] 是否要应用 StepAPICompatibility 包装器。默认为 False。
**kwargs – 传递给环境构造函数的任意关键字参数。

注意

TorchRL 的环境没有 "info" 字典的概念，因为 TensorDict 提供了在大多数训练设置中被认为必要的全部存储需求。尽管如此，您仍然可以使用 info_keys 参数来对被视为观察的内容和被视为信息的内容进行细粒度控制。

示例

>>> # 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

module 参数是当前注册了此挂钩的模块，而 incompatible_keys 参数是 NamedTuple，由属性 missing_keys 和 unexpected_keys 组成。missing_keys 是包含缺失键的 list，而 unexpected_keys 是包含意外键的 list。

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

请注意，使用 strict=True 调用 load_state_dict() 时执行的检查会受到挂钩对 missing_keys 或 unexpected_keys 的修改的影响，如预期的那样。向任一组键添加内容会导致在 strict=True 时抛出错误，而清除缺失键和意外键都会避免错误。

返回值:: 一个句柄，可用于通过调用 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 or None) – 要添加到模块的参数。如果为 None，则忽略在参数上运行的操作，例如 cuda。如果为 None，则不会将参数包含在模块的 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) → T¶

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

此方法就地设置参数的 requires_grad 属性。

此方法对于冻结模块的某些部分以进行微调或单独训练模型的某些部分（例如 GAN 训练）很有帮助。

请参阅在本地禁用梯度计算，以比较 .requires_grad_() 和可能与其混淆的几种类似机制。

参数：: requires_grad (bool) – 是否应记录此模块中参数上的自动梯度操作。默认值：True。
返回值:: self
返回类型:: Module

reset(tensordict: Optional[TensorDictBase] = None, **kwargs) → TensorDictBase¶

重置环境。

至于 step 和 _step，只有私有方法 _reset 应该被 EnvBase 子类覆盖。

参数：

tensordict (TensorDictBase, optional) – 用于包含生成的新的观察结果的张量字典。在某些情况下，此输入也可以用于将参数传递给 reset 函数。
kwargs (optional) – 要传递给本机 reset 函数的其他参数。

返回值:

就地修改的张量字典（如果有的话），其中包含生成的观察结果。

property reset_keys: List[NestedKey]¶

返回重置键的列表。

重置键是指示部分重置的键，在批处理、多任务或多代理设置中。它们被构造为 (*prefix, "_reset")，其中 prefix 是一个（可能是空的）字符串元组，指向可以在其中找到完成状态的张量字典位置。

键按数据树中的深度排序。

property reward_key¶

环境的奖励键。

默认情况下，这将是“reward”。

如果环境中有多个奖励键，则此函数将引发异常。

property reward_keys: List[NestedKey]¶

环境的奖励键。

默认情况下，只有一个名为“reward”的键。

键按数据树中的深度排序。

property reward_spec: TensorSpec¶

The reward spec.

The reward_spec is always stored as a composite spec.

If the reward spec is provided as a simple spec, this will be returned.

>>> env.reward_spec = UnboundedContinuousTensorSpec(1)
>>> env.reward_spec
UnboundedContinuousTensorSpec(
    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)

If the reward spec is provided as a composite spec and contains only one leaf, this function will return just the leaf.

>>> env.reward_spec = CompositeSpec({"nested": {"reward": UnboundedContinuousTensorSpec(1)}})
>>> env.reward_spec
UnboundedContinuousTensorSpec(
    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)

If the reward spec is provided as a composite spec and has more than one leaf, this function will return the whole spec.

>>> env.reward_spec = CompositeSpec({"nested": {"reward": UnboundedContinuousTensorSpec(1), "another_reward": DiscreteTensorSpec(1)}})
>>> env.reward_spec
CompositeSpec(
    nested: CompositeSpec(
        reward: UnboundedContinuousTensorSpec(
            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: DiscreteTensorSpec(
            shape=torch.Size([]),
            space=DiscreteBox(n=1),
            device=cpu,
            dtype=torch.int64,
            domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))

要检索传递的完整规范，请使用

>>> env.output_spec["full_reward_spec"]

此属性是可变的。

示例

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

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

在环境中执行一次 rollout。

当包含的环境之一返回 done=True 时，该函数将停止。

参数：

max_steps (int) – 要执行的最大步数。如果环境在执行 max_steps 之前达到 done 状态，则实际步数可能会更小。
policy (callable, optional) – 用于计算所需动作的可调用对象。如果未提供策略，则将使用 env.rand_step() 调用动作。策略可以是任何可调用对象，它读取 tensordict 或整个观测条目序列，这些条目 __sorted as__ env.observation_spec.keys()。默认值为 None。
callback (Callable[[TensorDict], Any], optional) – 在每次迭代时使用给定的 TensorDict 调用的函数。默认值为 None。 callback 的输出不会被收集，用户有责任在回调调用中保存任何结果，如果数据需要在对 rollout 的调用之外进行传递。
auto_reset (bool, optional) – 如果 True，则在启动 rollout 时，如果环境处于 done 状态，则会自动重置环境。默认值为 True。
auto_cast_to_device (bool, optional) – 如果 True，则在使用策略之前，tensordict 的设备会自动转换为策略设备。默认值为 False。
break_when_any_done (bool) – 如果任何 done 状态为 True，则中断。如果为 False，则对处于 done 状态的子环境调用 reset()。默认值为 True。
return_contiguous (bool) – 如果为 False，则返回 LazyStackedTensorDict。默认值为 True。
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。

返回值:

包含结果轨迹的 TensorDict 对象。

返回的数据将用 tensordict 的最后一个维度（在 env.ndim 索引处）的“时间”维度名称标记。

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

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

参数：: state (dict) – 来自 state_dict 的额外状态

set_seed(*args, **kwargs)[source]¶

设置环境的种子并返回要使用的下一个种子（如果存在单个环境，则为输入种子）。

参数：

seed (int) – 要设置的种子。种子仅在环境中本地设置。要处理全局种子，请参见 manual_seed()。
static_seed (bool, optional) – 如果 True，则种子不会递增。默认值为 False

返回值:

即，如果与该环境同时创建，则应用于另一个环境的种子。

返回类型:

表示“下一个种子”的整数

set_submodule(target: str, module: Module) → None¶

如果存在，则设置由 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 覆盖 Conv2d，您将调用 set_submodule("net_b.net_c.conv", nn.Linear(33, 16))。

参数：

target – 要查找的子模块的完全限定字符串名称。 (有关如何指定完全限定字符串，请参阅上面的示例)。
module – 要将子模块设置到的模块。

引发:

ValueError – 如果目标字符串为空
AttributeError – 如果目标字符串引用无效路径或解析为不是 nn.Module 的内容

property shape¶: 等效于 batch_size。

share_memory() → T¶: 见 torch.Tensor.share_memory_()。

property specs: CompositeSpec¶

返回一个包含所有环境的 Composite 容器。

此功能允许您创建一个环境，在单个数据容器中检索所有规范，然后从工作区中删除环境。

state_dict(*args, **kwargs)[source]¶

返回一个字典，其中包含对模块的整个状态的引用。

包括参数和持久缓冲区（例如运行平均值）。键是相应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

注意

返回的对象是一个浅层副本。它包含对模块的参数和缓冲区的引用。

警告

当前 state_dict() 也接受 destination、prefix 和 keep_vars 的位置参数。但是，这将被弃用，并且在将来的版本中将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数：

destination (dict, optional) – 如果提供，模块的状态将更新到字典中，并将返回同一个对象。否则，将创建一个 OrderedDict 并返回。默认值：None。
prefix (str, optional) – 添加到参数和缓冲区名称之前的后缀，以构成 state_dict 中的键。默认值：''。
keep_vars (bool, optional) – 默认情况下，state_dict 中返回的 Tensor s 与自动梯度分离。如果将其设置为 True，则不会执行分离。默认值：False。

返回值:

包含模块的整个状态的字典

返回类型:

dict

示例

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

property state_keys: List[NestedKey]¶

环境的状态键。

默认情况下，将只有一个名为“state”的键。

键按数据树中的深度排序。

property state_spec: CompositeSpec¶

状态规范。

必须是 torchrl.data.CompositeSpec 实例。此处列出的键应作为输入与动作一起提供给环境。

在 TorchRL 中，即使它们不是严格意义上的“状态”，所有输入到环境的，且不是动作的，都被存储在 state_spec 中。

因此，"state_spec" 应该被认为是环境输入的通用数据容器，而不是动作数据。

示例

>>> from torchrl.envs import BraxEnv
>>> for envname in BraxEnv.available_envs:
...     break
>>> env = BraxEnv(envname)
>>> env.state_spec
CompositeSpec(
    state: CompositeSpec(
        pipeline_state: CompositeSpec(
            q: UnboundedContinuousTensorSpec(
                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([]))

step(tensordict: TensorDictBase) → TensorDictBase¶

在环境中执行一步。

Step 接受一个参数 tensordict，它通常包含一个表示要采取的动作的“action”键。Step 将调用一个非就地私有方法 _step，它是 EnvBase 子类需要重写的方法。

参数：: tensordict (TensorDictBase) – 包含要采取的动作的 Tensordict。如果输入 tensordict 包含一个 "next" 条目，则其中包含的值将优先于新计算的值。这提供了一种覆盖底层计算机制。
返回值:: 修改后的输入 tensordict，其中包含生成的观察结果、完成状态和奖励（+ 其他需要的信息）。

step_and_maybe_reset(tensordict: TensorDictBase) → Tuple[TensorDictBase, TensorDictBase]¶

在环境中执行一步，并在需要时（部分）重置环境。

参数：: tensordict (TensorDictBase) – step() 方法的输入数据结构。

此方法允许轻松编写无停止的回滚函数。

示例

>>> 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)

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

移动和/或转换参数和缓冲区。

这可以被调用为

to(device=None, dtype=None, non_blocking=False)[source]

to(dtype, non_blocking=False)[source]

to(tensor, non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

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

请参阅以下示例。

注意

此方法在原地修改模块。

参数：

device (torch.device) – 此模块中参数和缓冲区的所需设备
dtype (torch.dtype) – 此模块中参数和缓冲区的所需浮点或复数数据类型
tensor (torch.Tensor) – 数据类型和设备是此模块中所有参数和缓冲区的所需数据类型和设备的 Tensor
memory_format (torch.memory_format) – 此模块中 4D 参数和缓冲区的所需内存格式（仅限关键字参数）

返回值:

self

返回类型:

Module

示例

>>> # 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[device, str, int]], recurse: bool = True) → T¶

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

参数：

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

返回值:

self

返回类型:

Module

train(mode: bool = True) → T¶

将模块设置为训练模式。

这仅对某些模块有效。有关特定模块在训练/评估模式下的行为细节，如果它们受到影响，例如 Dropout、BatchNorm 等，请参阅相关文档。

参数：: mode (bool) – 是否设置训练模式 (True) 或评估模式 (False)。默认值：True。
返回值:: self
返回类型:: Module

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

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

注意

此方法在原地修改模块。

参数：: dst_type (type or string) – 目标类型
返回值:: self
返回类型:: Module

update_kwargs(kwargs: Union[dict, List[dict]]) → None¶

根据字典或字典列表更新每个环境的 kwargs。

参数：: kwargs (dict or list of dict) – 要与环境一起使用的新的 kwargs

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

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

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

注意

此方法在原地修改模块。

参数：: device (int, optional) – if specified, all parameters will be copied to that device
返回值:: self
返回类型:: Module

zero_grad(set_to_none: bool = True) → None¶

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

有关更多上下文，请参见 torch.optim.Optimizer 中的类似函数。

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

SerialEnv¶

文档

教程

资源