Add visualization code

Author: josephmaa

bridger/go_explore_phase_1.py

@@ -444,10 +444,10 @@ def explore(
                 # iteration of exploratory rollouts.
                 state_sampler.update(state_sampler_cache_update)
 
-    object_logger.log(
-        f"state_cache-{hparams.env_width}.pkl",
-        OccurrenceLogEntry(batch_idx=0, object=state_sampler),
-    )
+        object_logger.log(
+            f"state_cache-{hparams.env_width}.pkl",
+            OccurrenceLogEntry(batch_idx=iteration, object=state_sampler),
+        )
 
     return success_entries
 

bridger/logging_utils/cache_entry_database.py

@@ -0,0 +1,55 @@
+from bridger.go_explore_phase_1 import CacheEntry
+
+
+class CacheEntryDatabase:
+    """
+    Provides different views into a cache. The cache is a list of CacheEntry objects and can be sorted by trajectory length, steps_since_led_to_something_new, visit count, or sample count.
+    """
+
+    def __init__(self, cache_entries: list[CacheEntry]):
+        self.cache_entries = cache_entries
+
+    def sort_by_key(self, sort_key: "SortKey"):
+        """
+        Sorts the cache entries using the provided sort key.
+        """
+        self.cache_entries = sorted(self.cache_entries, key=sort_key)
+
+    def get_top_n_by_sort_key(self, sort_key: "SortKey", n: int):
+        """
+        Returns the top n cache entries.
+        """
+        return self.sort_by_key(sort_key)[:n]
+
+
+class SortKey:
+    def __init__(self, key: str):
+        self.key = key
+
+    def __call__(self, cache_entry: CacheEntry):
+        return getattr(cache_entry, self.key)
+
+
+class TrajectorySortKey(SortKey):
+    def __init__(self):
+        super().__init__("trajectory")
+
+
+class StepsSinceLedToSomethingNewSortKey(SortKey):
+    def __init__(self):
+        super().__init__("steps_since_led_to_something_new")
+
+
+class StepsSinceLedToSomethingNewResetCountSortKey(SortKey):
+    def __init__(self):
+        super().__init__("steps_since_led_to_something_new_reset_count")
+
+
+class VisitCountSortKey(SortKey):
+    def __init__(self):
+        super().__init__("visit_count")
+
+
+class SampleCountSortKey(SortKey):
+    def __init__(self):
+        super().__init__("sample_count")

bridger/logging_utils/object_log_readers.py

@@ -31,7 +31,20 @@ def read_object_log(log_filepath: str):
         buffer = None
         while True:
             try:
-                buffer = pickle.load(f)
+                # Try to load with torch.load first to handle CUDA tensors
+                try:
+                    buffer = torch.load(f, map_location=torch.device("cpu"))
+                except:
+                    # If that fails, try regular pickle.load
+                    f.seek(0)  # Reset file pointer
+                    buffer = pickle.load(f)
+                    # Move any tensors to CPU
+                    if isinstance(buffer, list):
+                        for i, item in enumerate(buffer):
+                            if isinstance(item, torch.Tensor):
+                                buffer[i] = item.cpu()
+                    elif isinstance(buffer, torch.Tensor):
+                        buffer = buffer.cpu()
 
                 for element in buffer:
                     yield element

tools/web/object_log_cache.py

@@ -24,6 +24,7 @@
 
 ACTION_INVERSION_DATABASE_KEY = "action_inversion_database_key"
 TRAINING_HISTORY_DATABASE_KEY = "training_history_database_key"
+CACHE_ENTRY_DATABASE_KEY = "cache_entry_database_key"
 
 DatabaseType = (
     object_log_readers.TrainingHistoryDatabase
@@ -73,6 +74,13 @@ def _load_training_history_database_from_log(
 def _save_database(
     directory: str, experiment_name: str, database: DatabaseType
 ) -> None:
+    """Saves a database to disk.
+
+    Args:
+      directory: The directory to save the database in.
+      experiment_name: The name of the experiment to save.
+      database: The database to save.
+    """
     with open(os.path.join(directory, experiment_name), "wb") as f:
         pickle.dump(database, f)
 
@@ -82,8 +90,17 @@ def _database_exists(directory: str, experiment_name: str) -> bool:
 
 
 def _load_database(directory: str, experiment_name: str) -> DatabaseType:
+    """Loads a database from disk.
+
+    Args:
+      directory: The directory to load the database from.
+      experiment_name: The name of the experiment to load.
+
+    Returns:
+      The loaded database.
+    """
     with open(os.path.join(directory, experiment_name), "rb") as f:
-        return pickle.load(f)
+        return pickle.load(f, map_location=torch.device("cpu"))
 
 
 def _convert_log_to_saved_database_if_necessary(

tools/web/sibyl.py

@@ -11,6 +11,16 @@
 from tools.web import object_log_cache
 from tools.web import plot_utils
 from collections import Counter
+from bridger.go_explore_phase_1 import CacheEntry
+from bridger.logging_utils.cache_entry_database import (
+    CacheEntryDatabase,
+    TrajectorySortKey,
+    VisitCountSortKey,
+    SampleCountSortKey,
+    StepsSinceLedToSomethingNewSortKey,
+    StepsSinceLedToSomethingNewResetCountSortKey,
+)
+from bridger.logging_utils.object_log_readers import read_object_log
 
 app = flask.Flask(__name__)
 
@@ -65,6 +75,38 @@ def _get_experiment_names() -> list[str]:
     ]
 
 
+@app.route("/n_fewest_steps_since_led_to_something_new_go_explore", methods=["GET"])
+def n_fewest_steps_since_led_to_something_new_go_explore_plot_data():
+    """
+    Provides plot data for the n states that have most recently seen a new cell.
+    """
+    experiment_name = _get_string_or_default(_EXPERIMENT_NAME)
+    n = _get_int_or_default("n", 10)
+
+    cache_entry_database = CacheEntryDatabase(
+        list(read_object_log(os.path.join(_LOG_DIR, "go_explore", "start_entries.pkl")))
+    )
+
+    return {
+        "states": [
+            cache_entry.state for cache_entry in cache_entry_database.cache_entries
+        ],
+        "trajectory_length": cache_entry_database.get_top_n_by_sort_key(
+            TrajectorySortKey, n
+        ),
+        "steps_since_led_to_something_new": cache_entry_database.get_top_n_by_sort_key(
+            StepsSinceLedToSomethingNewSortKey, n
+        ),
+        "steps_since_led_to_something_new_reset_count": cache_entry_database.get_top_n_by_sort_key(
+            StepsSinceLedToSomethingNewResetCountSortKey, n
+        ),
+        "sample_count": cache_entry_database.get_top_n_by_sort_key(
+            SampleCountSortKey, n
+        ),
+        "visit_count": cache_entry_database.get_top_n_by_sort_key(VisitCountSortKey, n),
+    }
+
+
 @app.route("/training_history_plot_data", methods=["GET"])
 def training_history_plot_data():
     """Provides plot data on states and metrics based on filters.
@@ -374,6 +416,19 @@ def action_inversion():
     )
 
 
+@app.route("/go_explore")
+def go_explore():
+    experiment_names = _get_experiment_names()
+    selected_experiment_name = _get_string_or_default(
+        name=_EXPERIMENT_NAME, default=experiment_names[0]
+    )
+    return flask.render_template(
+        "go_explore.html",
+        experiment_names=experiment_names,
+        selected_experiment_name=selected_experiment_name,
+    )
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Load Sibyl debugger.")
     parser.add_argument(