step_mdp

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

创建一个新的 tensordict，它反映了输入 tensordict 在时间上的一个步骤。

给定一个在执行一个步骤后检索到的 tensordict，返回 "next" 索引的 tensordict。参数允许对应该保留什么以及应该从 "next" 条目中复制什么进行精确控制。默认行为是：移动观察条目、奖励和完成状态到根目录，排除当前动作并保留所有额外键（非动作、非完成、非奖励）。

参数:

• tensordict (TensorDictBase) – 具有要重命名的键的 tensordict

• next_tensordict (TensorDictBase, 可选) – 目标 tensordict

• keep_other (bool, 可选) – 如果 True，则将保留所有不以 'next_' 开头的键。默认值为 True。

• exclude_reward (bool, 可选) – 如果 True，则 "reward" 键将从结果 tensordict 中丢弃。如果 False，它将从 "next" 条目（如果存在）中复制（并替换）。默认值为 True。

• exclude_done (bool, 可选) – 如果 True，则 "done" 键将从结果 tensordict 中丢弃。如果 False，它将从 "next" 条目（如果存在）中复制（并替换）。默认值为 False。

• exclude_action (bool, 可选) – 如果 True，则 "action" 键将从结果 tensordict 中丢弃。如果 False，它将保留在根 tensordict 中（因为它不应出现在 "next" 条目中）。默认值为 True。

• reward_keys (NestedKey 或 NestedKey 列表, 可选) – 奖励写入的键。默认为“reward”。

• done_keys (NestedKey 或 NestedKey 列表, 可选) – 完成写入的键。默认为“done”。

• action_keys (NestedKey 或 NestedKey 列表, 可选) – 动作写入的键。默认为“action”。

返回值:

包含 t+1 步的张量的新的 tensordict（或 next_tensordict）。

示例：此函数允许使用这种循环

>>> 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）时，不应观察到此问题。