torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/envs/llm/transforms/kl.py

@@ -0,0 +1,1544 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import warnings
+
+from collections.abc import Callable
+from contextlib import nullcontext
+from copy import copy
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+from tensordict import NestedKey, set_list_to_stack, TensorDictBase, unravel_key
+from tensordict.utils import _zip_strict, is_seq_of_nested_key, logger as torchrl_logger
+from torch.nn.utils.rnn import pad_sequence
+from torchrl.data import Composite, Unbounded
+from torchrl.data.tensor_specs import DEVICE_TYPING
+from torchrl.envs import EnvBase, Transform
+from torchrl.envs.transforms.ray_service import _RayServiceMetaClass, RayTransform
+from torchrl.envs.transforms.transforms import Compose
+from torchrl.envs.transforms.utils import _set_missing_tolerance
+from torchrl.modules.llm.policies.common import LLMWrapperBase
+
+if TYPE_CHECKING:
+    import transformers
+
+
+class RayKLRewardTransform(RayTransform):
+    """A Ray-based implementation of :class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`.
+
+    This class creates a Ray remote actor from KLRewardTransform that can be shared across multiple workers.
+    All method calls are delegated to the remote actor, ensuring that multiple environments can
+    share the same KL computation resources.
+
+    To avoid serialization issues with large models, this class supports model factories
+    that create models on the remote actor rather than passing full models through Ray channels.
+
+    Args:
+        ref_model (LLMWrapperBase, optional): the reference model. Prefer using a model factory instead
+            to avoid serialization issues.
+
+    Keyword Args:
+        ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
+            This allows for explicit resource control and avoids serialization issues.
+        num_cpus (int, optional): Number of CPUs to allocate to the Ray actor. Defaults to 1.
+        num_gpus (int, optional): Number of GPUs to allocate to the Ray actor. Defaults to 0.
+        device (torch.device, optional): Device to use on the remote Ray actor for tensor operations.
+            The local Ray transform will handle CPU serialization and device restoration automatically.
+            Defaults to None.
+        actor_name (str, optional): Name of the Ray actor to use. If provided, the actor will be reused if it already exists.
+        **kwargs: Additional keyword arguments to pass to KLRewardTransform.
+
+    Note:
+        When using model factories, the corresponding model argument (ref_model) should be None.
+        Model factories are preferred for large models to avoid serialization overhead.
+
+    Examples:
+        >>> # Option 1: Using model factory for explicit resource control
+        >>> def create_ref_model():
+        ...     return TransformersWrapper(ref_model, tokenizer=tokenizer, generate=False, return_log_probs=True)
+        >>> transform = RayKLRewardTransform(
+        ...     ref_model=None,
+        ...     ref_model_factory=create_ref_model,
+        ...     num_gpus=1,
+        ...     device=torch.device("cuda")
+        ... )
+
+        >>> # Option 2: Pass model directly (Ray handles serialization)
+        >>> transform = RayKLRewardTransform(ref_model=ref_model, device=torch.device("cuda"))
+    """
+
+    def __init__(
+        self,
+        ref_model: LLMWrapperBase | None = None,
+        *,
+        ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int = 0,
+        device: DEVICE_TYPING | None = None,
+        actor_name: str | None = None,
+        **kwargs,
+    ):
+        # Validate arguments: model and factory should not both be provided
+        if ref_model is not None and ref_model_factory is not None:
+            raise ValueError(
+                "Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
+            )
+        if ref_model is None and ref_model_factory is None:
+            raise ValueError(
+                "Must provide exactly one of 'ref_model' or 'ref_model_factory'."
+            )
+
+        # Store creation parameters for actor creation
+        self._ref_model = ref_model
+        self._ref_model_factory = ref_model_factory
+        self._creation_kwargs = kwargs
+        # Store device separately for passing to remote actor
+        self._remote_device = device
+
+        # Default num_cpus
+        if num_cpus is None:
+            num_cpus = 1
+
+        # Call parent constructor without device (Ray transform handles CPU/device mapping)
+        super().__init__(
+            num_cpus=num_cpus,
+            num_gpus=num_gpus,
+            device=None,  # Don't store device locally
+            actor_name=actor_name,
+            **kwargs,
+        )
+
+    def _create_actor(self, **kwargs):
+        """Create the remote KLRewardTransform actor."""
+        # Create the remote KLRewardTransform with resource specifications
+        RemoteKLRewardTransform = self._ray.remote(
+            num_cpus=self._num_cpus, num_gpus=self._num_gpus
+        )(KLRewardTransform)
+
+        if self._actor_name is not None:
+            RemoteKLRewardTransform = RemoteKLRewardTransform.options(
+                name=self._actor_name
+            )
+
+        # Determine how to create model on the remote actor
+        ref_model_arg = self._ref_model
+
+        # If we have factory, we'll pass it and set model to None
+        creation_kwargs = self._creation_kwargs.copy()
+        if self._ref_model_factory is not None:
+            creation_kwargs["ref_model_factory"] = self._ref_model_factory
+            ref_model_arg = None
+
+        # Pass device to the remote actor
+        if self._remote_device is not None:
+            creation_kwargs["device"] = self._remote_device
+
+        # Create the shared actor
+        actor = RemoteKLRewardTransform.remote(
+            ref_model=ref_model_arg, **creation_kwargs
+        )
+
+        return actor
+
+    def __repr__(self):
+        """String representation."""
+        try:
+            if hasattr(self, "_actor") and self._actor is not None:
+                return self._ray.get(self._actor.__repr__.remote())
+            else:
+                return "RayKLRewardTransform(actor=None)"
+        except Exception:
+            return f"RayKLRewardTransform(actor={getattr(self, '_actor', 'None')})"
+
+
+class KLRewardTransform(Transform, metaclass=_RayServiceMetaClass):
+    """A legacy transform for computing KL divergence-based rewards.
+
+    **Deprecated**: This transform is maintained for backward compatibility but is no longer
+    the recommended approach. Use :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` instead,
+    which provides better modularity and integration with the new wrapper design.
+
+    **Recent Changes:**
+
+    - **Legacy Status**: This transform is now considered legacy and may not work optimally
+      with the new modular wrapper design.
+    - **ChatHistory Integration**: Limited support for the new :class:`~torchrl.modules.llm.policies.ChatHistory` objects.
+    - **Input Mode Support**: May not handle all input modes (`"history"`, `"text"`, `"tokens"`) consistently.
+
+    **Recommendation**:
+    Use :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` for new code, which provides:
+    - Better integration with the new wrapper design
+    - Consistent support for all input modes
+    - Proper handling of ChatHistory objects
+    - More modular and composable architecture
+
+    Args:
+        ref_model (LLMWrapperBase): the reference model.
+
+    Keyword Args:
+        ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
+        assistant_only (bool): whether to only compute KL on assistant tokens. Defaults to `True`.
+        tokenizer (transformers.AutoTokenizer): the tokenizer to use. Defaults to `None`.
+        detach (bool): whether to detach the KL from the computation graph. Defaults to `True`.
+        device (torch.device): the device to cast the tensors to. This is not the device of the specs, but the device
+            onto which the tensors will be moved. It allows to keep the model on a different device
+            than the upcoming data. When using Ray service, this device will be used on the remote actor.
+            Defaults to `None`.
+        padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
+        use_ray_service (bool, optional): whether to use Ray service. Defaults to `False`.
+        actor_name (str, optional): the name of the Ray actor to use. Defaults to `None`.
+
+    Examples:
+        >>> # Legacy usage (not recommended for new code)
+        >>> transform = KLRewardTransform(gen_model, ref_model)
+        >>>
+        >>> # Recommended approach using RetrieveKL
+        >>> from torchrl.envs.llm.transforms.kl import RetrieveKL
+        >>> transform = RetrieveKL(gen_model, ref_model, assistant_only=True)
+
+    .. seealso::
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: The recommended transform for KL divergence computation.
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: Base transform for retrieving log-probabilities.
+        :class:`~torchrl.envs.llm.transforms.kl.KLComputation`: Transform for computing KL divergence between log-prob tensors.
+    """
+
+    DEFAULT_IN_KEYS = ["reward"]
+    _RayServiceClass = RayKLRewardTransform
+
+    def __init__(
+        self,
+        ref_model: LLMWrapperBase | None = None,
+        *,
+        ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        coef=1.0,
+        in_keys=None,
+        out_keys=None,
+        log_prob_key: NestedKey = ("log_probs", "full"),
+        device: torch.device | None = None,
+        add_to_reward: bool = True,
+        tokenizer: transformers.AutoTokenizer | None = None,
+        assistant_only: bool = True,
+        padding_side: str = "left",
+        use_ray_service: bool = False,
+    ):
+        # Handle model factory - create model if factory is provided
+        if ref_model_factory is not None:
+            if ref_model is not None:
+                raise ValueError(
+                    "Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
+                )
+            ref_model = ref_model_factory()
+        elif ref_model is None:
+            raise ValueError(
+                "Must provide exactly one of 'ref_model' or 'ref_model_factory'."
+            )
+
+        if in_keys is None:
+            in_keys = self.DEFAULT_IN_KEYS
+        if out_keys is None:
+            out_keys = copy(in_keys)
+        if len(out_keys) == len(in_keys):
+            out_keys = out_keys + ["kl_penalty", "ref_log_probs"]
+        elif len(out_keys) != len(in_keys) + 2:
+            raise ValueError(
+                "The out_keys must have the same length as the in_keys (plus two additional optional kl entries for logging)."
+            )
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+        if not is_seq_of_nested_key(self.in_keys) or not is_seq_of_nested_key(
+            self.out_keys
+        ):
+            raise ValueError(
+                f"invalid in_keys / out_keys:\nin_keys={self.in_keys} \nout_keys={self.out_keys}"
+            )
+        if len(self.in_keys) != 1 or len(self.out_keys) != 3:
+            raise ValueError(
+                f"Only one in_key/out_key is allowed, got in_keys={self.in_keys}, out_keys={self.out_keys}."
+            )
+        self._out_keys = [unravel_key(out_key) for out_key in self._out_keys]
+
+        if getattr(ref_model, "generate", False):
+            raise ValueError(
+                "The actor is configured to generate text, not compute the log-probs."
+            )
+
+        # update the in_keys for dispatch etc
+        self.in_keys = self.in_keys + ref_model.in_keys
+        self.in_keys = [unravel_key(in_key) for in_key in self.in_keys]
+
+        self.add_to_reward = add_to_reward
+        # check that the model has parameters
+        self.__dict__["ref_model"] = ref_model
+
+        # self._buffers["actor_params"] = params.clone().detach()
+
+        self.device = device
+
+        # find the sample log-prob key
+        self.log_prob_full_key = log_prob_key
+
+        self._tokenizer = tokenizer
+        self.assistant_only = assistant_only
+        self.padding_side = padding_side
+
+        if not isinstance(coef, torch.Tensor):
+            coef = torch.as_tensor(coef)
+        self.register_buffer("coef", coef)
+        # sanity check for the ref_model
+        if not getattr(ref_model, "input_mode", "tokens") == "tokens":
+            raise ValueError(
+                "The ref_model must be configured to use tokens as input. Please set the `input_mode` argument to `tokens`."
+            )
+
+    @property
+    def pad_output(self):
+        # We need pad_output to match the pad_output of the inference model
+        return self.ref_model.pad_output
+
+    @property
+    def tokenizer(self):
+        tokenizer = self._tokenizer
+        if tokenizer is not None:
+            return tokenizer
+        try:
+            return self.ref_model.tokenizer
+        except AttributeError:
+            raise AttributeError(
+                "The ref_model does not have a tokenizer. Please pass the tokenizer to the constructor."
+            )
+
+    def set_container(self, container: Transform | EnvBase) -> None:
+        result = super().set_container(container)
+        if self.action_key is None:
+            parent = getattr(self, "parent", None)
+            if parent is not None:
+                action_keys = parent.action_keys
+                if len(action_keys) != 1:
+                    raise ValueError(
+                        f"More than one action_key found. Please pass the `action_key` argument directly to {type(self).__name__}."
+                    )
+                action_key = action_keys[0]
+                self.action_key = action_key
+        return result
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        with _set_missing_tolerance(self, True):
+            tensordict_reset = self._step(tensordict_reset, tensordict_reset)
+        return tensordict_reset
+
+    @property
+    def action_key(self) -> NestedKey:
+        # Get the action from the base env (a ChatEnv).
+        if self.parent.base_env.input_mode == "history":
+            return ("history", "full")
+        if self.parent.base_env.input_mode == "text":
+            return ("text", "full")
+        if self.parent.base_env.input_mode == "tokens":
+            return ("tokens", "full")
+        raise ValueError(f"Invalid input mode: {self.parent.base_env.input_mode}")
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        original_device = None
+        if self.device is not None:
+            original_device = tensordict.device
+            tensordict = tensordict.to(self.device)
+            next_tensordict = next_tensordict.to(self.device)
+        # tensordict = self._get_text_response(tensordict, next_tensordict)
+        response = tensordict.get(self.action_key, None)
+        if response is None:
+            if not self.missing_tolerance:
+                raise RuntimeError(
+                    f"Action with key {self.action_key} not found data {tensordict}"
+                )
+            # being called after reset or without action, skipping
+            if self.out_keys[0] != "reward" and self.parent is not None:
+                next_tensordict.set(self.out_keys[0], self.parent.reward_spec.zero())
+            return next_tensordict
+
+        # We use the ("tokens", "full") key to get the log-probs of the reference model
+        with torch.device(self.device) if self.device is not None else nullcontext():
+            td_input = tensordict.copy()
+            ref_log_prob_td = self.ref_model(td_input)
+        if self.pad_output:
+            ref_log_prob_padded = ref_log_prob_td.get(self.log_prob_full_key)
+        else:
+            ref_log_prob_unpadded = ref_log_prob_td.get(
+                self.log_prob_full_key, as_list=True  # type: ignore[misc]
+            )
+        if self.assistant_only:
+            # Get the assistant mask
+            mask = tensordict.get(("masks", "all_assistant_mask"))
+            # mask will often be None - fall back on prompt / response separation
+            if mask is None:
+                if self.pad_output:
+                    # simple case: just take the prompt length
+                    prompt_length = tensordict.get(("tokens", "prompt")).shape[-1]
+                    mask = tensordict.get(("masks", "all_attention_mask")).clone()
+                    mask[..., :prompt_length] = False
+                else:
+                    # simple case: just take the prompt length
+                    prompt_length = [
+                        t.size(-1)
+                        for t in tensordict.get(("tokens", "prompt"), as_list=True)  # type: ignore[misc]
+                    ]
+                    mask = tensordict.get(("masks", "all_attention_mask"), as_list=True)  # type: ignore[misc]
+                    for i in range(len(prompt_length)):
+                        mask[i] = mask[i].clone()
+                        mask[i][..., : prompt_length[i]] = False
+
+            # we want to keep the batch dimension
+            ref_log_prob_list = []
+            if self.pad_output:
+                for i in range(ref_log_prob_padded.size(0)):
+                    ref_log_prob_list.append(
+                        ref_log_prob_padded[i].masked_fill(~mask[i], 0)
+                    )
+            else:
+                for i in range(len(ref_log_prob_unpadded)):
+                    ref_log_prob_list.append(
+                        ref_log_prob_unpadded[i].masked_fill(~mask[i], 0)
+                    )
+            if self.pad_output:
+                ref_log_prob = pad_sequence(
+                    ref_log_prob_list,
+                    batch_first=True,
+                    padding_value=0,
+                    padding_side=self.padding_side,
+                )
+            else:
+                ref_log_prob = torch.nested.nested_tensor(
+                    ref_log_prob_list, layout=torch.strided
+                )
+
+        # we obtain the current log-probs (already computed) from the current tensordict
+        if self.pad_output:
+            curr_log_prob_padded = tensordict.get(self.log_prob_full_key)
+        else:
+            curr_log_prob_unpadded = tensordict.get(
+                self.log_prob_full_key, as_list=True  # type: ignore[misc]
+            )
+        if self.assistant_only:
+            # we want to keep the batch dimension
+            curr_log_prob_list = []
+            if self.pad_output:
+                for i in range(curr_log_prob_padded.size(0)):
+                    curr_log_prob_list.append(
+                        curr_log_prob_padded[i].masked_fill(~mask[i], 0)
+                    )
+            else:
+                for i in range(len(curr_log_prob_unpadded)):
+                    curr_log_prob_list.append(
+                        curr_log_prob_unpadded[i].masked_fill(~mask[i], 0)
+                    )
+            if self.pad_output:
+                curr_log_prob = pad_sequence(
+                    curr_log_prob_list,
+                    batch_first=True,
+                    padding_value=0,
+                    padding_side=self.padding_side,
+                )
+            else:
+                curr_log_prob = torch.nested.nested_tensor(
+                    curr_log_prob_list, layout=torch.strided
+                )
+
+        ref_log_prob = ref_log_prob.to(curr_log_prob.device)
+        # We want the log-probs to have a similar dim to the reward
+        curr_log_prob = curr_log_prob.unsqueeze(-1)
+        ref_log_prob = ref_log_prob.unsqueeze(-1)
+
+        for i in range(ref_log_prob.size(0)):
+            if ref_log_prob[i].shape != curr_log_prob[i].shape:
+                # Don't check shapes if nested
+                raise ValueError(
+                    f"the log-probability tensor shapes must match, got cur_log_prob.shape={curr_log_prob[i].shape} and log_prob.shape={ref_log_prob[i].shape}. "
+                    f"One possible reason is that the padding token is identical to the eos token, which means that the eos_token log_prob is truncated from the "
+                    f"reference model output."
+                )
+        kl = curr_log_prob - ref_log_prob
+        if self.add_to_reward:
+            reward_key = self.in_keys[0]
+            reward = next_tensordict.get(reward_key)
+            # we use the unbiased consistent estimator of the KL: log_p(x) - log_q(x) when x ~ p(x)
+            if not reward.is_nested and ref_log_prob.is_nested:
+                reward = torch.nested.nested_tensor(
+                    [rew.expand(lp.shape) for rew, lp in zip(reward, ref_log_prob)],
+                    layout=torch.strided,
+                )
+            if reward is not None and reward.ndim != curr_log_prob.ndim:
+                raise ValueError(
+                    "The number of dimensions of reward must be the same as the number of dimensions of the KL "
+                    f"term. Got ndim={reward.ndim} and {curr_log_prob.ndim} respectively."
+                )
+            if reward is None:
+                reward = 0
+            reward = reward - self.coef * kl
+            next_tensordict.set(self.out_keys[0], reward)
+        next_tensordict.set(self.out_keys[1], kl)
+        next_tensordict.set(self.out_keys[2], ref_log_prob)
+        if original_device is not None:
+            next_tensordict = next_tensordict.to(original_device)
+        return next_tensordict
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        next_td = tensordict.pop("next")
+        next_td = self._step(tensordict, next_td)
+        return tensordict.set("next", next_td)
+
+    def transform_output_spec(self, output_spec: Composite) -> Composite:
+        in_key = unravel_key(self.in_keys[0])
+        out_key = unravel_key(self.out_keys[0])
+
+        if "full_observation_spec" in output_spec.keys():
+            observation_spec = output_spec["full_observation_spec"]
+        else:
+            observation_spec = Composite(
+                shape=output_spec.shape, device=output_spec.device
+            )
+            output_spec["full_observation_spec"] = observation_spec
+
+        if in_key == "reward" and out_key == "reward":
+            parent = self.parent
+
+            reward_keys = parent.reward_keys
+            if len(reward_keys) == 1:
+                reward_key = reward_keys[0]
+                shape = output_spec["full_reward_spec"].shape
+            elif "reward" in reward_keys:
+                reward_key = "reward"
+                shape = output_spec["full_reward_spec"].shape
+            else:
+                shape = output_spec.shape
+                reward_key = "reward"
+            # For LLMs, the shape of the reward is (batch, -1, 1)
+            shape = torch.Size((*shape, -1, 1))
+            reward_spec = Unbounded(
+                device=output_spec.device,
+                shape=shape,
+            )
+            output_spec["full_reward_spec"] = Composite(
+                {reward_key: reward_spec},
+                shape=output_spec["full_reward_spec"].shape,
+            )
+        elif in_key == "reward":
+            # TODO: we should at least allow to make this a component of the reward specs, to avoid a call during reset
+            parent = self.parent
+            reward_spec = output_spec["full_reward_spec"][parent.reward_key]
+
+            shape = output_spec["full_reward_spec"].shape
+            # For LLMs, the shape of the reward is (batch, -1, 1)
+            shape = torch.Size((*shape, -1, 1))
+            reward_spec = reward_spec.clone()
+            reward_spec.shape = shape
+
+            # then we need to populate the output keys
+            observation_spec[out_key] = reward_spec
+        else:
+            observation_spec = output_spec["full_observation_spec"]
+            reward_spec = observation_spec[in_key]
+
+            shape = observation_spec.shape
+            shape = torch.Size((*shape, -1, 1))
+            reward_spec = reward_spec.clone()
+            reward_spec.shape = shape
+
+            # then we need to populate the output keys
+            observation_spec[out_key] = reward_spec
+
+        observation_spec[self.out_keys[1]] = reward_spec.clone()
+
+        return output_spec
+
+
+class RetrieveLogProb(Transform):
+    """A transform to retrieve log-probabilities from a model for KL divergence computation.
+
+    This transform computes log-probabilities from a reference model, which can then be used
+    to compute KL divergence with another model's log-probabilities. It's designed to work
+    with the :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` and :class:`~torchrl.envs.llm.transforms.kl.KLComputation` transforms.
+
+    Args:
+        model (LLMWrapperBase): the model to use to compute the log-probs.
+
+    Keyword Args:
+        log_probs_full_key (NestedKey): the key where the log-probs are stored.
+            If not provided, the key will be retrieved from the model's `log_probs_key` attribute
+            (i.e., `(model.log_probs_key, "full")`).
+        assistant_only (bool): whether to zero out the log-probs of the non-assistant tokens (i.e., steps of history
+            where the role is not `"assistant"`). Defaults to `True`.
+
+            .. note:: When `assistant_only=True`, the model must have `input_mode='history'` to properly identify
+                assistant tokens. For other input modes (`"text"` or `"tokens"`), set `assistant_only=False`.
+                This ensures users are conscious of the limitation that assistant token identification requires
+                structured conversation history.
+
+        tokenizer_kwargs (dict): the keyword arguments to pass to the tokenizer to be used to apply the chat template to the history when `assistant_only` is `True`.
+            To control the tokenization in the ref_model, pass the tokenizer kwargs to the ref_model constructor.
+            Defaults to `{"return_assistant_tokens_mask": True, "tokenize": True, "return_dict": True, "padding": False, "add_generation_prompt": False}`.
+        tokenizer (transformers.AutoTokenizer): the tokenizer to be used to tokenize the input and compute the assitant mask. If not provided, the tokenizer will be inferred from the `ref_model`.
+        detach (bool): whether to exclude the log-probs from the gradient computation. Defaults to `True`.
+        device (torch.device): the device to use for tensor creation. Defaults to `None`.
+        padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
+
+    Examples:
+        >>> from torchrl.data.llm import History
+        >>> from torchrl.modules.llm import TransformersWrapper
+        >>> from torchrl.modules.llm.policies import ChatHistory
+        >>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
+        >>> from tensordict import TensorDict, set_list_to_stack
+        >>> import torch
+        >>>
+        >>> # Set up list to stack for History
+        >>> set_list_to_stack(True).set()
+        >>>
+        >>> # Create chat data
+        >>> chats = [
+        ...     [
+        ...         {"role": "system", "content": "You are a helpful assistant."},
+        ...         {"role": "user", "content": "Hello, how are you?"},
+        ...         {"role": "assistant", "content": "I'm doing well, thank you!"},
+        ...     ],
+        ...     [
+        ...         {"role": "system", "content": "You are a helpful assistant."},
+        ...         {"role": "user", "content": "What's the weather like?"},
+        ...         {"role": "assistant", "content": "I can't check the weather for you."},
+        ...     ],
+        ... ]
+        >>> history = History.from_chats(chats)
+        >>> print(f"Created history with shape: {history.shape}")
+        Created history with shape: torch.Size([2, 3])
+        >>>
+        >>> # Setup tokenizer and model
+        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+        >>> tokenizer.pad_token = tokenizer.eos_token
+        >>> model = OPTForCausalLM(OPTConfig()).eval()
+        >>>
+        >>> # Create reference model
+        >>> ref_model = TransformersWrapper(
+        ...     model,
+        ...     tokenizer=tokenizer,
+        ...     input_mode="history",
+        ...     generate=False,
+        ...     return_log_probs=True,
+        ...     pad_output=True,
+        ... )
+        >>>
+        >>> # Create the RetrieveLogProb transform
+        >>> transform = RetrieveLogProb(
+        ...     ref_model,
+        ...     assistant_only=True,
+        ...     tokenizer=tokenizer,
+        ... )
+        >>>
+        >>> # Prepare data using ChatHistory
+        >>> chat_history = ChatHistory(full=history)
+        >>> data = TensorDict(history=chat_history, batch_size=(2,))
+        >>>
+        >>> # Apply the transform to get reference log probabilities
+        >>> result = transform(data)
+        >>> log_probs_key = (ref_model.log_probs_key, "full")
+        >>> ref_log_probs = result.get(log_probs_key)
+        >>> print(f"Log-probs shape: {ref_log_probs.shape}")
+        Log-probs shape: torch.Size([2, 26])
+
+    .. note::
+        By default, the log-probabilities are stored as a list of tensors (one per sample, with variable length).
+        Use `as_padded_tensor=True` in `.get()` to obtain a batchable tensor (with padding).
+        The reference log probabilities are computed only for assistant tokens when `assistant_only=True`.
+
+        **Input Mode Compatibility:**
+        - When `assistant_only=True` (default), the model must have `input_mode='history'` to properly identify assistant tokens.
+        - When `assistant_only=False`, the transform works with any input mode (`"history"`, `"text"`, or `"tokens"`).
+        - This design ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
+
+    .. seealso::
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: A higher-level transform that combines two `RetrieveLogProb` instances with `KLComputation`.
+        :class:`~torchrl.envs.llm.transforms.kl.KLComputation`: A transform that computes KL divergence between two log-prob tensors.
+        :class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
+    """
+
+    def __init__(
+        self,
+        model: LLMWrapperBase,
+        *,
+        log_probs_full_key: NestedKey | None = None,
+        assistant_only: bool = True,
+        tokenizer_kwargs: dict | None = None,
+        detach: bool = True,
+        device: torch.device | None = None,
+        tokenizer: transformers.AutoTokenizer | None = None,
+        padding_side: str = "left",
+    ):
+        # Set up keys
+        if log_probs_full_key is None:
+            log_probs_full_key = (model.log_probs_key, "full")
+        elif (
+            not isinstance(log_probs_full_key, tuple)
+            or log_probs_full_key[-1] != "full"
+        ):
+            warnings.warn(
+                f"The log_probs_full_key {log_probs_full_key} is not a tuple or does not end with 'full'. "
+                "This may cause issues with the KL computation. "
+                "Please use a tuple with the log_probs_key and 'full' as the last element."
+            )
+        self.log_probs_full_key = log_probs_full_key
+
+        # Set up input/output keys
+        in_keys = list(model.in_keys)
+        out_keys = [self.log_probs_full_key]
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+
+        # Store model and configuration
+        self.model = model
+        self.assistant_only = assistant_only
+        self.detach = detach
+        self.device = device
+        self.tokenizer = tokenizer
+        self.padding_side = padding_side
+
+        # Set up tokenizer kwargs
+        if tokenizer_kwargs is None:
+            tokenizer_kwargs = {}
+        tokenizer_kwargs.setdefault("return_assistant_tokens_mask", True)
+        tokenizer_kwargs.setdefault("tokenize", True)
+        tokenizer_kwargs.setdefault("return_dict", True)
+        tokenizer_kwargs.setdefault("padding", False)
+        tokenizer_kwargs.setdefault("add_generation_prompt", False)
+        self.tokenizer_kwargs = tokenizer_kwargs
+
+        # Validate model configuration (after setting assistant_only)
+        self._validate_model_config(model)
+
+    def _validate_model_config(self, model: LLMWrapperBase):
+        """Validate model configuration."""
+        if not getattr(model, "return_log_probs", True):
+            raise ValueError(
+                "The model must have `return_log_probs=True` to use the `RetrieveLogProb` transform."
+            )
+        if getattr(model, "generate", True):
+            raise ValueError(
+                "The model must have `generate=False` to use the `RetrieveLogProb` transform."
+            )
+
+        # Check input mode compatibility with assistant_only
+        input_mode = getattr(model, "input_mode", "history")
+        if self.assistant_only and input_mode != "history":
+            raise ValueError(
+                f"The model must have `input_mode='history'` when `assistant_only=True`. "
+                f"Current input_mode is '{input_mode}'. "
+                f"To use input_mode '{input_mode}', set `assistant_only=False`."
+            )
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        next_td = tensordict.get("next")
+        next_is_none = False
+        if next_td is None:
+            next_is_none = True
+            next_td = tensordict
+        output = self._step(tensordict, next_td)
+        if next_is_none:
+            return output
+        return tensordict.set("next", output)
+
+    def _mask_assistant_tokens(
+        self, td: TensorDictBase, lp_key: NestedKey
+    ) -> torch.Tensor:
+        """Mask log-probs to only include assistant tokens.
+
+        Args:
+            td: TensorDict containing the data
+            lp_key: Key for log-probs in the TensorDict
+
+        Returns:
+            Masked log-probs tensor
+        """
+        with torch.device(self.device) if self.device is not None else nullcontext():
+            # Get assistant mask
+            assistant_masks = td.get(("masks", "all_assistant_mask"), as_list=True)  # type: ignore[misc]
+            log_probs = td.get(lp_key, as_list=True)  # type: ignore[misc]
+            log_probs = [
+                torch.masked_fill(lp, ~mask, 0.0)
+                for lp, mask in _zip_strict(log_probs, assistant_masks)
+            ]
+            if self.model.pad_output:
+                log_probs = pad_sequence(
+                    log_probs,
+                    batch_first=True,
+                    padding_value=0.0,
+                    padding_side=self.padding_side,
+                )
+            else:
+                log_probs = torch.nested.as_nested_tensor(
+                    log_probs, layout=self.model.layout
+                )
+            return log_probs
+
+    @set_list_to_stack(True)
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        # Compute log-probs using the model
+        # Use tensordict since we want to process the "full" entry
+        ref_td = self.model(tensordict.copy())
+        tmp_log_probs_key = (self.model.log_probs_key, "full")
+
+        # Apply assistant masking if requested
+        if self.assistant_only:
+            log_probs = self._mask_assistant_tokens(ref_td, tmp_log_probs_key)
+            ref_td.set(tmp_log_probs_key, log_probs)
+
+        # Rename and store the log-probs
+        if tmp_log_probs_key != self.log_probs_full_key:
+            ref_td.rename_key_(tmp_log_probs_key, self.log_probs_full_key)
+        next_tensordict.update(ref_td, keys_to_update=(self.log_probs_full_key,))
+
+        return next_tensordict
+
+    def transform_observation_spec(self, observation_spec: Composite) -> Composite:
+        # Add kl to observation spec
+        observation_spec["kl_penalty"] = Unbounded(
+            device=observation_spec.device,
+            shape=observation_spec.shape,
+        )
+        return observation_spec
+
+
+class RayRetrieveKL(RayTransform):
+    """A Ray-based implementation of :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`.
+
+    This class creates a Ray remote actor from RetrieveKL that can be shared across multiple workers.
+    All method calls are delegated to the remote actor, ensuring that multiple environments can
+    share the same KL computation resources.
+
+    To avoid serialization issues with large models, this class supports model factories
+    that create models on the remote actor rather than passing full models through Ray channels.
+
+    Args:
+        gen_model (LLMWrapperBase | Literal["from_collector"]): the generation model, or "from_collector" for lazy initialization.
+            Prefer using a model factory instead to avoid serialization issues.
+        ref_model (LLMWrapperBase | None): the reference model. Prefer using a model factory instead
+            to avoid serialization issues.
+
+    Keyword Args:
+        gen_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a generation model.
+            This allows for explicit resource control and avoids serialization issues.
+        ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
+            This allows for explicit resource control and avoids serialization issues.
+        num_cpus (int, optional): Number of CPUs to allocate to the Ray actor. Defaults to 1.
+        num_gpus (int, optional): Number of GPUs to allocate to the Ray actor. Defaults to 0.
+        device (torch.device, optional): Device to use on the remote Ray actor for tensor operations.
+            The local Ray transform will handle CPU serialization and device restoration automatically.
+            Defaults to None.
+        actor_name (str, optional): Name of the Ray actor to use. If provided, the actor will be reused if it already exists.
+        **kwargs: Additional keyword arguments to pass to RetrieveKL.
+
+    Note:
+        When using model factories, the corresponding model arguments (gen_model, ref_model) should be None.
+        Model factories are preferred for large models to avoid serialization overhead.
+
+    Examples:
+        >>> # Option 1: Using model factories for explicit resource control
+        >>> def create_gen_model():
+        ...     return TransformersWrapper(model, tokenizer=tokenizer, generate=False, return_log_probs=True)
+        >>> def create_ref_model():
+        ...     return TransformersWrapper(ref_model, tokenizer=tokenizer, generate=False, return_log_probs=True)
+        >>> transform = RayRetrieveKL(
+        ...     gen_model=None, ref_model=None,
+        ...     gen_model_factory=create_gen_model,
+        ...     ref_model_factory=create_ref_model,
+        ...     num_gpus=1,
+        ...     device=torch.device("cuda")
+        ... )
+
+        >>> # Option 2: Pass models directly (Ray handles serialization)
+        >>> transform = RayRetrieveKL(gen_model=gen_model, ref_model=ref_model, device=torch.device("cuda"))
+    """
+
+    def __init__(
+        self,
+        gen_model: LLMWrapperBase | Literal["from_collector"] | None = "from_collector",
+        ref_model: LLMWrapperBase | None = None,
+        *,
+        gen_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int = 0,
+        device: DEVICE_TYPING | None = None,
+        actor_name: str | None = None,
+        **kwargs,
+    ):
+        # Validate arguments: models and factories should not both be provided
+        if gen_model is not None and gen_model_factory is not None:
+            raise ValueError(
+                "Cannot provide both 'gen_model' and 'gen_model_factory'. Choose one."
+            )
+        if ref_model is not None and ref_model_factory is not None:
+            raise ValueError(
+                "Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
+            )
+
+        # Store creation parameters for actor creation
+        self._gen_model = gen_model
+        self._ref_model = ref_model
+        self._gen_model_factory = gen_model_factory
+        self._ref_model_factory = ref_model_factory
+        self._creation_kwargs = kwargs
+        # Store device separately for passing to remote actor
+        self._remote_device = device
+
+        # Default num_cpus
+        if num_cpus is None:
+            num_cpus = 1
+
+        # Call parent constructor without device (Ray transform handles CPU/device mapping)
+        super().__init__(
+            num_cpus=num_cpus,
+            num_gpus=num_gpus,
+            device=None,  # Don't store device locally
+            actor_name=actor_name,
+            **kwargs,
+        )
+
+    def _create_actor(self, **kwargs):
+        """Create the remote RetrieveKL actor."""
+        # Create the remote RetrieveKL with resource specifications
+        RemoteRetrieveKL = self._ray.remote(
+            num_cpus=self._num_cpus, num_gpus=self._num_gpus
+        )(RetrieveKL)
+
+        if self._actor_name is not None:
+            RemoteRetrieveKL = RemoteRetrieveKL.options(name=self._actor_name)
+
+        # Determine how to create models on the remote actor
+        gen_model_arg = self._gen_model
+        ref_model_arg = self._ref_model
+
+        # If we have factories, we'll pass them and set models to None
+        creation_kwargs = self._creation_kwargs.copy()
+        if self._gen_model_factory is not None:
+            creation_kwargs["gen_model_factory"] = self._gen_model_factory
+            gen_model_arg = None
+        if self._ref_model_factory is not None:
+            creation_kwargs["ref_model_factory"] = self._ref_model_factory
+            ref_model_arg = None
+
+        # Pass device to the remote actor
+        if self._remote_device is not None:
+            creation_kwargs["device"] = self._remote_device
+
+        # Create the shared actor
+        actor = RemoteRetrieveKL.remote(
+            gen_model=gen_model_arg, ref_model=ref_model_arg, **creation_kwargs
+        )
+
+        return actor
+
+    def __repr__(self):
+        """String representation."""
+        try:
+            if hasattr(self, "_actor") and self._actor is not None:
+                return self._ray.get(self._actor.__repr__.remote())
+            else:
+                return "RayRetrieveKL(actor=None)"
+        except Exception:
+            return f"RayRetrieveKL(actor={getattr(self, '_actor', 'None')})"
+
+
+class RetrieveKL(Compose, metaclass=_RayServiceMetaClass):
+    """A transform to retrieve the KL divergence between two models' log-probabilities.
+
+    This transform combines two :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb` instances
+    with a :class:`~torchrl.envs.llm.transforms.kl.KLComputation` to compute KL divergence
+    between a generation model and a reference model.
+
+    .. note::
+        Both gen_model and ref_model must use the same pad_output value (True or False), otherwise KL computation will fail.
+
+    Args:
+        gen_model (LLMWrapperBase): the generation model, wrapped in such a way that it does not generate but computes the log-probs.
+            In cases where the transform is used within a :class:`~torchrl.collectors.llm.LLMCollector` run on a remote worker, the
+            policy may not be available ahead of time. In this case, the `gen_model` can be set to `"from_collector"` (default) to retrieve the
+            policy from the collector. See :meth:`~torchrl.modules.llm.policies.LLMWrapperBase.get_new_version` for more details
+            about generating a new version of the policy to gather the log-probs.
+        ref_model (LLMWrapperBase): the reference model, wrapped in such a way that it does not generate but computes the log-probs.
+
+    Keyword Args:
+        gen_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a generation model.
+            This allows for explicit resource control and avoids serialization issues when using Ray.
+        ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
+            This allows for explicit resource control and avoids serialization issues when using Ray.
+        assistant_only (bool): whether to only retrieve the log-probs of the assistant tokens (i.e., steps of history
+            where the role is `"assistant"`). Defaults to `True`.
+
+            .. note:: When `assistant_only=True`, both models must have `input_mode='history'` to properly identify assistant tokens.
+                For other input modes (`"text"` or `"tokens"`), set `assistant_only=False`.
+                This ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
+
+        gen_log_probs_full_key (str): the key where the log-probs of the generation model are stored. Defaults to `("log_probs", "full")`.
+        ref_log_probs_full_key (str): the key where the log-probs of the reference model are stored. Defaults to `("ref_log_probs", "full")`.
+        history_key (str): the key where the history is stored. Defaults to `"history"`.
+        tokenizer_kwargs (dict): the keyword arguments to pass to the tokenizer to be used to apply the chat template to the history when `assistant_only` is `True`.
+            To control the tokenization in the actor, pass the tokenizer kwargs to the actor constructor.
+            Defaults to `{"return_assistant_tokens_mask": True, "tokenize": True, "return_tensors": "pt", "padding": True, "add_generation_prompt": False}`.
+        detach (bool): whether to exclude the log-probs from the gradient computation. Defaults to `True`.
+        device (torch.device): the device to cast the tensors to. This is not the device of the specs, but the device
+            onto which the tensors will be moved. It allows to keep the model on a different device
+            than the upcoming data itself. When using Ray service, this device will be used on the remote actor.
+            Defaults to `None`.
+        tokenizer (transformers.AutoTokenizer): the tokenizer to be used to tokenize the input and compute the assitant mask. If not provided, the tokenizer will be inferred from the `actor`.
+        padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
+        kl_key (NestedKey): the key where the KL divergence is stored. Defaults to `"kl_penalty"`.
+        add_to_reward (bool): whether to add the KL divergence to the reward. Defaults to `True`.
+        coeff (float): the coefficient for the KL term when adding to reward. Defaults to `1.0`.
+        padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
+        use_ray_service (bool, optional): if ``True``, returns a :class:`RayRetrieveKL` instance instead,
+            which creates a Ray actor for shared KL computation across multiple environments.
+            Defaults to ``False``.
+        actor_name (str, optional): the name of the Ray actor to use. Defaults to `None`.
+        **kwargs: additional arguments to pass to the `RetrieveLogProb` transform.
+
+    Examples:
+        >>> from torchrl.data.llm import History
+        >>> from torchrl.modules.llm import TransformersWrapper
+        >>> from torchrl.modules.llm.policies import ChatHistory
+        >>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
+        >>> from tensordict import TensorDict, set_list_to_stack
+        >>> import torch
+        >>>
+        >>> # Set up list to stack for History
+        >>> set_list_to_stack(True).set()
+        >>>
+        >>> # Create chat data
+        >>> chats = [
+        ...     [
+        ...         {"role": "system", "content": "You are a helpful assistant."},
+        ...         {"role": "user", "content": "Hello, how are you?"},
+        ...         {"role": "assistant", "content": "I'm doing well, thank you!"},
+        ...     ],
+        ...     [
+        ...         {"role": "system", "content": "You are a helpful assistant."},
+        ...         {"role": "user", "content": "What's the weather like?"},
+        ...         {"role": "assistant", "content": "I can't check the weather for you."},
+        ...     ],
+        ... ]
+        >>> history = History.from_chats(chats)
+        >>> print(f"Created history with shape: {history.shape}")
+        Created history with shape: torch.Size([2, 3])
+        >>>
+        >>> # Setup tokenizer and model
+        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+        >>> tokenizer.pad_token = tokenizer.eos_token
+        >>> model = OPTForCausalLM(OPTConfig()).eval()
+        >>>
+        >>> # Create generation and reference models
+        >>> gen_model = TransformersWrapper(
+        ...     model,
+        ...     tokenizer=tokenizer,
+        ...     input_mode="history",
+        ...     generate=False,
+        ...     return_log_probs=True,
+        ...     pad_output=True,
+        ...     log_probs_key="gen_log_probs",
+        ... )
+        >>> ref_model = TransformersWrapper(
+        ...     model,
+        ...     tokenizer=tokenizer,
+        ...     input_mode="history",
+        ...     generate=False,
+        ...     return_log_probs=True,
+        ...     pad_output=True,
+        ...     log_probs_key="ref_log_probs",
+        ... )
+        >>>
+        >>> # Create RetrieveKL transform
+        >>> transform = RetrieveKL(
+        ...     gen_model=gen_model,
+        ...     ref_model=ref_model,
+        ...     assistant_only=True,
+        ...     tokenizer=tokenizer,
+        ... )
+        >>>
+        >>> # Prepare data with next tensordict using ChatHistory
+        >>> chat_history = ChatHistory(full=history)
+        >>> next_td = TensorDict(history=chat_history, batch_size=(2,))
+        >>> data = TensorDict(history=chat_history, next=next_td, batch_size=(2,))
+        >>>
+        >>> # Apply transform
+        >>> result = transform(data)
+        >>> kl = result["next"].get("kl_penalty")
+        >>> print(f"KL shape: {kl.shape}")
+        KL shape: torch.Size([2, 26])
+
+    Note:
+        **Input Mode Compatibility:**
+        - When `assistant_only=True`, both models must have `input_mode='history'` to properly identify assistant tokens.
+        - When `assistant_only=False`, the transform works with any input mode (`"history"`, `"text"`, or `"tokens"`).
+        - This design ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
+
+    .. seealso::
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: The base transform for retrieving log-probabilities from a single model.
+        :class:`~torchrl.envs.llm.transforms.kl.KLComputation`: The transform that computes KL divergence between two log-prob tensors.
+        :class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
+    """
+
+    _RayServiceClass = RayRetrieveKL
+
+    def __init__(
+        self,
+        gen_model: LLMWrapperBase | Literal["from_collector"] = "from_collector",
+        ref_model: LLMWrapperBase | None = None,
+        *,
+        gen_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
+        assistant_only: bool = True,
+        history_key: str = "history",
+        tokenizer_kwargs: dict[str, Any] | None = None,
+        detach: bool = True,
+        device: torch.device | None = None,
+        tokenizer: transformers.AutoTokenizer | None = None,
+        padding_side: str = "left",
+        gen_log_probs_full_key: NestedKey = ("log_probs", "full"),
+        ref_log_probs_full_key: NestedKey = ("ref_log_probs", "full"),
+        kl_key: NestedKey = "kl_penalty",
+        add_to_reward: bool = True,
+        coeff: float = 1.0,
+        use_ray_service: bool = False,
+        **kwargs,
+    ):
+        # Handle model factories - create models if factories are provided
+        if gen_model_factory is not None:
+            if gen_model is not None and gen_model != "from_collector":
+                raise ValueError(
+                    "Cannot provide both 'gen_model' and 'gen_model_factory'. Choose one."
+                )
+            gen_model = gen_model_factory()
+
+        if ref_model_factory is not None:
+            if ref_model is not None:
+                raise ValueError(
+                    "Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
+                )
+            ref_model = ref_model_factory()
+
+        if isinstance(gen_model, str) and gen_model == "from_collector":
+            # Lazy init
+            self._initialized = False
+            self._init_params = {
+                "ref_model": ref_model,
+                "gen_model_factory": gen_model_factory,
+                "ref_model_factory": ref_model_factory,
+                "assistant_only": assistant_only,
+                "history_key": history_key,
+                "tokenizer_kwargs": tokenizer_kwargs,
+                "detach": detach,
+                "device": device,
+                "tokenizer": tokenizer,
+                "gen_log_probs_full_key": gen_log_probs_full_key,
+                "ref_log_probs_full_key": ref_log_probs_full_key,
+                "kl_key": kl_key,
+                "add_to_reward": add_to_reward,
+                "coeff": coeff,
+                "padding_side": padding_side,
+                **kwargs,
+            }
+            super().__init__()
+            return
+
+        self._initialized = True
+
+        # Check pad_output consistency if both models are provided
+        if hasattr(gen_model, "pad_output") and hasattr(ref_model, "pad_output"):
+            if gen_model.pad_output != ref_model.pad_output:
+                raise ValueError(
+                    f"pad_output mismatch: gen_model.pad_output={gen_model.pad_output}, "
+                    f"ref_model.pad_output={ref_model.pad_output}. "
+                    "Both models must use the same padding strategy for KL computation."
+                )
+
+        if not getattr(gen_model, "return_log_probs", True):
+            raise ValueError(
+                "The generation model must have `return_log_probs=True` to use the `RetrieveKL` transform."
+            )
+        elif getattr(gen_model, "generate", False):
+            raise ValueError(
+                "The generation model must have `generate=False` to use the `RetrieveKL` transform."
+            )
+
+        if not getattr(ref_model, "return_log_probs", True):
+            raise ValueError(
+                "The reference model must have `return_log_probs=True` to use the `RetrieveKL` transform."
+            )
+        elif getattr(ref_model, "generate", False):
+            raise ValueError(
+                "The reference model must have `generate=False` to use the `RetrieveKL` transform."
+            )
+        if getattr(gen_model, "log_probs_key", "gen_log_probs") == getattr(
+            ref_model, "log_probs_key", "log_probs"
+        ):
+            raise ValueError(
+                "The generation and reference models must have different `log_prob_key` values to use the `RetrieveKL` transform."
+            )
+        if gen_model is None:
+            raise ValueError("gen_model cannot be None when not using 'from_collector'")
+        if ref_model is None:
+            raise ValueError("ref_model cannot be None")
+
+        t1 = RetrieveLogProb(
+            gen_model,
+            log_probs_full_key=gen_log_probs_full_key,
+            assistant_only=assistant_only,
+            tokenizer_kwargs=tokenizer_kwargs,
+            detach=detach,
+            device=device,
+            tokenizer=tokenizer,
+            padding_side=padding_side,
+            **kwargs,
+        )
+        t2 = RetrieveLogProb(
+            ref_model,
+            log_probs_full_key=ref_log_probs_full_key,
+            assistant_only=assistant_only,
+            tokenizer_kwargs=tokenizer_kwargs,
+            detach=detach,
+            device=device,
+            tokenizer=tokenizer,
+            padding_side=padding_side,
+            **kwargs,
+        )
+        t3 = KLComputation(
+            gen_log_probs_full_key=gen_log_probs_full_key,
+            ref_log_probs_full_key=ref_log_probs_full_key,
+            kl_key=kl_key,
+            add_to_reward=add_to_reward,
+            coeff=coeff,
+        )
+        super().__init__(t1, t2, t3)
+
+    def _init_deferred(self):
+        torchrl_logger.info("Initializing RetrieveKL transform")
+        container = self.container
+        if container is None:
+            # also logging, since this will be sometimes hidden within the AttributeError
+            torchrl_logger.warning(
+                "The container is not set. Please set the container before calling this method."
+            )
+            raise ValueError(
+                "The container is not set. Please set the container before calling this method."
+            )
+        container.empty_cache()
+        self.empty_cache()
+        collector = self.collector
+        if collector is None:
+            # also logging, since this will be sometimes hidden within the AttributeError
+            torchrl_logger.warning(
+                "The collector is not set. Please set the collector before calling this method."
+            )
+            raise ValueError(
+                "The collector is not set. Please set the collector before calling this method."
+            )
+        ref_model = self._init_params["ref_model"]
+        pad_output = getattr(ref_model, "pad_output", None)
+        gen_log_probs_full_key = self._init_params["gen_log_probs_full_key"]
+        if (
+            not isinstance(gen_log_probs_full_key, tuple)
+            or gen_log_probs_full_key[-1] != "full"
+        ):
+            raise ValueError(
+                f"The gen_log_probs_full_key {gen_log_probs_full_key} is not a tuple or does not end with 'full'. "
+                "This may cause issues with the KL computation. "
+                "Please use a tuple with the log_probs_key and 'full' as the last element."
+            )
+        log_probs_key = gen_log_probs_full_key[:-1]
+        gen_model = collector.policy.get_new_version(
+            generate=False,
+            return_log_probs=True,
+            log_probs_key=log_probs_key,
+            input_mode=ref_model.input_mode,
+            input_key=(ref_model.input_mode, "full"),
+            pad_output=pad_output,  # Pass pad_output from ref_model
+        )
+        # Create the transforms manually instead of calling __init__
+        t1 = RetrieveLogProb(
+            gen_model,
+            log_probs_full_key=gen_log_probs_full_key,
+            assistant_only=self._init_params["assistant_only"],
+            tokenizer_kwargs=self._init_params["tokenizer_kwargs"],
+            detach=self._init_params["detach"],
+            device=self._init_params["device"],
+            tokenizer=self._init_params["tokenizer"],
+            padding_side=self._init_params["padding_side"],
+        )
+        ref_log_probs_full_key = self._init_params["ref_log_probs_full_key"]
+        if (
+            not isinstance(ref_log_probs_full_key, tuple)
+            or ref_log_probs_full_key[-1] != "full"
+        ):
+            raise ValueError(
+                f"The ref_log_probs_full_key {ref_log_probs_full_key} is not a tuple or does not end with 'full'. "
+                "This may cause issues with the KL computation. "
+                "Please use a tuple with the log_probs_key and 'full' as the last element."
+            )
+        t2 = RetrieveLogProb(
+            ref_model,
+            log_probs_full_key=ref_log_probs_full_key,
+            assistant_only=self._init_params["assistant_only"],
+            tokenizer_kwargs=self._init_params["tokenizer_kwargs"],
+            detach=self._init_params["detach"],
+            device=self._init_params["device"],
+            tokenizer=self._init_params["tokenizer"],
+            padding_side=self._init_params["padding_side"],
+        )
+        t3 = KLComputation(
+            gen_log_probs_full_key=gen_log_probs_full_key,
+            ref_log_probs_full_key=ref_log_probs_full_key,
+            kl_key=self._init_params["kl_key"],
+            add_to_reward=self._init_params["add_to_reward"],
+            coeff=self._init_params["coeff"],
+        )
+        # Replace the transforms in the Compose
+        self.transforms.extend([t1, t2, t3])
+        del self._init_params
+        self._initialized = True
+        torchrl_logger.info("Successfully initialized")
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        if not self._initialized:
+            self._init_deferred()
+        return super()._step(tensordict, next_tensordict)
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        if not self._initialized:
+            self._init_deferred()
+        return super()._reset(tensordict, tensordict_reset)
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if not self._initialized:
+            self._init_deferred()
+        return super().forward(tensordict)
+
+    def transform_observation_spec(self, observation_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_observation_spec(observation_spec)
+
+    def transform_reward_spec(self, reward_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_reward_spec(reward_spec)
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if not self._initialized:
+            self._init_deferred()
+        return super()._inv_call(tensordict)
+
+    def transform_action_spec(self, action_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_action_spec(action_spec)
+
+    def transform_input_spec(self, input_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_input_spec(input_spec)
+
+    def transform_output_spec(self, output_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_output_spec(output_spec)
+
+    def transform_state_spec(self, state_spec: Composite) -> Composite:
+        if not self._initialized:
+            self._init_deferred()
+        return super().transform_state_spec(state_spec)
+
+
+class KLComputation(Transform):
+    """A transform to compute KL divergence between two log-prob tensors and optionally add it to the reward.
+
+    This transform computes KL divergence between generation and reference log-probabilities
+    and can optionally subtract it from the reward (for KL penalty). It's designed to work
+    with the :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb` and :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` transforms.
+
+    .. note::
+        Both input log-prob tensors must use the same padding strategy (pad_output) for correct KL computation.
+
+    Args:
+        gen_log_probs_full_key (NestedKey): the key where the generation model log-probs are stored.
+            Defaults to `("gen_log_probs", "full")`.
+        ref_log_probs_full_key (NestedKey): the key where the reference model log-probs are stored.
+            Defaults to `("ref_log_probs", "full")`.
+        kl_key (NestedKey): the key where the KL divergence is stored. Defaults to `"kl_penalty"`.
+        add_to_reward (bool): whether to add the KL divergence to the reward. Defaults to `True`.
+        coeff (float): the coefficient for the KL term when adding to reward. Defaults to `1.0`.
+        padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
+
+    Examples:
+        >>> from tensordict import TensorDict
+        >>> import torch
+        >>>
+        >>> # Create sample log-probs
+        >>> gen_log_probs = torch.randn(2, 10)  # 2 samples, 10 tokens each
+        >>> ref_log_probs = torch.randn(2, 10)
+        >>>
+        >>> # Create data with next tensordict
+        >>> next_td = TensorDict(
+        ...     {
+        ...         ("gen_log_probs", "full"): gen_log_probs,
+        ...         ("ref_log_probs", "full"): ref_log_probs,
+        ...         "reward": torch.randn(2, 10, 1),
+        ...     },
+        ...     batch_size=(2,)
+        ... )
+        >>> data = TensorDict(next=next_td, batch_size=(2,))
+        >>>
+        >>> # Create KLComputation transform
+        >>> kl_transform = KLComputation(
+        ...     gen_log_probs_key=("gen_log_probs", "full"),
+        ...     ref_log_probs_key=("ref_log_probs", "full"),
+        ...     kl_key="kl_penalty",
+        ...     add_to_reward=True,
+        ...     coef=1.0,
+        ... )
+        >>>
+        >>> # Apply transform
+        >>> result = kl_transform(data)
+        >>> kl = result["next"].get("kl_penalty")
+        >>> print(f"KL shape: {kl.shape}")
+        KL shape: torch.Size([2, 10])
+
+    .. seealso::
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: The base transform for retrieving log-probabilities from a single model.
+        :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: A higher-level transform that combines two `RetrieveLogProb` instances with `KLComputation`.
+        :class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
+
+    """
+
+    def __init__(
+        self,
+        gen_log_probs_full_key: NestedKey = ("log_probs", "full"),
+        ref_log_probs_full_key: NestedKey = ("ref_log_probs", "full"),
+        *,
+        kl_key: NestedKey = "kl_penalty",
+        add_to_reward: bool = True,
+        coeff: float = 1.0,
+        padding_side: str = "left",
+    ):
+        in_keys = [gen_log_probs_full_key, ref_log_probs_full_key]
+        if add_to_reward:
+            in_keys.append("reward")
+        out_keys = [kl_key]
+        if add_to_reward:
+            out_keys.append("reward")
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+
+        self.gen_log_probs_full_key = gen_log_probs_full_key
+        self.ref_log_probs_full_key = ref_log_probs_full_key
+        self.kl_key = kl_key
+        self.add_to_reward = add_to_reward
+        self.coeff = coeff
+        self.padding_side = padding_side
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        next_td = tensordict.get("next")
+        has_next_td = True
+        if next_td is None:
+            next_td = tensordict
+            has_next_td = False
+        next_td = self._step(tensordict, next_td)
+        if has_next_td:
+            return tensordict.set("next", next_td)
+        return next_td
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        # Get log-probs
+        gen_log_probs = next_tensordict.get(self.gen_log_probs_full_key, as_list=True)  # type: ignore[misc]
+        ref_log_probs = next_tensordict.get(self.ref_log_probs_full_key, as_list=True)  # type: ignore[misc]
+
+        if gen_log_probs is None or ref_log_probs is None:
+            raise ValueError(
+                f"Log-probs not found. Expected keys: {self.gen_log_probs_key}, {self.ref_log_probs_key}"
+            )
+
+        # Debug: Check lengths and shapes
+        if len(gen_log_probs) != len(ref_log_probs):
+            raise ValueError(
+                f"Batch size mismatch: gen_log_probs has {len(gen_log_probs)} samples, ref_log_probs has {len(ref_log_probs)} samples"
+            )
+
+        # Check individual sequence lengths
+        for i, (gen_lp, ref_lp) in enumerate(_zip_strict(gen_log_probs, ref_log_probs)):
+            if gen_lp.shape != ref_lp.shape:
+                raise ValueError(
+                    f"Sample {i} has different shapes: gen_log_probs[{i}].shape={gen_lp.shape}, ref_log_probs[{i}].shape={ref_lp.shape}"
+                )
+
+        # Compute KL divergence: KL(p||q) = E_p[log p - log q]
+        # Here gen_log_probs = log p, ref_log_probs = log q
+        kl = [
+            gen_lp - ref_lp
+            for gen_lp, ref_lp in _zip_strict(gen_log_probs, ref_log_probs)
+        ]
+
+        kl = torch.nested.as_nested_tensor(kl, layout=torch.strided)
+
+        next_tensordict.set(self.kl_key, kl)
+
+        # Add to reward if requested
+        if self.add_to_reward:
+            reward = next_tensordict.get("reward", as_list=True)  # type: ignore[misc]
+            if reward is not None:
+                if isinstance(reward, list):
+                    if reward[0].ndim != kl[0].ndim + 1:
+                        raise ValueError(
+                            f"The rewards have shape {reward[0].shape} but the kl has shape {kl[0].shape}. "
+                            f"The rewards should have one more dimension than the KL."
+                        )
+                    reward = [
+                        r - self.coeff * k.unsqueeze(-1)
+                        for r, k in _zip_strict(reward, kl)
+                    ]
+                    next_tensordict.set(
+                        "reward",
+                        torch.nested.as_nested_tensor(reward, layout=torch.strided),
+                    )
+                else:
+                    if reward.ndim != kl.ndim + 1:
+                        raise ValueError(
+                            f"The rewards have shape {reward.shape} but the kl has shape {kl.shape}. "
+                            f"The rewards should have one more dimension than the KL."
+                        )
+                    reward = reward - self.coeff * kl.unsqueeze(-1)
+                    next_tensordict.set("reward", reward)
+
+        return next_tensordict
+
+    def transform_observation_spec(self, observation_spec: Composite) -> Composite:
+        # Add kl to observation spec
+        observation_spec[self.kl_key] = Unbounded(
+            device=observation_spec.device,
+            shape=observation_spec.shape,
+        )
+        return observation_spec
+
+    def transform_reward_spec(self, reward_spec: Composite) -> Composite:
+        # Optionally adjust reward spec if KL is added to reward
+        if self.add_to_reward:
+            shape = reward_spec["reward"].shape
+            # For LLMs, the shape of the reward is (batch, -1, 1)
+            shape = torch.Size((*shape, -1, 1))
+            reward_spec["reward"] = reward_spec["reward"].clone()
+            reward_spec["reward"].shape = shape
+        return reward_spec