floaties (#4789)

facebookresearch · Sep 16, 2022 · accbd2f · accbd2f
1 parent 285f32a
commit accbd2f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/projects/bb3/agents/opt_bb3_agent.py b/projects/bb3/agents/opt_bb3_agent.py
@@ -314,7 +314,7 @@ def get_opening_memories(
         opening_memories = None
         if memories:
             if isinstance(memories, dict):
-                opening_memories = memories
+                opening_memories = {mem: int(turns) for mem, turns in memories.items()}
             elif isinstance(memories, list):
                 opening_memories = {}
                 for mem in memories:

diff --git a/projects/bb3/agents/utils.py b/projects/bb3/agents/utils.py
@@ -396,6 +396,7 @@ def get_available_memories(
         """
         available_memory = []
         for memory, turns_since_used in memories.items():
+            turns_since_used = int(turns_since_used)
             # check if we should ignore in session memories
             if ignore_in_session_memories and memory in in_session_memories:
                 continue

diff --git a/tests/nightly/gpu/test_bb3.py b/tests/nightly/gpu/test_bb3.py
@@ -611,6 +611,8 @@ def test_softblocking(self):
 
     def test_memory_utils(self):
         memories = self.memories
+        # test that it works with floats as well
+        memories = {m: float(t) for m, t in memories.items()}
         decay_factor = 0.99
         success = False
         for _ in range(10):