[BugFix] Wrappers stack fn (#3061)

Author: vmoens

test/llm/test_wrapper.py

@@ -1192,6 +1192,8 @@ def test_retrieve_kl_input_modes(
             ref_model=ref_model,
             assistant_only=assistant_only,
             tokenizer=tokenizer,
+            gen_log_probs_full_key=("gen_log_probs", "full"),
+            ref_log_probs_full_key=("ref_log_probs", "full"),
         )
 
         # Apply transform
@@ -1202,15 +1204,15 @@ def test_retrieve_kl_input_modes(
         # Check that both log-probs and KL are present
         assert ("gen_log_probs", "full") in result
         assert ("ref_log_probs", "full") in result
-        assert "kl" in result
+        assert "kl_penalty" in result
 
         # Check KL structure
         if pad_output:
-            kl = result.get("kl")
+            kl = result.get("kl_penalty")
             assert isinstance(kl, torch.Tensor)
             assert kl.shape[0] == 2  # batch size
         else:
-            kl = result.get("kl", as_list=True)
+            kl = result.get("kl_penalty", as_list=True)
             # For unpadded output, we get a list of tensors
             assert isinstance(kl, list)
             assert len(kl) == 2  # batch size
@@ -1391,7 +1393,7 @@ def test_kl_computation_transform(
 
         # Check that reward is modified
         assert "reward" in result
-        reward = result.get("reward")
+        reward = result.get("reward", as_list=True)
         assert reward is not None
 
 

torchrl/envs/llm/transforms/kl.py

@@ -680,8 +680,8 @@ class RetrieveKL(Compose):
                 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_prob_full_key (str): the key where the log-probs of the generation model are stored. Defaults to `("log_probs", "full")`.
-        ref_log_prob_full_key (str): the key where the log-probs of the reference model are stored. Defaults to `("ref_log_probs", "full")`.
+        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.

torchrl/modules/llm/policies/transformers_wrapper.py

@@ -686,7 +686,7 @@ def _from_transformers_generate_history(self, td, cfg, out) -> TensorDictBase:
             h_responses = _extract_responses_from_full_histories(
                 text_full, prompt_histories, self.chat_template_name, self.tokenizer
             )
-            history_chat_flat.response = torch.stack(h_responses)
+            history_chat_flat.response = h_responses
         result.set(self.history_key, history_chat)
         return result
 

torchrl/modules/llm/policies/vllm_wrapper.py

@@ -726,7 +726,7 @@ def _from_vllm_generate_history(
             h_responses = _extract_responses_from_full_histories(
                 text_full, prompt_histories, self.chat_template_name, self.tokenizer
             )
-            history_chat_flat.response = torch.stack(h_responses)
+            history_chat_flat.response = h_responses
         result.set(self.history_key, history_chat)
         return result