Add ability to save and load rational and basic agents to/from json.

- add methods for saving and loading to components
- add save / load to all main agent factories
- saving and loading of associative memories

PiperOrigin-RevId: 690960131
Change-Id: Ie323681fe63fbc966d70bbd21bcaf9d09f06c48d

Author: vezhnick

concordia/agents/entity_agent.py

@@ -107,6 +107,14 @@ def get_component(
     component = self._context_components[name]
     return cast(entity_component.ComponentT, component)
 
+  def get_act_component(self) -> entity_component.ActingComponent:
+    return self._act_component
+
+  def get_all_context_components(
+      self,
+  ) -> Mapping[str, entity_component.ContextComponent]:
+    return types.MappingProxyType(self._context_components)
+
   def _parallel_call_(
       self,
       method_name: str,

concordia/agents/entity_agent_with_logging.py

@@ -15,10 +15,12 @@
 """A modular entity agent using the new component system with side logging."""
 
 from collections.abc import Mapping
+import copy
 import types
 from typing import Any
 from absl import logging
 from concordia.agents import entity_agent
+from concordia.associative_memory import formative_memories
 from concordia.typing import agent
 from concordia.typing import entity_component
 from concordia.utils import measurements as measurements_lib
@@ -39,6 +41,7 @@ def __init__(
           types.MappingProxyType({})
       ),
       component_logging: measurements_lib.Measurements | None = None,
+      config: formative_memories.AgentConfig | None = None,
   ):
     """Initializes the agent.
 
@@ -56,6 +59,7 @@ def __init__(
         None, a NoOpContextProcessor will be used.
       context_components: The ContextComponents that will be used by the agent.
       component_logging: The channels where components publish events.
+      config: The agent configuration, used for checkpointing and debug.
     """
     super().__init__(agent_name=agent_name,
                      act_component=act_component,
@@ -75,6 +79,7 @@ def __init__(
           on_error=lambda e: logging.error('Error in component logging: %s', e))
     else:
       self._channel_names = []
+    self._config = copy.deepcopy(config)
 
   def _set_log(self, log: tuple[Any, ...]) -> None:
     """Set the logging object to return from get_last_log.
@@ -89,3 +94,6 @@ def _set_log(self, log: tuple[Any, ...]) -> None:
   def get_last_log(self):
     self._tick.on_next(None)  # Trigger the logging.
     return self._log
+
+  def get_config(self) -> formative_memories.AgentConfig | None:
+    return copy.deepcopy(self._config)

concordia/associative_memory/associative_memory.py

@@ -25,9 +25,11 @@
 import threading
 
 from concordia.associative_memory import importance_function
+from concordia.typing import entity_component
 import numpy as np
 import pandas as pd
 
+
 _NUM_TO_RETRIEVE_TO_CONTEXTUALIZE_IMPORTANCE = 25
 
 
@@ -79,6 +81,29 @@ def __init__(
     self._interval = clock_step_size
     self._stored_hashes = set()
 
+  def get_state(self) -> entity_component.ComponentState:
+    """Converts the AssociativeMemory to a dictionary."""
+
+    with self._memory_bank_lock:
+      output = {
+          'seed': self._seed,
+          'stored_hashes': list(self._stored_hashes),
+          'memory_bank': self._memory_bank.to_json(),
+      }
+      if self._interval:
+        output['interval'] = self._interval.total_seconds()
+    return output
+
+  def set_state(self, state: entity_component.ComponentState) -> None:
+    """Sets the AssociativeMemory from a dictionary."""
+
+    with self._memory_bank_lock:
+      self._seed = state['seed']
+      self._stored_hashes = set(state['stored_hashes'])
+      self._memory_bank = pd.read_json(state['memory_bank'])
+      if 'interval' in state:
+        self._interval = datetime.timedelta(seconds=state['interval'])
+
   def add(
       self,
       text: str,

concordia/associative_memory/formative_memories.py

@@ -15,7 +15,7 @@
 
 """This is a factory for generating memories for concordia agents."""
 
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Collection, Sequence
 import dataclasses
 import datetime
 import logging
@@ -58,9 +58,24 @@ class AgentConfig:
   specific_memories: str = ''
   goal: str = ''
   date_of_birth: datetime.datetime = DEFAULT_DOB
-  formative_ages: Iterable[int] = DEFAULT_FORMATIVE_AGES
+  formative_ages: Collection[int] = DEFAULT_FORMATIVE_AGES
   extras: dict[str, Any] = dataclasses.field(default_factory=dict)
 
+  def to_dict(self) -> dict[str, Any]:
+    """Converts the AgentConfig to a dictionary."""
+    result = dataclasses.asdict(self)
+    result['date_of_birth'] = self.date_of_birth.isoformat()
+    return result
+
+  @classmethod
+  def from_dict(cls, data: dict[str, Any]) -> 'AgentConfig':
+    """Initializes an AgentConfig from a dictionary."""
+    date_of_birth = datetime.datetime.fromisoformat(
+        data['date_of_birth']
+    )
+    data = data | {'date_of_birth': date_of_birth}
+    return cls(**data)
+
 
 class FormativeMemoryFactory:
   """Generator of formative memories."""

concordia/components/agent/action_spec_ignored.py

@@ -16,10 +16,11 @@
 
 import abc
 import threading
-from typing import Final
+from typing import Final, Any
 
 from concordia.typing import entity as entity_lib
 from concordia.typing import entity_component
+from typing_extensions import override
 
 
 class ActionSpecIgnored(
@@ -89,3 +90,11 @@ def get_named_component_pre_act_value(self, component_name: str) -> str:
     """Returns the pre-act value of a named component of the parent entity."""
     return self.get_entity().get_component(
         component_name, type_=ActionSpecIgnored).get_pre_act_value()
+
+  @override
+  def set_state(self, state: entity_component.ComponentState) -> Any:
+    return None
+
+  @override
+  def get_state(self) -> entity_component.ComponentState:
+    return {}

concordia/components/agent/all_similar_memories.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Return all memories similar to a prompt and filter them for relevance.
-"""
+"""Return all memories similar to a prompt and filter them for relevance."""
 
-from collections.abc import Mapping
 import types
+from typing import Mapping
 
 from concordia.components.agent import action_spec_ignored
 from concordia.components.agent import memory_component

concordia/components/agent/concat_act_component.py

@@ -164,3 +164,11 @@ def _log(self,
         'Value': result,
         'Prompt': prompt.view().text().splitlines(),
     })
+
+  def get_state(self) -> entity_component.ComponentState:
+    """Converts the component to a dictionary."""
+    return {}
+
+  def set_state(self, state: entity_component.ComponentState) -> None:
+    pass
+

concordia/components/agent/memory_component.py

@@ -57,6 +57,14 @@ def _check_phase(self) -> None:
           'You can only access the memory outside of the `UPDATE` phase.'
       )
 
+  def get_state(self) -> Mapping[str, Any]:
+    with self._lock:
+      return self._memory.get_state()
+
+  def set_state(self, state: Mapping[str, Any]) -> None:
+    with self._lock:
+      self._memory.set_state(state)
+
   def retrieve(
       self,
       query: str = '',

concordia/components/agent/plan.py

@@ -163,3 +163,15 @@ def _make_pre_act_value(self) -> str:
     })
 
     return result
+
+  def get_state(self) -> entity_component.ComponentState:
+    """Converts the component to JSON data."""
+    with self._lock:
+      return {
+          'current_plan': self._current_plan,
+      }
+
+  def set_state(self, state: entity_component.ComponentState) -> None:
+    """Sets the component state from JSON data."""
+    with self._lock:
+      self._current_plan = state['current_plan']

concordia/components/agent/question_of_query_associated_memories.py

@@ -189,6 +189,12 @@ def pre_act(
   def update(self) -> None:
     self._component.update()
 
+  def get_state(self) -> entity_component.ComponentState:
+    return self._component.get_state()
+
+  def set_state(self, state: entity_component.ComponentState) -> None:
+    self._component.set_state(state)
+
 
 class Identity(QuestionOfQueryAssociatedMemories):
   """Identity component containing a few characteristics.