step_mdp¶

torchrl.envs.utils.step_mdp(tensordict: TensorDictBase, next_tensordict: TensorDictBase = None, keep_other: bool = True, exclude_reward: bool = True, exclude_done: bool = False, exclude_action: bool = True, reward_keys: NestedKey | list[NestedKey] = 'reward', done_keys: NestedKey | list[NestedKey] = 'done', action_keys: NestedKey | list[NestedKey] = 'action') → TensorDictBase[源码]¶

创建一个新的 TensorDict，反映输入 TensorDict 的时间步进。

给定一个在步进后获取的 TensorDict，返回带 "next" 索引的 TensorDict。参数允许精确控制应保留哪些内容以及应从 "next" 条目中复制哪些内容。默认行为是：将观测条目、奖励和完成状态移动到根目录，排除当前动作，并保留所有额外键（非动作、非完成、非奖励）。

参数：

tensordict (TensorDictBase) – 包含要重命名键的 TensorDict。
next_tensordict (TensorDictBase, optional) – 目标 TensorDict。如果为 None，则创建一个新的 TensorDict。
keep_other (bool, optional) – 如果为 True，所有不以 'next_' 开头的键都将被保留。默认为 True。
exclude_reward (bool, optional) – 如果为 True，"reward" 键将从结果 TensorDict 中丢弃。如果为 False，它将从 "next" 条目中复制（并替换）（如果存在）。默认为 True。
exclude_done (bool, optional) – 如果为 True，"done" 键将从结果 TensorDict 中丢弃。如果为 False，它将从 "next" 条目中复制（并替换）（如果存在）。默认为 False。
exclude_action (bool, optional) – 如果为 True，"action" 键将从结果 TensorDict 中丢弃。如果为 False，它将保留在根 TensorDict 中（因为它不应该存在于 "next" 条目中）。默认为 True。
reward_keys (NestedKey 或 NestedKey 列表, optional) – 写入奖励的键。默认为 “reward”。
done_keys (NestedKey 或 NestedKey 列表, optional) – 写入完成状态的键。默认为 “done”。
action_keys (NestedKey 或 NestedKey 列表, optional) – 写入动作的键。默认为 “action”。

返回：

一个新的 TensorDict（如果提供了 next_tensordict，则为该 TensorDict），包含 t+1 步的张量。

返回类型：

TensorDictBase

另请参阅

EnvBase.step_mdp() 是此自由函数的基于类别的版本。它将尝试缓存键值以减少在 MDP 中执行一步的开销。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict({
...     "done": torch.zeros((), dtype=torch.bool),
...     "reward": torch.zeros(()),
...     "extra": torch.zeros(()),
...     "next": TensorDict({
...         "done": torch.zeros((), dtype=torch.bool),
...         "reward": torch.zeros(()),
...         "obs": torch.zeros(()),
...     }, []),
...     "obs": torch.zeros(()),
...     "action": torch.zeros(()),
... }, [])
>>> print(step_mdp(td))
TensorDict(
    fields={
        done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
        extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(step_mdp(td, exclude_done=True))  # "done" is dropped
TensorDict(
    fields={
        extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(step_mdp(td, exclude_reward=False))  # "reward" is kept
TensorDict(
    fields={
        done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
        extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(step_mdp(td, exclude_action=False))  # "action" persists at the root
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
        extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(step_mdp(td, keep_other=False))  # "extra" is missing
TensorDict(
    fields={
        done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
        obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

警告

如果奖励键在被排除时也是输入键的一部分，此函数将无法正常工作。这就是为什么 RewardSum 变换默认将回合奖励注册到观测中而不是奖励规范中。使用此函数的快速缓存版本 (_StepMDP) 时，不应出现此问题。

step_mdp¶

文档

教程

资源