Upload aether/agents.py

aether/agents.py

@@ -0,0 +1,363 @@
+"""
+AETHER Agent Orchestration.
+Integrates:
+- smolagents multi-agent hierarchy (Manager + Workers)
+- MLPO: Multi-agent guided Leader Policy Optimization
+- BabyAGI task creation/prioritization/execution loop
+- Agentic Neural Networks: textual backpropagation
+- Yunjue Agent: Manager/Executor/Developer/Integrator/Merger/Aggregator roles
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, List, Any, Optional, Callable
+import logging
+import time
+from collections import deque
+
+logger = logging.getLogger("AETHER.Agents")
+
+
+class AgentRole:
+    """Role definitions inspired by Yunjue Agent multi-agent system."""
+    MANAGER = "manager"
+    EXECUTOR = "executor"
+    DEVELOPER = "developer"
+    INTEGRATOR = "integrator"
+    MERGER = "merger"
+    AGGREGATOR = "aggregator"
+    RESEARCHER = "researcher"
+
+
+class BaseAgent(nn.Module):
+    """Base agent with policy network. Implements MLPO-style leader policy."""
+    
+    def __init__(self, role: str, hidden_dim: int = 128, 
+                 vocab_size: int = 32000):
+        super().__init__()
+        self.role = role
+        self.hidden_dim = hidden_dim
+        
+        self.encoder = nn.Sequential(
+            nn.Embedding(vocab_size, hidden_dim),
+            nn.LSTM(hidden_dim, hidden_dim, batch_first=True),
+        )
+        self.policy_head = nn.Linear(hidden_dim, hidden_dim)
+        self.value_head = nn.Linear(hidden_dim, 1)
+        
+        self.task_history: deque = deque(maxlen=100)
+        self.performance_log: List[float] = []
+        
+    def forward(self, input_ids: torch.Tensor) -> Dict[str, torch.Tensor]:
+        embeds = self.encoder[0](input_ids)
+        lstm_out, _ = self.encoder[1](embeds)
+        hidden = lstm_out[:, -1, :]
+        
+        return {
+            "policy_logits": self.policy_head(hidden),
+            "value": self.value_head(hidden),
+            "hidden": hidden,
+        }
+    
+    def act(self, observation: str) -> str:
+        self.task_history.append({
+            "observation": observation,
+            "timestamp": time.time(),
+        })
+        
+        role_actions = {
+            AgentRole.MANAGER: f"[MANAGER] Decomposing task: '{observation[:50]}...'",
+            AgentRole.EXECUTOR: f"[EXECUTOR] Executing: '{observation[:50]}...'",
+            AgentRole.DEVELOPER: f"[DEVELOPER] Synthesizing tool for: '{observation[:50]}...'",
+            AgentRole.INTEGRATOR: f"[INTEGRATOR] Integrating components for: '{observation[:50]}...'",
+            AgentRole.MERGER: f"[MERGER] Consolidating tools for: '{observation[:50]}...'",
+            AgentRole.AGGREGATOR: f"[AGGREGATOR] Aggregating results for: '{observation[:50]}...'",
+            AgentRole.RESEARCHER: f"[RESEARCHER] Exploring knowledge for: '{observation[:50]}...'",
+        }
+        
+        return role_actions.get(self.role, f"[{self.role.upper()}] Processing: '{observation}'")
+    
+    def update(self, reward: float):
+        self.performance_log.append(reward)
+
+
+class HierarchicalAgent(nn.Module):
+    """
+    HiMAC-style hierarchical agent with Macro-Policy and Micro-Policy.
+    Macro: generates blueprint (sub-goals)
+    Micro: executes atomic actions conditioned on blueprint
+    """
+    
+    def __init__(self, macro_dim: int = 256, micro_dim: int = 128,
+                 num_subgoals: int = 5):
+        super().__init__()
+        self.macro_dim = macro_dim
+        self.micro_dim = micro_dim
+        self.num_subgoals = num_subgoals
+        
+        self.macro_encoder = nn.LSTM(macro_dim, macro_dim, batch_first=True)
+        self.macro_decoder = nn.LSTM(macro_dim, macro_dim, batch_first=True)
+        self.subgoal_head = nn.Linear(macro_dim, num_subgoals)
+        self.termination_token = nn.Parameter(torch.randn(macro_dim))
+        
+        self.micro_encoder = nn.LSTM(micro_dim + macro_dim, micro_dim, batch_first=True)
+        self.action_head = nn.Linear(micro_dim, 50)
+        
+        self.current_blueprint: Optional[List[str]] = None
+        self.active_subgoal_idx = 0
+    
+    def generate_blueprint(self, task_embedding: torch.Tensor) -> List[str]:
+        batch_size = task_embedding.size(0)
+        hidden = (torch.zeros(1, batch_size, self.macro_dim),
+                  torch.zeros(1, batch_size, self.macro_dim))
+        
+        blueprints = []
+        input_token = task_embedding.unsqueeze(1)
+        
+        for _ in range(self.num_subgoals):
+            out, hidden = self.macro_decoder(input_token, hidden)
+            subgoal_logits = self.subgoal_head(out.squeeze(1))
+            subgoal_id = torch.argmax(subgoal_logits, dim=-1)
+            
+            similarity = torch.cosine_similarity(out.squeeze(1), 
+                                                  self.termination_token.unsqueeze(0))
+            if similarity.item() > 0.9:
+                break
+            
+            blueprints.append(f"subgoal_{subgoal_id.item()}")
+            input_token = out
+        
+        self.current_blueprint = blueprints
+        self.active_subgoal_idx = 0
+        return blueprints
+    
+    def execute_action(self, observation: torch.Tensor,
+                      blueprint: Optional[List[str]] = None) -> torch.Tensor:
+        if blueprint is not None:
+            self.current_blueprint = blueprint
+        
+        if not self.current_blueprint:
+            return torch.zeros(1, 50)
+        
+        active_subgoal = self.current_blueprint[
+            min(self.active_subgoal_idx, len(self.current_blueprint) - 1)
+        ]
+        
+        subgoal_embed = torch.randn(1, self.macro_dim)
+        combined = torch.cat([observation, subgoal_embed], dim=-1)
+        
+        out, _ = self.micro_encoder(combined.unsqueeze(1))
+        action_logits = self.action_head(out.squeeze(1))
+        
+        return action_logits
+    
+    def advance_subgoal(self):
+        self.active_subgoal_idx += 1
+    
+    def reset(self):
+        self.current_blueprint = None
+        self.active_subgoal_idx = 0
+
+
+class BabyAGILoop:
+    """BabyAGI-inspired task-driven autonomous loop."""
+    
+    def __init__(self, objective: str, max_iterations: int = 50):
+        self.objective = objective
+        self.max_iterations = max_iterations
+        self.task_list: deque = deque()
+        self.completed_tasks: List[Dict] = []
+        self.results: Dict[int, Any] = {}
+        self.iteration = 0
+    
+    def create_tasks(self, previous_result: str, task_description: str) -> List[str]:
+        new_tasks = [
+            f"Sub-task {len(self.task_list) + i}: Analyze {previous_result[:30]}..."
+            for i in range(3)
+        ]
+        return new_tasks
+    
+    def prioritize_tasks(self) -> List[str]:
+        tasks = list(self.task_list)
+        scores = []
+        for task in tasks:
+            overlap = sum(1 for word in self.objective.lower().split() 
+                         if word in task.lower())
+            scores.append(overlap)
+        
+        sorted_tasks = [t for _, t in sorted(zip(scores, tasks), reverse=True)]
+        return sorted_tasks
+    
+    def execute_task(self, task: str, agent: BaseAgent) -> str:
+        result = agent.act(task)
+        self.completed_tasks.append({
+            "task": task,
+            "result": result,
+            "iteration": self.iteration,
+        })
+        return result
+    
+    def run(self, execution_agent: BaseAgent) -> Dict[str, Any]:
+        self.task_list.append(self.objective)
+        
+        while self.iteration < self.max_iterations and self.task_list:
+            prioritized = self.prioritize_tasks()
+            self.task_list = deque(prioritized)
+            
+            current_task = self.task_list.popleft()
+            previous_result = self.completed_tasks[-1]["result"] if self.completed_tasks else ""
+            
+            result = self.execute_task(current_task, execution_agent)
+            self.results[self.iteration] = result
+            
+            new_tasks = self.create_tasks(result, current_task)
+            for t in new_tasks:
+                if t not in self.task_list:
+                    self.task_list.append(t)
+            
+            self.iteration += 1
+            
+            logger.info(f"BabyAGI iteration {self.iteration}: "
+                       f"tasks_remaining={len(self.task_list)}, "
+                       f"completed={len(self.completed_tasks)}")
+        
+        return {
+            "completed_tasks": self.completed_tasks,
+            "results": self.results,
+            "iterations": self.iteration,
+            "objective": self.objective,
+        }
+
+
+class AetherAgentOrchestrator(nn.Module):
+    """
+    Multi-agent orchestrator combining:
+    - smolagents hierarchical delegation
+    - MLPO: train single leader, peers untrained
+    - Agentic Neural Networks: textual backpropagation
+    - CoMAS: co-evolving via interaction rewards
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        self.agents: Dict[str, BaseAgent] = nn.ModuleDict({
+            "manager": BaseAgent(AgentRole.MANAGER, hidden_dim=config.macro_policy_dim),
+            "executor": BaseAgent(AgentRole.EXECUTOR, hidden_dim=config.micro_policy_dim),
+            "developer": BaseAgent(AgentRole.DEVELOPER, hidden_dim=config.micro_policy_dim),
+            "researcher": BaseAgent(AgentRole.RESEARCHER, hidden_dim=config.micro_policy_dim),
+        })
+        
+        self.leader = BaseAgent(AgentRole.MANAGER, hidden_dim=config.macro_policy_dim)
+        
+        self.hierarchical = HierarchicalAgent(
+            macro_dim=config.macro_policy_dim,
+            micro_dim=config.micro_policy_dim,
+        )
+        
+        self.routing_weights = nn.Parameter(torch.ones(len(self.agents)))
+        self.aggregation_gate = nn.Softmax(dim=0)
+        
+        self.agent_tasks: Dict[str, BabyAGILoop] = {}
+        
+        self.task_count = 0
+        self.agent_interactions: List[Dict] = []
+    
+    def forward(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        task_embed = torch.randn(1, self.config.macro_policy_dim)
+        blueprint = self.hierarchical.generate_blueprint(task_embed)
+        
+        routing_probs = self.aggregation_gate(self.routing_weights)
+        
+        agent_outputs = {}
+        for i, (name, agent) in enumerate(self.agents.items()):
+            if name == "manager":
+                continue
+            
+            weight = routing_probs[i].item()
+            if weight < 0.15:
+                continue
+            
+            sub_task = blueprint[min(i, len(blueprint) - 1)] if blueprint else task
+            output = agent.act(f"[{name}] {sub_task}")
+            agent_outputs[name] = {
+                "output": output,
+                "weight": weight,
+                "sub_task": sub_task,
+            }
+        
+        synthesized = self.leader.act(
+            f"Synthesize: {task} with inputs: {list(agent_outputs.keys())}"
+        )
+        
+        self.agent_interactions.append({
+            "task": task,
+            "blueprint": blueprint,
+            "agent_outputs": agent_outputs,
+            "leader_synthesis": synthesized,
+            "routing_probs": routing_probs.detach().cpu().tolist(),
+            "timestamp": time.time(),
+        })
+        
+        self.task_count += 1
+        
+        return {
+            "output": synthesized,
+            "blueprint": blueprint,
+            "agent_outputs": agent_outputs,
+            "routing_weights": routing_probs.detach().cpu().tolist(),
+        }
+    
+    def execute(self, task: str, kg_context: Any, context: Dict[str, Any]) -> Dict[str, Any]:
+        return self.forward(task, context)
+    
+    def textual_backprop(self, global_gradient: str, 
+                        performance_feedback: float,
+                        beta: float = 0.5) -> Dict[str, str]:
+        updates = {}
+        for name, agent in self.agents.items():
+            local_grad = f"{global_gradient} + Agent {name} performance: {performance_feedback}"
+            
+            if hasattr(agent, 'previous_gradient'):
+                blended = f"0.7*{local_grad} + 0.3*{agent.previous_gradient}"
+            else:
+                blended = local_grad
+            
+            agent.previous_gradient = blended
+            updates[name] = blended
+        
+        self.routing_weights.data += performance_feedback * 0.01
+        
+        return updates
+    
+    def co_evolve_interactions(self) -> List[Dict]:
+        rewards = []
+        
+        for interaction in self.agent_interactions[-10:]:
+            num_agents_involved = len(interaction.get("agent_outputs", {}))
+            blueprint_complexity = len(interaction.get("blueprint", []))
+            
+            reward = num_agents_involved * 0.1 + min(blueprint_complexity * 0.05, 0.5)
+            rewards.append({
+                "interaction_id": id(interaction),
+                "reward": reward,
+                "agents_involved": num_agents_involved,
+            })
+        
+        return rewards
+    
+    def run_babyagi(self, objective: str, max_iterations: int = 20) -> Dict[str, Any]:
+        loop = BabyAGILoop(objective, max_iterations)
+        result = loop.run(self.agents["manager"])
+        self.agent_tasks[objective] = loop
+        return result
+    
+    def stats(self) -> Dict[str, Any]:
+        return {
+            "total_tasks": self.task_count,
+            "num_agents": len(self.agents),
+            "total_interactions": len(self.agent_interactions),
+            "routing_weights": self.routing_weights.detach().cpu().tolist(),
+            "active_tasks": len(self.agent_tasks),
+        }