VecNormV2¶
- class torchrl.envs.transforms.VecNormV2(in_keys: Sequence[NestedKey], out_keys: Sequence[NestedKey] | None = None, *, lock: mp.Lock = None, stateful: bool = True, decay: float = 0.9999, eps: float = 0.0001, shared_data: TensorDictBase | None = None, reduce_batch_dims: bool = False)[source]¶
用于归一化强化学习环境中向量化观测值和奖励的类。
VecNormV2 可以以有状态或无状态模式运行。在有状态模式下,它维护内部统计数据(均值和方差)来归一化输入。在无状态模式下,它需要提供外部统计数据来进行归一化。
注意
此类设计为
VecNorm
的几乎即插即用替代品。它不应直接构造,而应使用 new_api=True 关键字参数通过VecNorm
变换来构造。在 v0.10 中,VecNorm
变换将默认切换到新 API。- 有状态 vs. 无状态
有状态模式 (stateful=True)
维护内部统计数据(loc、var、count)以进行归一化。
除非被冻结,否则每次调用都会更新统计数据。
state_dict 返回当前统计数据。
load_state_dict 使用提供的状态更新内部统计数据。
无状态模式 (stateful=False)
需要提供外部统计数据来进行归一化。
不维护或更新内部统计数据。
state_dict 返回一个空字典。
load_state_dict 不影响内部状态。
- 参数:
in_keys (Sequence[NestedKey]) – 要归一化的数据的输入键。
out_keys (Sequence[NestedKey] | None) – 归一化数据的输出键。如果未提供,则默认为 in_keys。
lock (mp.Lock, optional) – 用于线程安全的锁。
stateful (bool, optional) – VecNorm 是否有状态。此变换的无状态版本需要将数据包含在输入/输出 tensordicts 中。默认为 True。
decay (float, optional) – 更新统计数据的衰减率。默认为 0.9999。如果使用 decay=1,归一化统计数据具有无限内存(每个项目的权重相同)。较低的值会使近期数据比旧数据更重要。
eps (float, optional) – 防止除零的微小值。默认为 1e-4。
shared_data (TensorDictBase | None, optional) – 用于初始化的共享数据。默认为 None。
reduce_batch_dims (bool, optional) – 如果为 True,则通过平均数据来减小批次维度,然后再更新统计数据。这在接收批次样本时很有用,因为它允许移动平均在整个批次上计算,而不是在单个元素上计算。请注意,此选项仅在有状态模式(stateful=True)下支持。默认为 False。
- 变量:
stateful (bool) – 指示 VecNormV2 是有状态还是无状态。
lock (mp.Lock) – 用于确保更新统计数据时的线程安全的进程锁。
decay (float) – 更新统计数据的衰减率。
eps (float) – 在归一化过程中防止除零的微小值。
frozen (bool) – 指示 VecNormV2 是否已冻结,防止统计数据被更新。
_cast_int_to_float (bool) – 指示是否将整数输入转换为浮点数。
- loc()¶
返回用于归一化的位置(均值)。
- scale()¶
返回用于归一化的尺度(标准差)。
- standard_normal()¶
指示归一化是否遵循标准正态分布。
状态字典行为
在有状态模式下,state_dict 返回一个包含当前 loc、var 和 count 的字典。这些可用于在进程之间共享张量(此方法由
VecNorm
自动触发,以在进程之间共享 VecNorm 状态)。在无状态模式下,state_dict 返回一个空字典,因为没有维护内部状态。
加载状态字典行为
在有状态模式下,load_state_dict 使用提供的状态更新内部的 loc、var 和 count。
在无状态模式下,load_state_dict 不会修改任何内部状态,因为没有需要更新的状态。
另请参阅
此变换的第一个版本
VecNorm
。示例
>>> import torch >>> from torchrl.envs import EnvCreator, GymEnv, ParallelEnv, SerialEnv, VecNormV2 >>> >>> torch.manual_seed(0) >>> env = GymEnv("Pendulum-v1") >>> env_trsf = env.append_transform( >>> VecNormV2(in_keys=["observation", "reward"], out_keys=["observation_norm", "reward_norm"]) >>> ) >>> r = env_trsf.rollout(10) >>> print("Unnormalized rewards", r["next", "reward"]) Unnormalized rewards tensor([[ -1.7967], [ -2.1238], [ -2.5911], [ -3.5275], [ -4.8585], [ -6.5028], [ -8.2505], [-10.3169], [-12.1332], [-13.1235]]) >>> print("Normalized rewards", r["next", "reward_norm"]) Normalized rewards tensor([[-1.6596e-04], [-8.3072e-02], [-1.9170e-01], [-3.9255e-01], [-5.9131e-01], [-7.4671e-01], [-8.3760e-01], [-9.2058e-01], [-9.3484e-01], [-8.6185e-01]]) >>> # Aggregate values when using batched envs >>> env = SerialEnv(2, [lambda: GymEnv("Pendulum-v1")] * 2) >>> env_trsf = env.append_transform( >>> VecNormV2( >>> in_keys=["observation", "reward"], >>> out_keys=["observation_norm", "reward_norm"], >>> # Use reduce_batch_dims=True to aggregate values across batch elements >>> reduce_batch_dims=True, ) >>> ) >>> r = env_trsf.rollout(10) >>> print("Unnormalized rewards", r["next", "reward"]) Unnormalized rewards tensor([[[-0.1456], [-0.1862], [-0.2053], [-0.2605], [-0.4046], [-0.5185], [-0.8023], [-1.1364], [-1.6183], [-2.5406]],
- [[-0.0920],
[-0.1492], [-0.2702], [-0.3917], [-0.5001], [-0.7947], [-1.0160], [-1.3347], [-1.9082], [-2.9679]]])
>>> print("Normalized rewards", r["next", "reward_norm"]) Normalized rewards tensor([[[-0.2199], [-0.2918], [-0.1668], [-0.2083], [-0.4981], [-0.5046], [-0.7950], [-0.9791], [-1.1484], [-1.4182]],
- [[ 0.2201],
[-0.0403], [-0.5206], [-0.7791], [-0.8282], [-1.2306], [-1.2279], [-1.2907], [-1.4929], [-1.7793]]])
>>> print("Loc / scale", env_trsf.transform.loc["reward"], env_trsf.transform.scale["reward"]) Loc / scale tensor([-0.8626]) tensor([1.1832]) >>> >>> # Share values between workers >>> def make_env(): ... env = GymEnv("Pendulum-v1") ... env_trsf = env.append_transform( ... VecNormV2(in_keys=["observation", "reward"], out_keys=["observation_norm", "reward_norm"]) ... ) ... return env_trsf ... ... >>> if __name__ == "__main__": ... # EnvCreator will share the loc/scale vals ... make_env = EnvCreator(make_env) ... # Create a local env to track the loc/scale ... local_env = make_env() ... env = ParallelEnv(2, [make_env] * 2) ... r = env.rollout(10) ... # Non-zero loc and scale testify that the sub-envs share their summary stats with us ... print("Remotely updated loc / scale", local_env.transform.loc["reward"], local_env.transform.scale["reward"]) Remotely updated loc / scale tensor([-0.4307]) tensor([0.9613]) ... env.close()
- freeze() VecNormV2 [source]¶
冻结 VecNorm,避免在调用时更新统计数据。
请参阅
unfreeze()
。
- get_extra_state() OrderedDict [source]¶
返回要包含在模块 state_dict 中的任何额外状态。
如果您的模块需要存储额外状态,请实现此函数以及对应的
set_extra_state()
。在构建模块的 state_dict() 时会调用此函数。请注意,额外状态应该是可拾取的,以确保 state_dict 的有效序列化。我们仅为序列化张量提供向后兼容性保证;其他对象的序列化拾取形式如果发生更改,可能会破坏向后兼容性。
- 返回:
要存储在模块 state_dict 中的任何额外状态
- 返回类型:
对象
- property loc¶
返回一个 TensorDict,其中包含可用于仿射变换的 loc。
- property scale¶
返回一个 TensorDict,其中包含可用于仿射变换的 scale。
- set_extra_state(state: OrderedDict) None [source]¶
设置加载的 state_dict 中包含的额外状态。
此函数从
load_state_dict()
调用,以处理 state_dict 中找到的任何额外状态。如果您的模块需要在其 state_dict 中存储额外状态,请实现此函数以及对应的get_extra_state()
。- 参数:
state (dict) – 来自 state_dict 的额外状态
- property standard_normal¶
由 loc 和 scale 提供的仿射变换是否遵循标准正态方程。
类似于
ObservationNorm
的 standard_normal 属性。始终返回
True
。
- transform_observation_spec(observation_spec: Composite) Composite [source]¶
转换观察规范,使结果规范与转换映射匹配。
- 参数:
observation_spec (TensorSpec) – 转换前的规范
- 返回:
转换后的预期规范
- transform_output_spec(output_spec: Composite) Composite [source]¶
转换输出规范,使结果规范与转换映射匹配。
此方法通常应保持不变。更改应通过
transform_observation_spec()
、transform_reward_spec()
和transform_full_done_spec()
来实现。 :param output_spec: 变换前的规范 :type output_spec: TensorSpec- 返回:
转换后的预期规范
- transform_reward_spec(reward_spec: Composite, observation_spec) Composite [source]¶
转换奖励的 spec,使其与变换映射匹配。
- 参数:
reward_spec (TensorSpec) – 变换前的 spec
- 返回:
转换后的预期规范