[BugFix] Tree make node fix (#2839)

Author: rolo

test/test_storage_map.py

@@ -350,6 +350,17 @@ def test_edges(self):
         edges_check = {(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)}
         assert edges == edges_check
 
+    def test_make_node(self):
+        td = TensorDict({"obs": torch.tensor([0])})
+        tree = Tree(node_data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(td)
+        assert tree.node_data is not None
+
 
 class TestMCTSForest:
     def dummy_rollouts(self) -> Tuple[TensorDict, ...]:

torchrl/data/llm/__init__.py

@@ -11,7 +11,14 @@
 )
 from .prompt import PromptData, PromptTensorDictTokenizer
 from .reward import PairwiseDataset, RewardData
-from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel, LLMData, LLMOutput, LLMInput
+from .utils import (
+    AdaptiveKLController,
+    ConstantKLController,
+    LLMData,
+    LLMInput,
+    LLMOutput,
+    RolloutFromModel,
+)
 
 __all__ = [
     "AdaptiveKLController",

torchrl/data/llm/utils.py

@@ -543,8 +543,10 @@ def step_scheduler(self):
             while len(self._kl_queue):
                 self._kl_queue.remove(self._kl_queue[0])
 
+
 LLMInpOut = TypeVar("LLMInpOut")
 
+
 class LLMInput(TensorClass["nocast"]):
     """Represents the input to a Large Language Model (LLM).
 
@@ -557,11 +559,13 @@ class LLMInput(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMOutput` and :class:`~torchrl.data.LLMData`.
 
     """
+
     tokens: torch.Tensor
     attention_mask: torch.Tensor | None = None
     token_list: list[int] | list[list[int]] | None = None
     text: str | list[str] | None = None
 
+
 class LLMOutput(TensorClass["nocast"]):
     """Represents the output from a Large Language Model (LLM).
 
@@ -581,6 +585,7 @@ class LLMOutput(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMInput` and :class:`~torchrl.data.LLMData`.
 
     """
+
     tokens: torch.Tensor
     tokens_response: torch.Tensor | None = None
     token_list: list[int] | list[list[int]] | None = None
@@ -594,6 +599,7 @@ def from_vllm_output(cls: type[LLMInpOut], vllm_output) -> LLMInpOut:
         # placeholder
         raise NotImplementedError
 
+
 class LLMData(TensorClass["nocast"]):
     """Represents the input or output of a Large Language Model (LLM).
 
@@ -619,6 +625,7 @@ class LLMData(TensorClass["nocast"]):
     .. seealso:: :class:`~torchrl.data.LLMInput` and :class:`~torchrl.data.LLMOutput`.
 
     """
+
     tokens: torch.Tensor
     tokens_response: torch.Tensor | None = None
     attention_mask: torch.Tensor | None = None

torchrl/data/map/tree.py

@@ -122,7 +122,7 @@ def make_node(
         return cls(
             count=torch.zeros(()),
             wins=torch.zeros(()),
-            node=data.exclude("action", "next"),
+            node_data=data.exclude("action", "next"),
             rollout=rollout,
             subtree=subtree,
             device=device,