Add purpose_agent/experience_replay.py

purpose_agent/experience_replay.py

@@ -0,0 +1,407 @@
+"""
+Experience Replay — Trajectory storage and retrieval.
+
+Stores completed trajectories with their scores and supports retrieval
+ranked by a combination of:
+  1. Semantic similarity (embedding-based) — from MemRL (arxiv:2601.03192)
+  2. Learned Q-value utility scores — from REMEMBERER (arxiv:2306.07929)
+
+The two-phase retrieval (recall by similarity → re-rank by Q-value) separates
+"semantically similar" from "functionally useful" — a key insight from MemRL.
+
+This module is the "database" — it stores but doesn't analyze. The Optimizer
+module reads from here and writes heuristics back.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import math
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+from purpose_agent.types import (
+    Heuristic,
+    MemoryRecord,
+    MemoryTier,
+    Trajectory,
+    TrajectoryStep,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ExperienceReplay:
+    """
+    Experience Replay buffer with two-phase retrieval.
+    
+    Phase 1 (Recall): Retrieve top-k records by semantic similarity to query
+    Phase 2 (Re-rank): Re-order by learned Q-value utility scores
+    
+    The buffer supports:
+    - Adding trajectories with automatic scoring
+    - Retrieving similar past experiences for the Actor's context
+    - Q-value updates after heuristics are applied (Bellman-style)
+    - Persistence to disk (JSON)
+    - Capacity management (evict lowest Q-value records when full)
+    
+    Args:
+        capacity: Maximum number of records to store
+        similarity_weight: λ in retrieval score = λ·similarity + (1-λ)·q_value
+        persistence_path: If set, auto-save/load buffer to this file
+    """
+
+    def __init__(
+        self,
+        capacity: int = 500,
+        similarity_weight: float = 0.6,
+        persistence_path: str | Path | None = None,
+    ):
+        self.capacity = capacity
+        self.similarity_weight = similarity_weight
+        self.persistence_path = Path(persistence_path) if persistence_path else None
+        self.records: list[MemoryRecord] = []
+
+        # Load from disk if available
+        if self.persistence_path and self.persistence_path.exists():
+            self._load()
+
+    # ------------------------------------------------------------------
+    # Core Operations
+    # ------------------------------------------------------------------
+
+    def add(self, trajectory: Trajectory) -> MemoryRecord:
+        """
+        Add a completed trajectory to the buffer.
+        
+        Automatically computes a task embedding (simple TF-IDF-style hash)
+        and initial Q-value based on trajectory performance.
+        """
+        # Compute initial Q-value from trajectory performance
+        initial_q = self._compute_initial_q(trajectory)
+
+        record = MemoryRecord(
+            trajectory=trajectory,
+            heuristics=[],
+            task_embedding=self._compute_embedding(trajectory.task_description),
+            retrieval_q_value=initial_q,
+        )
+
+        # Capacity management: evict lowest Q-value if full
+        if len(self.records) >= self.capacity:
+            self._evict()
+
+        self.records.append(record)
+        logger.info(
+            f"Experience Replay: Added trajectory '{trajectory.id}' "
+            f"(q={initial_q:.3f}, steps={len(trajectory.steps)}, "
+            f"Σreward={trajectory.cumulative_reward:.2f})"
+        )
+
+        if self.persistence_path:
+            self._save()
+
+        return record
+
+    def retrieve(
+        self,
+        query: str,
+        top_k: int = 5,
+        min_q_value: float = 0.0,
+    ) -> list[MemoryRecord]:
+        """
+        Two-phase retrieval (per MemRL arxiv:2601.03192):
+        
+        Phase 1: Recall candidates by semantic similarity
+        Phase 2: Re-rank by Q-value utility
+        
+        Returns top-k records sorted by combined score.
+        """
+        if not self.records:
+            return []
+
+        query_embedding = self._compute_embedding(query)
+
+        # Phase 1: Compute similarity scores for all records
+        scored: list[tuple[float, MemoryRecord]] = []
+        for record in self.records:
+            if record.retrieval_q_value < min_q_value:
+                continue
+            sim = self._cosine_similarity(
+                query_embedding, record.task_embedding or []
+            )
+            # Phase 2: Combined score
+            combined = (
+                self.similarity_weight * sim
+                + (1 - self.similarity_weight) * record.retrieval_q_value
+            )
+            scored.append((combined, record))
+
+        # Sort descending by combined score
+        scored.sort(key=lambda x: -x[0])
+
+        results = [record for _, record in scored[:top_k]]
+        logger.debug(
+            f"Experience Replay: Retrieved {len(results)} records for query "
+            f"(top score={scored[0][0]:.3f})" if scored else "no records"
+        )
+        return results
+
+    def update_q_value(
+        self,
+        record_id: str,
+        reward: float,
+        alpha: float = 0.1,
+    ) -> None:
+        """
+        Update a record's retrieval Q-value using Monte Carlo update.
+        
+        Q_new = Q_old + α * (reward - Q_old)
+        
+        From REMEMBERER (arxiv:2306.07929): α = 1/N where N = number of
+        updates. We use a fixed α for simplicity; override for REMEMBERER-exact.
+        """
+        for record in self.records:
+            if record.id == record_id:
+                old_q = record.retrieval_q_value
+                record.retrieval_q_value += alpha * (reward - old_q)
+                record.retrieval_q_value = max(0.0, min(1.0, record.retrieval_q_value))
+                logger.debug(
+                    f"Experience Replay: Q-value update for {record_id}: "
+                    f"{old_q:.3f} → {record.retrieval_q_value:.3f}"
+                )
+                if self.persistence_path:
+                    self._save()
+                return
+        logger.warning(f"Experience Replay: Record {record_id} not found for Q-update")
+
+    def attach_heuristics(
+        self, record_id: str, heuristics: list[Heuristic]
+    ) -> None:
+        """Attach extracted heuristics to a memory record."""
+        for record in self.records:
+            if record.id == record_id:
+                record.heuristics = heuristics
+                if self.persistence_path:
+                    self._save()
+                return
+
+    # ------------------------------------------------------------------
+    # Statistics & Queries
+    # ------------------------------------------------------------------
+
+    def get_top_trajectories(
+        self,
+        n: int = 10,
+        min_success_rate: float = 0.5,
+    ) -> list[Trajectory]:
+        """Get the n best trajectories by cumulative reward."""
+        candidates = [
+            r.trajectory for r in self.records
+            if r.trajectory.success_rate >= min_success_rate
+        ]
+        candidates.sort(key=lambda t: -t.cumulative_reward)
+        return candidates[:n]
+
+    def get_all_heuristics(self, tier: MemoryTier | None = None) -> list[Heuristic]:
+        """Get all extracted heuristics, optionally filtered by tier."""
+        heuristics = []
+        for record in self.records:
+            for h in record.heuristics:
+                if tier is None or h.tier == tier:
+                    heuristics.append(h)
+        return heuristics
+
+    @property
+    def size(self) -> int:
+        return len(self.records)
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        if not self.records:
+            return {"size": 0}
+        q_values = [r.retrieval_q_value for r in self.records]
+        rewards = [r.trajectory.cumulative_reward for r in self.records]
+        return {
+            "size": len(self.records),
+            "avg_q_value": sum(q_values) / len(q_values),
+            "max_q_value": max(q_values),
+            "avg_cumulative_reward": sum(rewards) / len(rewards),
+            "total_heuristics": sum(len(r.heuristics) for r in self.records),
+        }
+
+    # ------------------------------------------------------------------
+    # Embedding & Similarity (lightweight — no external deps)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _compute_embedding(text: str) -> list[float]:
+        """
+        Lightweight text embedding using character n-gram hashing.
+        
+        This is intentionally simple — for production, swap in a real
+        embedding model (sentence-transformers, OpenAI embeddings, etc.).
+        
+        To use real embeddings, subclass ExperienceReplay and override
+        _compute_embedding() and _cosine_similarity().
+        """
+        # Character trigram hashing into a fixed-size vector
+        dim = 128
+        vec = [0.0] * dim
+        text_lower = text.lower()
+        for i in range(len(text_lower) - 2):
+            trigram = text_lower[i:i + 3]
+            h = hash(trigram) % dim
+            vec[h] += 1.0
+
+        # L2 normalize
+        magnitude = math.sqrt(sum(x * x for x in vec))
+        if magnitude > 0:
+            vec = [x / magnitude for x in vec]
+        return vec
+
+    @staticmethod
+    def _cosine_similarity(a: list[float], b: list[float]) -> float:
+        """Cosine similarity between two vectors."""
+        if not a or not b or len(a) != len(b):
+            return 0.0
+        dot = sum(x * y for x, y in zip(a, b))
+        mag_a = math.sqrt(sum(x * x for x in a))
+        mag_b = math.sqrt(sum(x * x for x in b))
+        if mag_a == 0 or mag_b == 0:
+            return 0.0
+        return dot / (mag_a * mag_b)
+
+    # ------------------------------------------------------------------
+    # Initial Q-Value Estimation
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _compute_initial_q(trajectory: Trajectory) -> float:
+        """
+        Compute initial Q-value from trajectory performance.
+        
+        Uses a combination of:
+        - Success rate (fraction of steps that improved state)
+        - Total delta (net improvement)
+        - Trajectory length efficiency (shorter = better for same delta)
+        """
+        if not trajectory.steps:
+            return 0.3  # Uninformative prior
+
+        success_rate = trajectory.success_rate
+        total_delta = trajectory.total_delta
+        length = len(trajectory.steps)
+
+        # Normalize total_delta to 0-1 (assuming max meaningful delta is ~10)
+        delta_normalized = max(0.0, min(1.0, total_delta / 10.0))
+
+        # Efficiency bonus: more progress per step = higher Q
+        efficiency = delta_normalized / max(length, 1)
+
+        q = 0.4 * success_rate + 0.4 * delta_normalized + 0.2 * min(efficiency * 5, 1.0)
+        return max(0.0, min(1.0, q))
+
+    # ------------------------------------------------------------------
+    # Capacity Management
+    # ------------------------------------------------------------------
+
+    def _evict(self) -> None:
+        """Evict the lowest Q-value record."""
+        if not self.records:
+            return
+        worst = min(self.records, key=lambda r: r.retrieval_q_value)
+        self.records.remove(worst)
+        logger.debug(
+            f"Experience Replay: Evicted record {worst.id} "
+            f"(q={worst.retrieval_q_value:.3f})"
+        )
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def _save(self) -> None:
+        """Save buffer to disk as JSON."""
+        if not self.persistence_path:
+            return
+        self.persistence_path.parent.mkdir(parents=True, exist_ok=True)
+
+        data = []
+        for record in self.records:
+            data.append({
+                "id": record.id,
+                "retrieval_q_value": record.retrieval_q_value,
+                "task_embedding": record.task_embedding,
+                "trajectory": {
+                    "id": record.trajectory.id,
+                    "task_description": record.trajectory.task_description,
+                    "purpose": record.trajectory.purpose,
+                    "created_at": record.trajectory.created_at,
+                    "cumulative_reward": record.trajectory.cumulative_reward,
+                    "total_delta": record.trajectory.total_delta,
+                    "success_rate": record.trajectory.success_rate,
+                    "num_steps": len(record.trajectory.steps),
+                },
+                "heuristics": [
+                    {
+                        "id": h.id,
+                        "pattern": h.pattern,
+                        "strategy": h.strategy,
+                        "steps": h.steps,
+                        "tier": h.tier.value,
+                        "q_value": h.q_value,
+                        "times_used": h.times_used,
+                        "times_succeeded": h.times_succeeded,
+                    }
+                    for h in record.heuristics
+                ],
+            })
+
+        with open(self.persistence_path, "w") as f:
+            json.dump(data, f, indent=2, default=str)
+
+    def _load(self) -> None:
+        """Load buffer from disk."""
+        if not self.persistence_path or not self.persistence_path.exists():
+            return
+        try:
+            with open(self.persistence_path) as f:
+                data = json.load(f)
+
+            for entry in data:
+                traj_data = entry["trajectory"]
+                trajectory = Trajectory(
+                    task_description=traj_data["task_description"],
+                    purpose=traj_data["purpose"],
+                    id=traj_data["id"],
+                    created_at=traj_data.get("created_at", time.time()),
+                )
+                heuristics = [
+                    Heuristic(
+                        id=h["id"],
+                        pattern=h["pattern"],
+                        strategy=h["strategy"],
+                        steps=h["steps"],
+                        tier=MemoryTier(h["tier"]),
+                        q_value=h["q_value"],
+                        times_used=h.get("times_used", 0),
+                        times_succeeded=h.get("times_succeeded", 0),
+                    )
+                    for h in entry.get("heuristics", [])
+                ]
+                record = MemoryRecord(
+                    id=entry["id"],
+                    trajectory=trajectory,
+                    heuristics=heuristics,
+                    task_embedding=entry.get("task_embedding"),
+                    retrieval_q_value=entry.get("retrieval_q_value", 0.5),
+                )
+                self.records.append(record)
+
+            logger.info(f"Experience Replay: Loaded {len(self.records)} records from disk")
+        except Exception as e:
+            logger.error(f"Experience Replay: Failed to load from disk: {e}")