[Feature] Support lazy tensordict inputs in ppo loss

ghstack-source-id: 89098ba
Pull Request resolved: #2883

Author: Vincent Moens

torchrl/modules/llm/vllm_wrapper.py

@@ -223,7 +223,7 @@ def __init__(
         if from_text:
             self.out_keys += [self.text_response_key, self.token_key]
         if self.return_log_probs:
-            self.out_keys += ["log_probs"]
+            self.out_keys += [self.log_prob_key]
 
     def forward(
         self,
@@ -303,7 +303,7 @@ def _from_vllm_generate_text(self, td):
                 ),
             )
         in_keys = [
-            "log_probs",
+            self.log_prob_key,
             self.token_response_key,
             self.text_response_key,
             self.token_key,
@@ -394,7 +394,7 @@ def _from_vllm_logprobs_text(self, td):
         if isinstance(input_ids_response, list):
             input_ids_response = torch.nested.nested_tensor(input_ids_response)
         out["tokens_response"] = input_ids_response
-        out["log_probs"] = lps
+        out[self.log_prob_key] = lps
         inputs = td.select(*self.in_keys, strict=False)
         if inputs.ndim < out.ndim:
             # This happens when n > 1
@@ -423,18 +423,19 @@ def _from_vllm_generate_tokens(self, td):
             ).to_padded_tensor(padding=self.padding_value)
         tokens_response_td.rename_key_("token_ids", "tokens_response")
         if self.return_log_probs:
-            tokens_response_td.rename_key_("logprobs", "log_probs")
+            tokens_response_td.rename_key_("logprobs", self.log_prob_key)
             if self.pad_output:
                 padded_values = (
                     tokens_response_td["tokens_response"] == self.padding_value
                 )
                 if padded_values.any():
-                    lps = tokens_response_td["log_probs"]
+                    lps = tokens_response_td[self.log_prob_key]
                     lps = torch.where(expand_as_right(~padded_values, lps), lps, 0.0)
-                    tokens_response_td["log_probs"] = lps
+                    tokens_response_td[self.log_prob_key] = lps
         out = tokens_response_td.empty(recurse=True)
         out.update(
-            tokens_response_td, keys_to_update=(self.token_response_key, "log_probs")
+            tokens_response_td,
+            keys_to_update=(self.token_response_key, self.log_prob_key),
         )
         inputs = td.select(*self.in_keys, strict=False)
         if inputs.ndim < out.ndim:
@@ -467,7 +468,7 @@ def _from_vllm_logprobs_tokens(self, td):
         padded = tokens_response == self.padding_value
         prompt_logprobs = torch.where(~padded, prompt_logprobs, 0.0)
         out = tokens_out._tensordict.empty(recurse=True)
-        out.set("log_probs", prompt_logprobs)
+        out.set(self.log_prob_key, prompt_logprobs)
         out.set(self.token_response_key, tokens_response)
         inputs = td.select(*self.in_keys, strict=False)
         if inputs.ndim < out.ndim:
@@ -501,13 +502,13 @@ def _get_output_tokens_and_log_probs(self, tokens_out):
             )
 
         if self.return_log_probs or "logprobs" in tokens_response_td:
-            tokens_response_td.rename_key_("logprobs", "log_probs")
+            tokens_response_td.rename_key_("logprobs", self.log_prob_key)
             if self.pad_output:
                 padded_values = tokens_response_td["tokens_response"] == padding_value
                 if padded_values.any():
-                    lps = tokens_response_td["log_probs"]
+                    lps = tokens_response_td[self.log_prob_key]
                     lps = torch.where(expand_as_right(~padded_values, lps), lps, 0.0)
-                    tokens_response_td["log_probs"] = lps
+                    tokens_response_td[self.log_prob_key] = lps
         return tokens_response_td
 
     def _to_list(self, tokens, attention_mask):

torchrl/objectives/ppo.py

@@ -533,7 +533,7 @@ def _get_cur_log_prob(self, tensordict):
         if isinstance(
             self.actor_network,
             (ProbabilisticTensorDictSequential, ProbabilisticTensorDictModule),
-        ):
+        ) or hasattr(self.actor_network, "get_dist"):
             # assert tensordict['log_probs'].requires_grad
             # assert tensordict['logits'].requires_grad
             with self.actor_network_params.to_module(
@@ -987,7 +987,9 @@ def out_keys(self, values):
     @dispatch
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         tensordict = tensordict.clone(False)
-        advantage = tensordict.get(self.tensor_keys.advantage, None)
+        advantage = tensordict.get(
+            self.tensor_keys.advantage, None, as_padded_tensor=True
+        )
         if advantage is None:
             if self.critic_network is None:
                 raise RuntimeError(

torchrl/objectives/utils.py

@@ -624,7 +624,7 @@ def _sum_td_features(data: TensorDictBase) -> torch.Tensor:
 
 def _maybe_get_or_select(td, key_or_keys, target_shape=None):
     if isinstance(key_or_keys, (str, tuple)):
-        return td.get(key_or_keys)
+        return td.get(key_or_keys, as_padded_tensor=True)
     result = td.select(*key_or_keys)
     if target_shape is not None and result.shape != target_shape:
         result.batch_size = target_shape