RewardSum¶
- class torchrl.envs.transforms.RewardSum(in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = None, reset_keys: Sequence[NestedKey] | None = None, *, reward_spec: bool = False)[源代码]¶
跟踪每轮的累积奖励。
此转换器接受 tensordict 奖励键的列表(即 'in_keys'),并在每轮的时间维度上跟踪它们的累积值。
调用时,转换器会为每个 `in_key` 写入一个新的 tensordict 条目,命名为 `episode_{in_key}`,其中包含累积值。
- 参数:
in_keys(NestedKeys 列表,可选)– 输入奖励键。所有 'in_keys' 都应是环境 reward_spec 的一部分。如果未指定 `in_keys`,则此转换器假定 'reward' 是输入键。但是,也可以指定多个奖励(例如,“reward1” 和 “reward2”)。
out_keys(NestedKeys 列表,可选)– 输出总和键,每个输入键应有一个。
reset_keys(NestedKeys 列表,可选)– 要使用的 reset_keys 列表,如果找不到父环境。如果提供,此值将优先于环境 `reset_keys`。
- 关键字参数:
reward_spec(bool,可选)– 如果为 `True`,则新的奖励条目将在 reward specs 中注册。默认为 `False`(在 `observation_specs` 中注册)。
示例
>>> from torchrl.envs.transforms import RewardSum, TransformedEnv >>> from torchrl.envs.libs.gym import GymEnv >>> env = TransformedEnv(GymEnv("CartPole-v1"), RewardSum()) >>> env.set_seed(0) >>> torch.manual_seed(0) >>> td = env.reset() >>> print(td["episode_reward"]) tensor([0.]) >>> td = env.rollout(3) >>> print(td["next", "episode_reward"]) tensor([[1.], [2.], [3.]])
- forward(tensordict: TensorDictBase) TensorDictBase [源代码]¶
读取输入 tensordict,并对选定的键应用转换。
默认情况下,此方法
直接调用
_apply_transform()
。不调用
_step()
或_call()
。
此方法不会在任何时候在 env.step 中调用。但是,它会在
sample()
中调用。注意
forward
也可以使用dispatch
将参数名称转换为键,并使用常规关键字参数。示例
>>> class TransformThatMeasuresBytes(Transform): ... '''Measures the number of bytes in the tensordict, and writes it under `"bytes"`.''' ... def __init__(self): ... super().__init__(in_keys=[], out_keys=["bytes"]) ... ... def forward(self, tensordict: TensorDictBase) -> TensorDictBase: ... bytes_in_td = tensordict.bytes() ... tensordict["bytes"] = bytes ... return tensordict >>> t = TransformThatMeasuresBytes() >>> env = env.append_transform(t) # works within envs >>> t(TensorDict(a=0)) # Works offline too.
- transform_input_spec(input_spec: TensorSpec) TensorSpec [源代码]¶
转换输入规范,使结果规范与转换映射匹配。
- 参数:
input_spec (TensorSpec) – 转换前的规范
- 返回:
转换后的预期规范
- transform_observation_spec(observation_spec: TensorSpec) TensorSpec [源代码]¶
转换 observation spec,添加 RewardSum 生成的新键。
- transform_reward_spec(reward_spec: TensorSpec) TensorSpec [源代码]¶
转换奖励的 spec,使其与变换映射匹配。
- 参数:
reward_spec (TensorSpec) – 变换前的 spec
- 返回:
转换后的预期规范