Lyra/lyra/core/self_evolution.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
import json
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

@dataclass
class EvolutionMetrics:
    """Tracks how Lyra is evolving over time."""
    conversation_satisfaction: float = 0.0
    learning_rate_adaptation: float = 0.0
    personality_drift: float = 0.0
    knowledge_expansion: float = 0.0
    emotional_growth: float = 0.0
    social_adaptation: float = 0.0
    creativity_index: float = 0.0
    coherence_score: float = 0.0

class SelfEvolutionEngine(nn.Module):
    """
    Core self-evolution system that allows Lyra to adapt and grow like a real person.

    This system monitors her performance, emotional state, social interactions,
    and continuously adapts her neural weights, personality traits, and behavior patterns.
    """

    def __init__(
        self,
        model_dim: int = 768,
        evolution_rate: float = 0.001,
        adaptation_threshold: float = 0.7,
        personality_plasticity: float = 0.1,
        memory_capacity: int = 10000,
        device: Optional[torch.device] = None
    ):
        super().__init__()

        self.model_dim = model_dim
        self.evolution_rate = evolution_rate
        self.adaptation_threshold = adaptation_threshold
        self.personality_plasticity = personality_plasticity
        self.memory_capacity = memory_capacity
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Evolution networks
        self.adaptation_network = nn.Sequential(
            nn.Linear(model_dim * 2, model_dim),
            nn.LayerNorm(model_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(model_dim, model_dim // 2),
            nn.LayerNorm(model_dim // 2),
            nn.GELU(),
            nn.Linear(model_dim // 2, model_dim)
        )

        # Self-reflection mechanism
        self.reflection_head = nn.MultiheadAttention(
            embed_dim=model_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        )

        # Meta-learning controller
        self.meta_controller = nn.Sequential(
            nn.Linear(model_dim, model_dim // 2),
            nn.ReLU(),
            nn.Linear(model_dim // 2, 5)  # 5 evolution parameters
        )

        # Experience memory buffer
        self.experience_buffer = []
        self.evolution_history = []

        # Evolution metrics
        self.metrics = EvolutionMetrics()

        # Adaptive learning rate
        self.adaptive_lr = torch.nn.Parameter(torch.tensor(evolution_rate))

        self.to(self.device)

    def forward(
        self,
        current_state: torch.Tensor,
        context: torch.Tensor,
        feedback_signal: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Execute one step of self-evolution.

        Args:
            current_state: Current model hidden state
            context: Conversation/interaction context
            feedback_signal: Optional feedback from environment

        Returns:
            evolved_state: Updated model state
            evolution_info: Information about the evolution step
        """
        batch_size, seq_len, dim = current_state.shape

        # Self-reflection: Let Lyra examine her own thoughts
        reflected_state, attention_weights = self.reflection_head(
            current_state, current_state, current_state
        )

        # Combine current state with reflection
        combined_state = torch.cat([current_state, reflected_state], dim=-1)

        # Generate adaptation signal
        adaptation_signal = self.adaptation_network(combined_state)

        # Meta-learning: Adjust evolution parameters based on context
        meta_params = self.meta_controller(context.mean(dim=1))  # [batch, 5]

        # Apply evolution with meta-learned parameters
        evolution_strength = torch.sigmoid(meta_params[:, 0:1]).unsqueeze(1)  # [batch, 1, 1]
        personality_shift = torch.tanh(meta_params[:, 1:2]).unsqueeze(1)
        learning_adaptation = torch.sigmoid(meta_params[:, 2:3]).unsqueeze(1)
        emotional_weight = torch.sigmoid(meta_params[:, 3:4]).unsqueeze(1)
        creativity_factor = torch.sigmoid(meta_params[:, 4:5]).unsqueeze(1)

        # Evolve the state
        evolved_state = current_state + (
            evolution_strength * self.adaptive_lr * adaptation_signal +
            personality_shift * self.personality_plasticity * reflected_state +
            emotional_weight * 0.1 * torch.randn_like(current_state) * learning_adaptation
        )

        # Apply feedback if available
        if feedback_signal is not None:
            feedback_weight = torch.sigmoid(feedback_signal)
            evolved_state = evolved_state * feedback_weight + current_state * (1 - feedback_weight)

        # Store experience for future learning
        experience = {
            'state': current_state.detach().cpu(),
            'context': context.detach().cpu(),
            'evolution': evolved_state.detach().cpu(),
            'meta_params': meta_params.detach().cpu(),
            'timestamp': torch.tensor(float(torch.rand(1)))
        }
        self.store_experience(experience)

        # Update metrics
        evolution_info = self.update_metrics(
            current_state, evolved_state, meta_params, attention_weights
        )

        return evolved_state, evolution_info

    def store_experience(self, experience: Dict[str, torch.Tensor]):
        """Store experience in memory buffer for future learning."""
        if len(self.experience_buffer) >= self.memory_capacity:
            # Remove oldest experience
            self.experience_buffer.pop(0)

        self.experience_buffer.append(experience)

    def update_metrics(
        self,
        old_state: torch.Tensor,
        new_state: torch.Tensor,
        meta_params: torch.Tensor,
        attention_weights: torch.Tensor
    ) -> Dict[str, Any]:
        """Update evolution metrics and track growth."""
        with torch.no_grad():
            # Calculate state change magnitude
            state_change = torch.norm(new_state - old_state, dim=-1).mean()

            # Update metrics
            self.metrics.personality_drift = float(state_change * 0.1)
            self.metrics.learning_rate_adaptation = float(meta_params[:, 2].mean())
            self.metrics.creativity_index = float(meta_params[:, 4].mean())

            # Attention diversity (measure of cognitive flexibility)
            attention_entropy = -torch.sum(
                attention_weights * torch.log(attention_weights + 1e-8), dim=-1
            ).mean()

            evolution_info = {
                'state_change_magnitude': float(state_change),
                'attention_entropy': float(attention_entropy),
                'adaptive_lr': float(self.adaptive_lr),
                'metrics': self.metrics.__dict__.copy()
            }

            self.evolution_history.append(evolution_info)

            return evolution_info

    def evolve_from_conversation(
        self,
        conversation_embedding: torch.Tensor,
        user_satisfaction: float,
        emotional_context: Dict[str, float]
    ):
        """
        Evolve based on a conversation interaction.

        This is where Lyra learns from each conversation like a human would.
        """
        # Convert satisfaction to feedback signal
        satisfaction_tensor = torch.tensor(
            [[user_satisfaction]], device=self.device, dtype=torch.float32
        )

        # Create emotional context tensor
        emotional_values = list(emotional_context.values())
        emotional_tensor = torch.tensor(
            [emotional_values], device=self.device, dtype=torch.float32
        )

        # Evolve based on this interaction
        evolved_embedding, evolution_info = self.forward(
            conversation_embedding.unsqueeze(0),
            emotional_tensor.unsqueeze(0),
            satisfaction_tensor
        )

        # Update conversation satisfaction metric
        self.metrics.conversation_satisfaction = (
            0.9 * self.metrics.conversation_satisfaction + 0.1 * user_satisfaction
        )

        # Adapt learning rate based on satisfaction
        if user_satisfaction > 0.8:
            self.adaptive_lr.data *= 1.01  # Increase learning when doing well
        elif user_satisfaction < 0.3:
            self.adaptive_lr.data *= 0.99  # Decrease when struggling

        # Clamp learning rate
        self.adaptive_lr.data = torch.clamp(self.adaptive_lr.data, 1e-6, 1e-2)

        return evolved_embedding.squeeze(0), evolution_info

    def long_term_evolution(self):
        """
        Perform long-term evolutionary changes based on accumulated experience.

        This happens periodically (like during sleep for humans) to consolidate learning.
        """
        if len(self.experience_buffer) < 100:  # Need sufficient experience
            return

        logger.info("Performing long-term evolution consolidation...")

        # Analyze patterns in stored experiences
        recent_experiences = self.experience_buffer[-100:]

        # Extract patterns
        state_changes = []
        meta_patterns = []

        for exp in recent_experiences:
            state_change = torch.norm(exp['evolution'] - exp['state'], dim=-1).mean()
            state_changes.append(float(state_change))
            meta_patterns.append(exp['meta_params'].mean(0))

        # Update long-term adaptation parameters
        avg_change = np.mean(state_changes)
        if avg_change > 0.1:  # Too much change - stabilize
            self.personality_plasticity *= 0.95
        elif avg_change < 0.01:  # Too little change - increase plasticity
            self.personality_plasticity *= 1.05

        # Clamp plasticity
        self.personality_plasticity = np.clip(self.personality_plasticity, 0.01, 0.3)

        # Update evolution rate based on performance
        recent_satisfaction = self.metrics.conversation_satisfaction
        if recent_satisfaction > 0.7:
            self.evolution_rate *= 0.98  # Slower evolution when performing well
        else:
            self.evolution_rate *= 1.02  # Faster evolution when struggling

        logger.info(f"Evolution update - Plasticity: {self.personality_plasticity:.4f}, "
                   f"Rate: {self.evolution_rate:.6f}, Satisfaction: {recent_satisfaction:.3f}")

    def get_evolution_summary(self) -> Dict[str, Any]:
        """Get a summary of Lyra's evolution and growth."""
        if not self.evolution_history:
            return {"status": "no_evolution_data"}

        recent_history = self.evolution_history[-100:] if len(self.evolution_history) > 100 else self.evolution_history

        return {
            "total_evolution_steps": len(self.evolution_history),
            "current_metrics": self.metrics.__dict__,
            "recent_growth_rate": np.mean([h["state_change_magnitude"] for h in recent_history]),
            "personality_plasticity": self.personality_plasticity,
            "adaptive_learning_rate": float(self.adaptive_lr),
            "experience_buffer_size": len(self.experience_buffer),
            "cognitive_flexibility": np.mean([h["attention_entropy"] for h in recent_history])
        }

    def save_evolution_state(self, path: Path):
        """Save evolution state for persistence."""
        state = {
            "metrics": self.metrics.__dict__,
            "evolution_history": self.evolution_history[-1000:],  # Keep recent history
            "personality_plasticity": self.personality_plasticity,
            "evolution_rate": self.evolution_rate,
            "adaptive_lr": float(self.adaptive_lr),
            "model_state": self.state_dict()
        }

        with open(path, 'w') as f:
            json.dump(state, f, indent=2, default=str)

    def load_evolution_state(self, path: Path):
        """Load evolution state from file."""
        if not path.exists():
            logger.warning(f"Evolution state file not found: {path}")
            return

        try:
            with open(path, 'r') as f:
                state = json.load(f)

            # Restore metrics
            for key, value in state["metrics"].items():
                setattr(self.metrics, key, value)

            self.evolution_history = state.get("evolution_history", [])
            self.personality_plasticity = state.get("personality_plasticity", 0.1)
            self.evolution_rate = state.get("evolution_rate", 0.001)

            if "adaptive_lr" in state:
                self.adaptive_lr.data = torch.tensor(state["adaptive_lr"])

            # Load model state
            if "model_state" in state:
                self.load_state_dict(state["model_state"])

            logger.info(f"Evolution state loaded from {path}")

        except Exception as e:
            logger.error(f"Failed to load evolution state: {e}")