KLRewardTransform

class torchrl.envs.transforms.KLRewardTransform(actor: ProbabilisticTensorDictModule, coef=1.0, in_keys=None, out_keys=None, requires_grad=False)

一個轉換，用於將 KL[pi_current||pi_0] 校正項添加到獎勵中。

此轉換用於約束策略使其保持接近其原始配置，這限制了在使用 RLHF 進行微調時的過度擬合。

參數:

• actor (ProbabilisticTensorDictModule) – 一個機率性 actor。它必須具有以下特徵：它必須有一組輸入 (in_keys) 和輸出鍵 (out_keys)。它必須有一個 get_dist 方法，該方法輸出動作的分配。

• coef (float) – KL 項的係數。預設為 1.0。

• in_keys (str 或 str/str 組成的列表/元組) – 應提取獎勵的輸入鍵。預設為 "reward"。

• out_keys (str 或 str/str 組成的列表/元組) – 應寫入獎勵的輸出鍵。預設為 "reward"。

• requires_grad (bool, optional) – 如果 True，則凍結的參數將由原始參數的可微克隆組成。預設為 False。

注意

如果參數不可微（預設），則在調用 dtype 或設備轉換操作時（例如 cuda(), to() 等），它們不會遵循模組。當 requires_grad=True 時，轉換操作將按預期工作。

範例

>>> from torchrl.envs.libs.gym import GymEnv
>>> from torchrl.envs import TransformedEnv
>>> from tensordict.nn import TensorDictModule as Mod, NormalParamExtractor
>>> from torchrl.modules import ProbabilisticActor
>>> from tensordict import TensorDict
>>> from torchrl.modules.distributions import TanhNormal
>>> from torch import nn
>>> base_env = GymEnv("Pendulum-v1")
>>> n_obs = base_env.observation_spec["observation"].shape[-1]
>>> n_act = base_env.action_spec.shape[-1]
>>> module = Mod(
...     nn.Sequential(nn.Linear(n_obs, n_act * 2), NormalParamExtractor()),
...     in_keys=["observation"],
...     out_keys=["loc", "scale"],
... )
>>> actor = ProbabilisticActor(
...     module,
...     in_keys=["loc", "scale"],
...     distribution_class=TanhNormal,
...     return_log_prob=True,
... )
>>> transform = KLRewardTransform(actor, out_keys="reward_kl")
>>> env = TransformedEnv(base_env, transform)
>>> with torch.no_grad():
...     # modify the actor parameters
...     _ = TensorDict(dict(actor.named_parameters()), []).apply_(lambda x: x.data.copy_(x.data + 1))
...     td = env.rollout(3, actor)
>>> # check that rewards have been modified
>>> assert (td.get(("next", "reward")) != td.get(("next", "reward_kl"))).all()

注意

由於 KL 公式並非總是可用，並且原始分佈的參數可能未被記錄，因此我們使用 KL 散度的隨機估計。

forward(tensordict: TensorDictBase) → TensorDictBase

讀取輸入的 tensordict，並針對選定的鍵應用轉換。

transform_output_spec(output_spec: Composite) → Composite

轉換輸出規格，使產生的規格符合轉換映射。

這個方法通常應該保持不動。 更改應該使用 transform_observation_spec()、 transform_reward_spec() 和 transformfull_done_spec() 來實作。 :param output_spec: 轉換前的規格 :type output_spec: TensorSpec

回傳:

轉換後預期的規格