QValueModule

class torchrl.modules.tensordict_module.QValueModule(*args, **kwargs)

Q 值 TensorDictModule，用于 Q 值策略。

此模块处理包含动作值的张量，将其转换为其 argmax 分量（即，由此产生的贪婪动作），遵循给定的动作空间（one-hot、binary 或 categorical）。它可以与 tensordict 和常规张量一起使用。

参数:

• action_space (str, optional) – 动作空间。必须是以下之一："one-hot"、"mult-one-hot"、"binary" 或 "categorical"。此参数与 spec 互斥，因为 spec 会条件化 action_space。

• action_value_key (str 或 tuple of str, optional) – 表示动作值的输入键。默认为 "action_value"。

• action_mask_key (str 或 tuple of str, optional) – 表示动作掩码的输入键。默认为 "None"（等同于无掩码）。

• out_keys (list of str 或 tuple of str, optional) – 表示动作、动作值和选定动作值的输出键。默认为 ["action", "action_value", "chosen_action_value"]。

• var_nums (int, optional) – 如果 action_space = "mult-one-hot"，此值表示每个动作分量的基数。

• spec (TensorSpec, optional) – 如果提供，则为动作（和/或其它输出）的规范。这与 action_space 互斥，因为 spec 会条件化 action_space。

• safe (bool) – 如果为 True，则输出值会根据输入 spec 进行检查。由于探索策略或数值下溢/上溢问题，可能会出现域外采样。如果此值超出范围，它将使用 TensorSpec.project 方法投影回期望的空间。默认为 False。

返回:

如果输入是单个张量，则会返回一个包含选定动作、值和选定动作值的元组。如果提供了 tensordict，则会使用 out_keys 字段中指定的键更新该 tensordict。

示例

>>> from tensordict import TensorDict
>>> action_space = "categorical"
>>> action_value_key = "my_action_value"
>>> actor = QValueModule(action_space, action_value_key=action_value_key)
>>> # This module works with both tensordict and regular tensors:
>>> value = torch.zeros(4)
>>> value[-1] = 1
>>> actor(my_action_value=value)
(tensor(3), tensor([0., 0., 0., 1.]), tensor([1.]))
>>> actor(value)
(tensor(3), tensor([0., 0., 0., 1.]), tensor([1.]))
>>> actor(TensorDict({action_value_key: value}, []))
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        action_value: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
        chosen_action_value: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        my_action_value: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

forward(tensordict: Tensor = None) → TensorDictBase

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

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

注意

虽然前向传播的实现需要在该函数中定义，但用户应该在之后调用 Module 实例而不是此函数，因为前者会负责运行已注册的钩子，而后者会静默忽略它们。