[Tutorial] Beam search with GPT models

ghstack-source-id: 36de30d
Pull Request resolved: #2623

Author: Vincent Moens

docs/requirements.txt

@@ -28,3 +28,7 @@ vmas
 onnxscript
 onnxruntime
 onnx
+plotly
+igraph
+transformers
+datasets

docs/source/_static/img/rollout-llm.png

@@ -105,6 +105,7 @@ Intermediate
    tutorials/dqn_with_rnn
    tutorials/rb_tutorial
    tutorials/export
+   tutorials/beam_search_with_gpt
 
 Advanced
 --------

docs/source/index.rst

@@ -798,6 +798,10 @@ def done_keys(self) -> List[NestedKey]:
 
     @done_keys.setter
     def done_keys(self, value):
+        if isinstance(value, (str, tuple)):
+            value = [value]
+        if value is not None:
+            value = [unravel_key(val) for val in value]
         self._done_keys = _make_list_of_nestedkeys(value, "done_keys")
 
     @property
@@ -818,6 +822,10 @@ def reward_keys(self) -> List[NestedKey]:
 
     @reward_keys.setter
     def reward_keys(self, value):
+        if isinstance(value, (str, tuple)):
+            value = [value]
+        if value is not None:
+            value = [unravel_key(val) for val in value]
         self._reward_keys = _make_list_of_nestedkeys(value, "reward_keys")
 
     @property
@@ -838,6 +846,10 @@ def action_keys(self) -> List[NestedKey]:
 
     @action_keys.setter
     def action_keys(self, value):
+        if isinstance(value, (str, tuple)):
+            value = [value]
+        if value is not None:
+            value = [unravel_key(val) for val in value]
         self._action_keys = _make_list_of_nestedkeys(value, "action_keys")
 
     @property
@@ -857,6 +869,10 @@ def observation_keys(self) -> List[NestedKey]:
 
     @observation_keys.setter
     def observation_keys(self, value):
+        if isinstance(value, (str, tuple)):
+            value = [value]
+        if value is not None:
+            value = [unravel_key(val) for val in value]
         self._observation_keys = _make_list_of_nestedkeys(value, "observation_keys")
 
     @property
@@ -1012,6 +1028,27 @@ def add(self, step, *, return_node: bool = False):
         if return_node:
             return self.get_tree(step)
 
+    def add(self, step):
+        source, dest = (
+            step.exclude("next").copy(),
+            step.select("next", *self.action_keys).copy(),
+        )
+
+        if self.data_map is None:
+            self._make_storage(source, dest)
+
+        # We need to set the action somewhere to keep track of what action lead to what child
+        # # Set the action in the 'next'
+        # dest[1:] = source[:-1].exclude(*self.done_keys)
+
+        # Add ('observation', 'action') -> ('next, observation')
+        self.data_map[source] = dest
+        value = source
+        if self.node_map is None:
+            self._make_storage_branches(source, dest)
+        # map ('observation',) -> ('indices',)
+        self.node_map[source] = value
+
     def get_child(self, root: TensorDictBase) -> TensorDictBase:
         return self.data_map[root]