XiuchengWu
300e86956b
- Add ConsensusSnapshot and ConsensusSlot classes for structured consensus state representation. - Introduce OnlineLanguageLearner for stateful online learning with experience replay and EWC regularization. - Create WorkspaceEvent and ActionEvent classes for serializable event types. - Develop InMemoryVectorMemoryStore for storing and querying workspace events. - Implement ActionPolicy to manage action selection and motor parameterization. - Establish a unified WorkspaceRuntime for managing the agent's lifecycle and interactions. - Enhance tests in test_smoke.py to cover new functionalities and ensure correctness.
116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
"""
|
|
Memory stores for workspace events.
|
|
|
|
The default store is local and in-memory, but it exposes a small backend
|
|
interface that can later be replaced by a vector database or remote service.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
import time
|
|
from typing import Protocol
|
|
|
|
import numpy as np
|
|
|
|
from .events import WorkspaceEvent
|
|
|
|
|
|
@dataclass
|
|
class MemoryRecord:
|
|
event: WorkspaceEvent
|
|
similarity: float | None = None
|
|
|
|
@property
|
|
def w(self) -> np.ndarray:
|
|
return self.event.workspace_array()
|
|
|
|
@property
|
|
def context(self) -> dict:
|
|
return self.event.context
|
|
|
|
@property
|
|
def timestamp(self) -> float:
|
|
return self.event.timestamp
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"w": self.w,
|
|
"context": self.context,
|
|
"timestamp": self.timestamp,
|
|
"similarity": self.similarity,
|
|
"event": self.event.to_dict(),
|
|
}
|
|
|
|
|
|
class MemoryStore(Protocol):
|
|
def put(self, event: WorkspaceEvent) -> None:
|
|
...
|
|
|
|
def query(self, event_or_workspace, top_k: int = 3) -> list[MemoryRecord]:
|
|
...
|
|
|
|
def size(self) -> int:
|
|
...
|
|
|
|
|
|
class InMemoryVectorMemoryStore:
|
|
"""Small cosine-similarity memory store for workspace events."""
|
|
|
|
def __init__(self, capacity: int = 1000, workspace_dim: int = 64):
|
|
self.capacity = capacity
|
|
self.workspace_dim = workspace_dim
|
|
self.records: deque[MemoryRecord] = deque(maxlen=capacity)
|
|
|
|
@property
|
|
def memories(self):
|
|
return self.records
|
|
|
|
def put(self, event: WorkspaceEvent) -> None:
|
|
self.records.append(MemoryRecord(event=event))
|
|
|
|
def query(self, event_or_workspace, top_k: int = 3) -> list[MemoryRecord]:
|
|
if not self.records:
|
|
return []
|
|
query = self._as_workspace(event_or_workspace)
|
|
similarities = [
|
|
float(np.dot(query, record.w) / (np.linalg.norm(query) * np.linalg.norm(record.w) + 1e-8))
|
|
for record in self.records
|
|
]
|
|
top_idx = np.argsort(similarities)[-top_k:][::-1]
|
|
return [
|
|
MemoryRecord(event=self.records[int(idx)].event, similarity=similarities[int(idx)])
|
|
for idx in top_idx
|
|
]
|
|
|
|
def size(self) -> int:
|
|
return len(self.records)
|
|
|
|
def store(self, w: np.ndarray, context: dict | None = None):
|
|
"""Backward-compatible write API."""
|
|
event = WorkspaceEvent(
|
|
step=int((context or {}).get("step", 0)),
|
|
modality=str((context or {}).get("modality", "unknown")),
|
|
workspace=self._as_workspace(w).astype(float).tolist(),
|
|
consensus=(context or {}).get("consensus"),
|
|
context=context or {},
|
|
timestamp=time.time(),
|
|
)
|
|
self.put(event)
|
|
|
|
def recall(self, w_query: np.ndarray, top_k: int = 3) -> list[dict]:
|
|
"""Backward-compatible query API."""
|
|
return [record.to_dict() for record in self.query(w_query, top_k)]
|
|
|
|
def _as_workspace(self, event_or_workspace) -> np.ndarray:
|
|
if isinstance(event_or_workspace, WorkspaceEvent):
|
|
return event_or_workspace.workspace_array()
|
|
data = np.asarray(event_or_workspace, dtype=np.float32)
|
|
if data.ndim > 1:
|
|
data = data[0]
|
|
return data.reshape(-1)
|
|
|
|
|
|
# Backward-compatible brain-region name.
|
|
Hippocampus = InMemoryVectorMemoryStore
|