Determine whether to wrap tree loaded from pickle or if it's already in the correct format

python/skills/skills/episodic_verbalization.py

@@ -2,7 +2,7 @@ from pathlib import Path
 import pickle
 import logging
 
-from episodic_verbalization_core.em.em_tree import HigherLevelSummary
+from episodic_verbalization_core.em.em_tree import HigherLevelSummary, GoalBasedSummary
 from episodic_verbalization_core.lmp.setup import instantiate_llm
 from episodic_verbalization_core.em.armarx_lt_mem import load_episode_from_armarx_lt_mem, extend_existing_history_from_memory_snapshots
 from episodic_verbalization_core.em.llm_summary import LLMBasedSummarizer
@@ -28,7 +28,16 @@ class EM_Verbalization_Tree:
         # load tree:
         if self.history_cache.is_file():  
             print('loading tree from', self.history_cache)
-            self.tree = HigherLevelSummary('', children=[pickle.loads(self.history_cache.read_bytes())])
+            tree = pickle.loads(self.history_cache.read_bytes())
+            assert type(tree).__name__ == HigherLevelSummary.__name__
+            if type(tree.children[-1]).__name__ == HigherLevelSummary.__name__:
+                # this tree already seems to be in the correct format
+                self.tree = tree
+            elif all(type(c).__name__ == GoalBasedSummary.__name__ for c in tree.children):
+                # this tree should be wrapped
+                self.tree = HigherLevelSummary('', children=[])
+            else:
+                raise ValueError('Unsupported tree file ' + str(self.history_cache))
         else:
             print('could not load tree from', self.history_cache)
             self.create_tree()