diff --git a/jspaceai/__init__.py b/jspaceai/__init__.py index ee26e7b..e7aba7d 100644 --- a/jspaceai/__init__.py +++ b/jspaceai/__init__.py @@ -35,6 +35,20 @@ from .jlens import ( DirectedModulation, CounterfactualReflection, ) +from .consensus import ( + ConsensusSlot, + ConsensusSnapshot, +) +from .events import ( + WorkspaceEvent, + ActionEvent, +) +from .memory import ( + MemoryRecord, + MemoryStore, + InMemoryVectorMemoryStore, + Hippocampus, +) from .multimodal import ( MultimodalConfig, MultimodalJSpaceModel, @@ -42,6 +56,19 @@ from .multimodal import ( AudioEncoder, AudioDecoder, TextEncoder, TextDecoder, ) +from .continual import ( + OnlineLanguageLearner, +) +from .policy import ( + ACTION_LABELS, + compose_action_params, + ReflexRule, + ActionGate, + ActionValueModel, + MotorController, + ActionDecision, + ActionPolicy, +) from .realtime import ( Frame, CameraStream, @@ -66,12 +93,12 @@ from .platform import ( ) from .embodied import ( MouseActuator, KeyboardActuator, AudioActuator, ScreenActuator, - Cerebellum, CentralNervousSystem, Hippocampus, BasalGanglia, + Cerebellum, CentralNervousSystem, BasalGanglia, EmbodiedAgent, ) -from .autonomous import ( - CuriosityDrive, PersistentState, SelfModel, MetaLearner, - AutonomousMind, +from .runtime import ( + CuriosityDrive, RuntimeStateStore, PersistentState, SelfModel, MetaLearner, + WorkspaceRuntime, AutonomousMind, ) from .modules import ( ExternalModule, SmallModelModule, KnowledgeBaseModule, @@ -94,11 +121,18 @@ __all__ = [ # J-lens 可解释性 "JLensConfig", "JLensProbe", "JLensSuite", "WorkspaceAblator", "DirectedModulation", "CounterfactualReflection", + "ConsensusSlot", "ConsensusSnapshot", + "WorkspaceEvent", "ActionEvent", + "MemoryRecord", "MemoryStore", "InMemoryVectorMemoryStore", # 多模态 "MultimodalConfig", "MultimodalJSpaceModel", "VisualEncoder", "VisualDecoder", "AudioEncoder", "AudioDecoder", "TextEncoder", "TextDecoder", + # 动作策略 + "ACTION_LABELS", "compose_action_params", "ReflexRule", + "ActionGate", "ActionValueModel", "MotorController", + "ActionDecision", "ActionPolicy", # 实时 I/O "Frame", "CameraStream", "MicrophoneStream", "AudioPlayer", "MultimodalStream", "SensoryMotorLoop", @@ -114,11 +148,12 @@ __all__ = [ "Cerebellum", "CentralNervousSystem", "Hippocampus", "BasalGanglia", "EmbodiedAgent", # 自主心智（最重要的能力） - "CuriosityDrive", "PersistentState", "SelfModel", "MetaLearner", - "AutonomousMind", + "CuriosityDrive", "RuntimeStateStore", "PersistentState", + "SelfModel", "MetaLearner", "WorkspaceRuntime", "AutonomousMind", # 外挂模块系统（可热插拔） "ExternalModule", "SmallModelModule", "KnowledgeBaseModule", "ToolModule", "ModuleDock", # 自主进化 "EvolutionTrainer", + "OnlineLanguageLearner", ] diff --git a/jspaceai/autonomous.py b/jspaceai/autonomous.py index 04a4e27..808fc57 100644 --- a/jspaceai/autonomous.py +++ b/jspaceai/autonomous.py @@ -1,389 +1,25 @@ """ -自主心智（AutonomousMind）—— 智慧最重要的能力 +Backward-compatible autonomous runtime exports. -四个核心能力： - 1. CuriosityDrive: 好奇心驱动的主动探索（内在奖励） - 2. PersistentState: 跨会话状态持久化（海洋不蒸发） - 3. SelfModel: 自我模型（知道自己会什么不会什么） - 4. MetaLearner: 元学习（学会如何学习） - -合起来 = AutonomousMind，永不停止的自主进化。 +The primary implementation now lives in runtime.py so the workspace loop has a +single home. This module keeps the older import path stable. """ -from __future__ import annotations +from .runtime import ( + CuriosityDrive, + RuntimeStateStore, + PersistentState, + SelfModel, + MetaLearner, + WorkspaceRuntime, + AutonomousMind, +) -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -import json -import time -from pathlib import Path -from typing import Optional -from dataclasses import dataclass -from collections import deque - - -class CuriosityDrive(nn.Module): - """好奇心驱动——基于预测误差的内在奖励 - - 模型有世界模型预测下一状态，预测误差=好奇心=内在奖励。 - 模型被驱动去探索"预测不准"的区域。学会后好奇心降低，转向新区域。 - """ - - def __init__(self, workspace_dim: int, action_dim: int = 5, hidden_dim: int = 64): - super().__init__() - self.workspace_dim = workspace_dim - self.world_model = nn.Sequential( - nn.Linear(workspace_dim + action_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, workspace_dim), - ) - self.state_history: deque = deque(maxlen=500) - self.prediction_error_ema = 0.1 - - def predict_next(self, w, action): - return self.world_model(torch.cat([w, action], dim=-1)) - - def compute_curiosity(self, w_current, action, w_next): - with torch.no_grad(): - w_pred = self.predict_next(w_current, action) - pred_error = F.mse_loss(w_pred, w_next).item() - w_np = w_current[0].cpu().numpy() - novelty = self._compute_novelty(w_np) - progress = max(0, pred_error - self.prediction_error_ema * 0.9) - self.prediction_error_ema = 0.95 * self.prediction_error_ema + 0.05 * pred_error - curiosity = progress + 0.3 * novelty - self.state_history.append(w_np.copy()) - return curiosity - - def _compute_novelty(self, w): - if len(self.state_history) < 5: - return 1.0 - history = list(self.state_history)[-100:] - distances = [np.linalg.norm(w - h) for h in history] - return float(min(1.0, min(distances) / 2.0)) - - def train_world_model(self, w_current, action, w_next): - w_pred = self.predict_next(w_current, action.detach()) - return F.mse_loss(w_pred, w_next.detach()) - - -class PersistentState: - """跨会话状态持久化——让海洋不蒸发 - - 保存 workspace + 专家状态 + 海马体 + 基底神经节 + 好奇心历史 + 自我模型。 - """ - - def __init__(self, save_dir: Path): - self.save_dir = Path(save_dir) - self.save_dir.mkdir(parents=True, exist_ok=True) - self.state_file = self.save_dir / 'mind_state.json' - self.tensors_file = self.save_dir / 'mind_tensors.npz' - - def save(self, state: dict): - tensors = {} - if 'w' in state: - tensors['w'] = state['w'].cpu().numpy() - if 'm' in state: - for i, m in enumerate(state['m']): - if m is not None: - tensors[f'm_{i}'] = m.cpu().numpy() - if 'basal_ganglia' in state and state['basal_ganglia'] is not None: - tensors['bg_weights'] = state['basal_ganglia'] - if 'curiosity_history' in state: - tensors['curiosity_history'] = np.array(state['curiosity_history']) - if tensors: - np.savez(self.tensors_file, **tensors) - - json_state = { - 'step_count': state.get('step_count', 0), - 'total_runtime': state.get('total_runtime', 0.0), - 'self_model': state.get('self_model', {}), - 'saved_at': time.time(), - } - self.state_file.write_text(json.dumps(json_state, indent=2, default=str)) - - def load(self) -> Optional[dict]: - if not self.state_file.exists(): - return None - result = {} - if self.tensors_file.exists(): - data = np.load(self.tensors_file, allow_pickle=True) - if 'w' in data: - result['w'] = torch.tensor(data['w']) - ms = {} - for key in data.files: - if key.startswith('m_'): - idx = int(key.split('_')[1]) - ms[idx] = torch.tensor(data[key]) - if ms: - result['m'] = [ms[i] for i in sorted(ms.keys())] - if 'bg_weights' in data: - result['basal_ganglia'] = data['bg_weights'] - if 'curiosity_history' in data: - result['curiosity_history'] = data['curiosity_history'].tolist() - json_state = json.loads(self.state_file.read_text()) - result.update(json_state) - return result - - -class SelfModel: - """自我模型——知道自己会什么、不会什么 - - 对每个能力域维护置信度（0-1），通过历史成功率更新。 - 不知道时主动学习（好奇心驱动）。 - """ - - def __init__(self, capabilities=None): - if capabilities is None: - capabilities = ['visual', 'audio', 'text', 'motor_mouse', - 'motor_keyboard', 'memory', 'prediction'] - self.capabilities = capabilities - self.confidence = {c: 0.0 for c in capabilities} - self.attempts = {c: 0 for c in capabilities} - self.recent_results = {c: deque(maxlen=20) for c in capabilities} - - def record_attempt(self, capability, success): - if capability not in self.confidence: - return - self.attempts[capability] += 1 - self.recent_results[capability].append(success) - recent = list(self.recent_results[capability]) - if recent: - weights = np.linspace(0.5, 1.0, len(recent)) - self.confidence[capability] = float(np.average(recent, weights=weights)) - - def get_weakness(self): - return min(self.confidence, key=self.confidence.get) - - def get_strength(self): - return max(self.confidence, key=self.confidence.get) - - def knows(self, capability, threshold=0.5): - return self.confidence.get(capability, 0.0) > threshold - - def summary(self): - return { - 'capabilities': dict(self.confidence), - 'attempts': dict(self.attempts), - 'strength': self.get_strength(), - 'weakness': self.get_weakness(), - } - - -class MetaLearner: - """元学习——学会如何学习 - - 每个能力域有自适应学习率：进步快→加快，停滞→减慢。 - 记录什么学习策略有效。 - """ - - def __init__(self, capabilities=None): - if capabilities is None: - capabilities = ['visual', 'audio', 'text', 'motor_mouse', - 'motor_keyboard', 'memory', 'prediction'] - self.learning_rates = {c: 1e-3 for c in capabilities} - self.loss_history = {c: deque(maxlen=20) for c in capabilities} - self.strategy_scores = { - 'predict_next': 0.5, 'replay': 0.5, - 'explore': 0.5, 'imitate': 0.5, - } - - def get_lr(self, capability): - return self.learning_rates.get(capability, 1e-3) - - def record_loss(self, capability, loss): - if capability not in self.loss_history: - return - self.loss_history[capability].append(loss) - history = list(self.loss_history[capability]) - if len(history) < 5: - return - recent_avg = np.mean(history[-5:]) - old_avg = np.mean(history[-10:-5]) if len(history) >= 10 else recent_avg - improvement = (old_avg - recent_avg) / max(old_avg, 1e-8) - lr = self.learning_rates[capability] - if improvement > 0.05: - lr *= 1.1 - elif improvement < 0.01: - lr *= 0.9 - self.learning_rates[capability] = max(1e-5, min(1e-2, lr)) - - def best_strategy(self): - return max(self.strategy_scores, key=self.strategy_scores.get) - - def reward_strategy(self, strategy, reward): - if strategy in self.strategy_scores: - self.strategy_scores[strategy] = ( - 0.9 * self.strategy_scores[strategy] + 0.1 * reward - ) - - -class AutonomousMind: - """自主心智——永不停止的自主进化 - - 整合好奇心 + 持久化 + 自我模型 + 元学习。 - 核心循环：感知→自我评估→好奇探索→行动→观察→学习→存盘 - """ - - def __init__(self, agent, save_dir='outputs/mind', device='cpu'): - self.agent = agent - self.device = device - self.config = agent.config - - self.curiosity = CuriosityDrive( - workspace_dim=self.config.workspace_dim, action_dim=5, - ).to(device) - self.persistence = PersistentState(Path(save_dir)) - self.self_model = SelfModel() - self.meta_learner = MetaLearner() - - self.step_count = 0 - self.total_runtime = 0.0 - self.start_time = time.time() - self.running = False - self.curiosity_history = deque(maxlen=1000) - self.curiosity_optimizer = torch.optim.Adam( - self.curiosity.parameters(), lr=1e-3 - ) - - self._load_state() - - def _load_state(self): - state = self.persistence.load() - if state is None: - print(" [心智] 全新启动") - return - print(f" [心智] 恢复状态: step={state.get('step_count', 0)}") - if 'w' in state: - self.agent.state['w'] = state['w'].to(self.device) - if 'm' in state: - for i, m in enumerate(state['m']): - if m is not None and i < len(self.agent.state['m']): - self.agent.state['m'][i] = m.to(self.device) - if 'basal_ganglia' in state and hasattr(self.agent, 'basal_ganglia'): - self.agent.basal_ganglia.action_weights = state['basal_ganglia'] - self.step_count = state.get('step_count', 0) - self.total_runtime = state.get('total_runtime', 0.0) - - def save_state(self): - state = { - 'w': self.agent.state['w'], - 'm': self.agent.state['m'], - 'basal_ganglia': getattr(self.agent.basal_ganglia, 'action_weights', None) - if hasattr(self.agent, 'basal_ganglia') else None, - 'curiosity_history': list(self.curiosity.state_history), - 'step_count': self.step_count, - 'total_runtime': self.total_runtime + (time.time() - self.start_time), - 'self_model': {'confidence': self.self_model.confidence}, - } - self.persistence.save(state) - - def step(self) -> dict: - # 1. 感知 + 思考 - sensory_data = self.agent.perceive() - w_before = self.agent.state['w'].clone() - weakness = self.self_model.get_weakness() - strength = self.self_model.get_strength() - - # 2. 行动 - action_info = self.agent.decide_and_act(w_before, sensory_data.get('modality', 'idle')) - w_after, modality = self.agent.think(sensory_data) - - # 3. 好奇心 - action_tensor = torch.tensor(action_info['action_params'], - dtype=torch.float32).unsqueeze(0).to(self.device) - curiosity_reward = self.curiosity.compute_curiosity(w_before, action_tensor, w_after) - self.curiosity_history.append(curiosity_reward) - - # 4. 训练世界模型 - world_loss = self.curiosity.train_world_model(w_before, action_tensor, w_after) - self.curiosity_optimizer.zero_grad() - world_loss.backward() - self.curiosity_optimizer.step() - - # 5. 评估成功度 - w_stability = 1.0 - min(1.0, abs(w_after.norm().item() - w_before.norm().item())) - success = (0.3 * float(action_info['executed']) + - 0.4 * min(1.0, curiosity_reward) + 0.3 * w_stability) - - # 6. 更新自我模型 - cap_map = {'image': 'visual', 'screen': 'visual', 'audio': 'audio', - 'text': 'text', 'keyboard': 'text', 'mouse': 'motor_mouse', 'idle': 'prediction'} - cap = cap_map.get(modality, 'prediction') - self.self_model.record_attempt(cap, success) - self.self_model.record_attempt('prediction', 1.0 - min(1.0, world_loss.item())) - - # 7. 元学习 - self.meta_learner.record_loss(cap, world_loss.item()) - if curiosity_reward > 0.3: - self.meta_learner.reward_strategy('explore', curiosity_reward) - - # 8. 记忆 + 基底神经节 - self.agent.remember(w_after, { - 'modality': modality, 'curiosity': curiosity_reward, - 'success': success, 'step': self.step_count, - }) - self.agent.learn(w_after, np.array(action_info['action_params']), reward=curiosity_reward) - - self.step_count += 1 - return { - 'step': self.step_count, 'modality': modality, - 'w_norm': w_after.norm().item(), 'curiosity': curiosity_reward, - 'world_loss': world_loss.item(), 'success': success, - 'weakness': weakness, 'strength': strength, - 'self_confidence': dict(self.self_model.confidence), - 'best_strategy': self.meta_learner.best_strategy(), - 'memory_count': self.agent.hippocampus.size() if self.agent.hippocampus else 0, - } - - def run(self, n_steps=100, interval=0.2, save_every=50, on_step=None): - self.running = True - self.start_time = time.time() - self.agent.senses.start() - print(f"\n自主心智启动 | 总步数: {self.step_count} | 保存间隔: {save_every}") - print("=" * 60) - - log = [] - try: - for _ in range(n_steps): - if not self.running: - break - info = self.step() - log.append(info) - if on_step: - on_step(info) - elif info['step'] % 10 == 0: - print(f" step {info['step']:4d} | mod {info['modality']:8s} | " - f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | " - f"success {info['success']:.2f} | weak={info['weakness']} | " - f"mem {info['memory_count']}") - if info['step'] % save_every == 0: - self.save_state() - time.sleep(interval) - except KeyboardInterrupt: - print("\n用户中断") - finally: - self.running = False - self.agent.senses.stop() - if hasattr(self.agent, 'audio_actuator'): - self.agent.audio_actuator.stop() - self.save_state() - self.total_runtime += time.time() - self.start_time - return log - - def introspect(self) -> str: - """内省——自我报告""" - sm = self.self_model.summary() - avg_curio = np.mean(list(self.curiosity_history)) if self.curiosity_history else 0 - report = f"=== 自主心智内省 ===\n步数: {self.step_count}\n运行: {self.total_runtime:.0f}s\n" - report += f"记忆: {self.agent.hippocampus.size() if self.agent.hippocampus else 0}\n" - report += f"平均好奇心: {avg_curio:.3f}\n\n自我认知:\n" - for cap, conf in sm['capabilities'].items(): - bar = '█' * int(conf * 20) - report += f" {cap:15s}: {conf:.2f} {bar}\n" - report += f"\n最强: {sm['strength']}\n最弱: {sm['weakness']}\n" - report += f"最佳策略: {self.meta_learner.best_strategy()}\n" - return report +__all__ = [ + "CuriosityDrive", + "RuntimeStateStore", + "PersistentState", + "SelfModel", + "MetaLearner", + "WorkspaceRuntime", + "AutonomousMind", +] diff --git a/jspaceai/consensus.py b/jspaceai/consensus.py new file mode 100644 index 0000000..addd205 --- /dev/null +++ b/jspaceai/consensus.py @@ -0,0 +1,83 @@ +""" +Consensus state utilities. + +The workspace vector is still the primary state, but these helpers expose a +small structured summary that other subsystems can share without needing to +decode the full latent every time. +""" +from __future__ import annotations + +from dataclasses import dataclass +import math +import torch + + +@dataclass +class ConsensusSlot: + index: int + label: str + weight: float + + +@dataclass +class ConsensusSnapshot: + modality: str + workspace_norm: float + confidence: float + attention_entropy: float + slots: list[ConsensusSlot] + + @classmethod + def from_workspace( + cls, + w: torch.Tensor, + alpha: torch.Tensor | None, + modality: str, + labels: list[str] | None = None, + top_k: int = 3, + ) -> "ConsensusSnapshot": + workspace_norm = float(w.norm(dim=-1).mean().item()) + slots: list[ConsensusSlot] = [] + attention_entropy = 0.0 + confidence = 0.0 + + if alpha is not None and alpha.numel() > 0: + probs = alpha[0].detach().cpu() + top_vals, top_idx = probs.topk(min(top_k, probs.numel())) + slots = [ + ConsensusSlot( + index=int(idx.item()), + label=labels[int(idx.item())] if labels and int(idx.item()) < len(labels) + else f"expert_{idx.item()}", + weight=float(val.item()), + ) + for val, idx in zip(top_vals, top_idx) + ] + + probs_clamped = probs.clamp_min(1e-8) + attention_entropy = float((-(probs_clamped * probs_clamped.log()).sum()).item()) + max_entropy = math.log(max(2, probs.numel())) + confidence = float(max(0.0, 1.0 - attention_entropy / max_entropy)) + + return cls( + modality=modality, + workspace_norm=workspace_norm, + confidence=confidence, + attention_entropy=attention_entropy, + slots=slots, + ) + + def primary_slot(self) -> ConsensusSlot | None: + return self.slots[0] if self.slots else None + + def to_dict(self) -> dict: + return { + "modality": self.modality, + "workspace_norm": self.workspace_norm, + "confidence": self.confidence, + "attention_entropy": self.attention_entropy, + "slots": [ + {"index": slot.index, "label": slot.label, "weight": slot.weight} + for slot in self.slots + ], + } diff --git a/jspaceai/continual.py b/jspaceai/continual.py new file mode 100644 index 0000000..f64597c --- /dev/null +++ b/jspaceai/continual.py @@ -0,0 +1,104 @@ +""" +Continual learning utilities shared by interactive runtimes. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from .language_model import ( + EWCOptimizer, + ExperienceReplay, + ExpertPlasticity, +) + + +class OnlineLanguageLearner: + """ + Stateful online learner for chat/runtime usage. + + Unlike the old ad hoc SGD update, this keeps replay, EWC regularization, + and periodic consolidation alive across the whole session. + """ + + def __init__( + self, + model, + config, + tokenizer, + device: str = "cpu", + lr: float = 5e-3, + ewc_lambda: float = 0.05, + seq_len: int = 48, + replay_batch_size: int = 4, + replay_weight: float = 0.5, + consolidate_every: int = 20, + ): + self.model = model + self.config = config + self.tokenizer = tokenizer + self.device = device + self.seq_len = seq_len + self.replay_batch_size = replay_batch_size + self.replay_weight = replay_weight + self.consolidate_every = consolidate_every + + self.optimizer = EWCOptimizer(model, lr=lr, ewc_lambda=ewc_lambda) + self.replay_buffer = ExperienceReplay(capacity=500, seq_len=seq_len) + self.plasticity = ExpertPlasticity(num_experts=config.num_experts) + self.step_count = 0 + + def _prepare_sequence(self, text: str) -> torch.Tensor | None: + token_ids = self.tokenizer.encode(text) + if len(token_ids) < 4: + return None + if len(token_ids) < self.seq_len: + token_ids = token_ids * (self.seq_len // len(token_ids) + 1) + token_ids = token_ids[:self.seq_len] + return torch.tensor([token_ids], dtype=torch.long, device=self.device) + + def learn_text(self, text: str) -> dict | None: + token_seq = self._prepare_sequence(text) + if token_seq is None: + return None + + self.model.train() + logits, info = self.model(token_seq) + pred = logits[:, :-1] + target = token_seq[:, 1:] + loss = F.cross_entropy( + pred.reshape(-1, self.config.vocab_size), + target.reshape(-1), + ) + + replay_loss = torch.tensor(0.0, device=self.device) + replay_seq = self.replay_buffer.sample(self.replay_batch_size) + if replay_seq is not None: + replay_seq = replay_seq.to(self.device) + replay_logits, _ = self.model(replay_seq) + replay_pred = replay_logits[:, :-1] + replay_target = replay_seq[:, 1:] + replay_loss = F.cross_entropy( + replay_pred.reshape(-1, self.config.vocab_size), + replay_target.reshape(-1), + ) + + total_task_loss = loss + self.replay_weight * replay_loss + total_loss = self.optimizer.step(total_task_loss) + + self.plasticity.update(info["alpha"].detach(), token_seq.detach()) + self.replay_buffer.push(token_seq.detach().cpu()) + self.step_count += 1 + + if self.step_count % self.consolidate_every == 0: + self.optimizer.consolidate(token_seq.detach(), n_samples=10) + + self.model.eval() + return { + "loss": float(loss.item()), + "replay_loss": float(replay_loss.item()), + "total_loss": float(total_loss), + "w_norm_mean": float(info["w_norm"].mean().item()), + "usage": self.plasticity.usage.tolist(), + "step": self.step_count, + } diff --git a/jspaceai/embodied.py b/jspaceai/embodied.py index f47b827..eb3276c 100644 --- a/jspaceai/embodied.py +++ b/jspaceai/embodied.py @@ -20,22 +20,22 @@ workspace 是模态无关的"意图空间"。 from __future__ import annotations import torch -import torch.nn as nn -import torch.nn.functional as F import numpy as np import time -import json -from pathlib import Path -from typing import Optional -from dataclasses import dataclass, field -from collections import deque from pynput import mouse as pynput_mouse from pynput import keyboard as pynput_keyboard - -# ============================================================ -# 执行器层（对应手脚口） -# ============================================================ +from .consensus import ConsensusSnapshot +from .events import WorkspaceEvent +from .memory import Hippocampus, InMemoryVectorMemoryStore +from .policy import ( + ActionPolicy, + BasalGanglia, + Cerebellum, + CentralNervousSystem, + ReflexArc, + compose_action_params, +) class MouseActuator: """ @@ -172,248 +172,6 @@ class ScreenActuator: if self.enabled: self._cv2.destroyWindow(self.window_name) - -# ============================================================ -# 小脑：运动控制器（前向模型 + 逆模型） -# ============================================================ - -class Cerebellum(nn.Module): - """ - 小脑——运动协调与精细控制。 - - 前向模型：预测"如果执行动作 A，鼠标会到哪里" - 逆模型：给定"目标位置"，计算"需要什么动作" - - 人类小脑学习动作的精细映射，让动作平滑准确。 - 我们这里学习 workspace 意图 → 精确动作参数的映射。 - - 逆模型：workspace → 动作参数 - 前向模型：动作参数 → 预测结果（用于误差反馈学习） - """ - - def __init__(self, workspace_dim: int, action_dim: int = 5): - super().__init__() - # 逆模型：workspace → action - self.inverse_model = nn.Sequential( - nn.Linear(workspace_dim, 64), - nn.ReLU(), - nn.Linear(64, 32), - nn.ReLU(), - nn.Linear(32, action_dim), - nn.Tanh(), # 动作在 [-1, 1] - ) - # 前向模型：action + 当前状态 → 预测下一状态 - self.forward_model = nn.Sequential( - nn.Linear(action_dim + workspace_dim, 64), - nn.ReLU(), - nn.Linear(64, workspace_dim), - ) - self.action_dim = action_dim - - def compute_action(self, w: torch.Tensor) -> torch.Tensor: - """逆模型：从 workspace 意图计算动作""" - return self.inverse_model(w) - - def predict_next(self, w: torch.Tensor, action: torch.Tensor) -> torch.Tensor: - """前向模型：预测执行动作后的 workspace 状态""" - return self.forward_model(torch.cat([action, w], dim=-1)) - - def compute_motor_error(self, w: torch.Tensor, action: torch.Tensor, - w_actual_next: torch.Tensor) -> torch.Tensor: - """计算运动误差——用于小脑学习""" - w_pred = self.predict_next(w, action) - return F.mse_loss(w_pred, w_actual_next) - - -# ============================================================ -# 中枢神经：动作调度器 -# ============================================================ - -@dataclass -class ReflexArc: - """反射弧——不经过大脑的快速反应""" - trigger: str # 触发条件描述 - condition: callable # 检查函数 - action: callable # 执行函数 - priority: int = 0 # 优先级 - - -class CentralNervousSystem: - """ - 中枢神经系统——动作调度。 - - 功能： - 1. 反射弧：快速反应，不经过 workspace - - 如：突然大声音 → 退缩 - - 如：屏幕突然变暗 → 警觉 - 2. 决策门控：决定是否让 workspace 的意图执行 - - 高风险动作需要"确认" - - 习惯化动作直接执行 - 3. 动作序列：把复杂意图拆成动作序列 - - 如"点击按钮"→ 移动到位置 + 点击 - - 这是"自由意志"的工程对应——不是所有意图都执行， - 系统有一个门控机制决定哪些意图变成行动。 - """ - - def __init__(self): - self.reflexes: list[ReflexArc] = [] - self.action_history: deque = deque(maxlen=100) - self.inhibit_score: float = 0.0 # 抑制分数，高时阻止动作 - - def add_reflex(self, reflex: ReflexArc): - self.reflexes.append(reflex) - self.reflexes.sort(key=lambda r: -r.priority) - - def check_reflexes(self, sensory_state: dict) -> Optional[callable]: - """检查是否有反射触发""" - for reflex in self.reflexes: - try: - if reflex.condition(sensory_state): - return reflex.action - except Exception: - continue - return None - - def should_execute(self, action_strength: float, risk: float = 0.0) -> bool: - """决策门控：是否执行动作 - - Args: - action_strength: 动作强度（workspace 驱动） - risk: 风险评估（0-1） - - Returns: - 是否执行 - """ - # 抑制分数高时不执行 - threshold = 0.3 + risk * 0.5 + self.inhibit_score - return action_strength > threshold - - def record_action(self, action: np.ndarray, modality: str): - """记录执行的动作""" - self.action_history.append({ - 'time': time.time(), - 'action': action.tolist() if hasattr(action, 'tolist') else action, - 'modality': modality, - }) - - -# ============================================================ -# 海马体：外部情景记忆库 -# ============================================================ - -class Hippocampus: - """ - 海马体——情景记忆。 - - 存储历史 workspace 快照 + 时间戳 + 上下文。 - 当前 workspace 可以"回忆"相似的历史状态。 - - 人类的情景记忆："我记得昨天在那个房间里说了什么" - 对应：检索与当前 workspace 相似的历史 workspace。 - - 实现用简单的向量数据库（numpy + cosine similarity）。 - """ - - def __init__(self, capacity: int = 1000, workspace_dim: int = 64): - self.capacity = capacity - self.workspace_dim = workspace_dim - self.memories: deque = deque(maxlen=capacity) - - def store(self, w: np.ndarray, context: dict = None): - """存储一个 workspace 快照""" - self.memories.append({ - 'w': w.copy(), - 'context': context or {}, - 'timestamp': time.time(), - }) - - def recall(self, w_query: np.ndarray, top_k: int = 3) -> list[dict]: - """检索相似的历史记忆 - - Args: - w_query: 当前 workspace - top_k: 返回最相似的 k 个 - - Returns: - list of {w, context, timestamp, similarity} - """ - if not self.memories: - return [] - - # 计算相似度 - similarities = [] - for mem in self.memories: - sim = np.dot(w_query, mem['w']) / ( - np.linalg.norm(w_query) * np.linalg.norm(mem['w']) + 1e-8 - ) - similarities.append(sim) - - # 取 top-k - top_idx = np.argsort(similarities)[-top_k:][::-1] - results = [] - for idx in top_idx: - mem = self.memories[idx] - results.append({ - 'w': mem['w'], - 'context': mem['context'], - 'timestamp': mem['timestamp'], - 'similarity': similarities[idx], - }) - return results - - def size(self) -> int: - return len(self.memories) - - -# ============================================================ -# 基底神经节：动作价值学习 -# ============================================================ - -class BasalGanglia: - """ - 基底神经节——习惯学习与动作选择。 - - 学习"在什么状态下执行什么动作价值多少"。 - 高频执行的（workspace, action）对会"习惯化"——直接执行不经过思考。 - - 对应人类的习惯：开车的动作熟练后不需要思考， - 就是基底神经节接管了动作选择。 - """ - - def __init__(self, workspace_dim: int = 64, n_actions: int = 5, - learning_rate: float = 0.01): - self.workspace_dim = workspace_dim - self.n_actions = n_actions - self.lr = learning_rate - # Q-table 的近似：用线性函数 Q(s, a) = w_a · s - self.action_weights = np.zeros((n_actions, workspace_dim)) - # 习惯化计数 - self.habit_counts = np.zeros(n_actions) - - def compute_values(self, w: np.ndarray) -> np.ndarray: - """计算各动作的价值 Q(s, a)""" - return self.action_weights @ w # (n_actions,) - - def select_action(self, w: np.ndarray, exploration: float = 0.1) -> int: - """选择动作（ε-greedy）""" - values = self.compute_values(w) - if np.random.random() < exploration: - return np.random.randint(self.n_actions) - return np.argmax(values) - - def update(self, w: np.ndarray, action: int, reward: float): - """更新动作价值（TD learning 简化版）""" - values = self.compute_values(w) - td_error = reward - values[action] - self.action_weights[action] += self.lr * td_error * w - self.habit_counts[action] += 1 - - def is_habitual(self, action: int, threshold: int = 10) -> bool: - """判断动作是否已习惯化""" - return self.habit_counts[action] >= threshold - - # ============================================================ # 完整的具身 Agent # ============================================================ @@ -422,18 +180,14 @@ class EmbodiedAgent: """ 完整的具身 Agent——感知-思考-行动闭环。 - 结构（对应人类神经系统）： + 结构： 感知层（眼耳皮肤）→ FullSensoryStream ↓ 编码 大脑皮层（思考）→ MultimodalJSpaceModel ↓ workspace w 海马体（记忆）→ Hippocampus 存储和检索 ↓ - 基底神经节（动作选择）→ BasalGanglia 选动作 - ↓ - 小脑（运动控制）→ Cerebellum 精细化动作 - ↓ - 中枢神经（门控）→ CentralNervousSystem 决定是否执行 + ActionPolicy（价值 + 电机参数 + 门控） ↓ 执行器（手足口）→ MouseActuator + KeyboardActuator + AudioActuator @@ -441,11 +195,9 @@ class EmbodiedAgent: 1. 感知：从五感获取输入 2. 思考：workspace 更新 3. 回忆：海马体检索相关记忆 - 4. 决策：基底神经节选动作 - 5. 精化：小脑计算动作参数 - 6. 门控：中枢神经决定执行 + 4. 决策：ActionPolicy 选动作并门控 7. 行动：执行器执行 - 8. 学习：更新基底神经节、小脑、海马体 + 8. 学习：更新策略与海马体 """ def __init__(self, model, device: str = 'cpu', @@ -474,25 +226,27 @@ class EmbodiedAgent: self.audio_actuator = AudioActuator(enabled=enable_audio_output) self.screen_actuator = ScreenActuator(enabled=enable_screen_output) - # 神经系统 - self.cerebellum = Cerebellum( + # 策略层 + self.policy = ActionPolicy( workspace_dim=self.config.workspace_dim, - action_dim=5, # (dx, dy, click_l, click_r, scroll) - ).to(device) - self.cns = CentralNervousSystem() - self.hippocampus = Hippocampus( + action_dim=5, + n_actions=5, + base_threshold=risk_threshold, + device=device, + ) + self.cerebellum = self.policy.motor_controller + self.cns = self.policy.gate + self.hippocampus = InMemoryVectorMemoryStore( workspace_dim=self.config.workspace_dim, ) if enable_memory else None - self.basal_ganglia = BasalGanglia( - workspace_dim=self.config.workspace_dim, - n_actions=5, - ) + self.basal_ganglia = self.policy.value_model # 内部状态 self.state = model.init_state(1, torch.device(device)) self.step_count = 0 self.running = False self.risk_threshold = risk_threshold + self.last_consensus: ConsensusSnapshot | None = None # 设置反射弧 self._setup_reflexes() @@ -509,7 +263,7 @@ class EmbodiedAgent: def loud_noise_response(): self.cns.inhibit_score = 0.5 - self.cns.add_reflex(ReflexArc( + self.policy.add_reflex(ReflexArc( trigger='loud_noise', condition=loud_noise_condition, action=loud_noise_response, @@ -572,14 +326,24 @@ class EmbodiedAgent: if x.dim() == 1: x = x.unsqueeze(0) if x.dim() == 3: x = x[:, -1, :] if x.shape[0] != 1: x = x[-1:] - self.state, _ = self.model.step(self.state, x) + self.state, alpha, _ = self.model.step(self.state, x) + self.last_consensus = ConsensusSnapshot.from_workspace( + self.state['w'], alpha, modality, self.model.expert_modality + ) return self.state['w'], modality def remember(self, w: torch.Tensor, context: dict): - """存储到海马体""" + """存储 workspace event 到记忆层""" if self.hippocampus: - self.hippocampus.store(w[0].cpu().numpy(), context) + event = WorkspaceEvent.from_tensor( + w, + modality=str(context.get('modality', 'unknown')), + step=int(context.get('step', self.step_count)), + consensus=context.get('consensus', self.last_consensus), + context=context, + ) + self.hippocampus.put(event) def recall_memories(self, w: torch.Tensor, top_k: int = 3) -> list: """从海马体回忆""" @@ -595,26 +359,11 @@ class EmbodiedAgent: 3. 中枢神经门控 4. 执行器执行 """ - w_np = w[0].cpu().numpy() - - # 1. 基底神经节：选动作类别 - action_idx = self.basal_ganglia.select_action(w_np, exploration=0.2) - - # 2. 小脑：计算精确动作参数 - with torch.no_grad(): - action_params = self.cerebellum.compute_action(w)[0].cpu().numpy() - - # 3. 中枢神经：门控 - action_strength = np.abs(action_params).max() - risk = 0.0 - # 鼠标点击风险较高 - if action_params[2] > 0.5 or action_params[3] > 0.5: - risk = 0.5 - - execute = self.cns.should_execute(action_strength, risk) + decision = self.policy.decide(w) + action_params = decision.action_params + execute = decision.should_execute # 4. 执行 - action_taken = None if execute: # 鼠标动作 self.mouse_actuator.execute(action_params) @@ -630,24 +379,13 @@ class EmbodiedAgent: img_out = ((img_out + 1) / 2 * 255).clip(0, 255).astype(np.uint8) self.screen_actuator.show_image(img_out) - action_taken = action_params.tolist() - self.cns.record_action(action_params, modality) + self.policy.record_action(action_params, modality) - return { - 'action_idx': action_idx, - 'action_params': action_params.tolist(), - 'executed': execute, - 'action_strength': float(action_strength), - 'risk': float(risk), - } + return decision.to_dict(executed=execute) - def learn(self, w: torch.Tensor, action_params: np.ndarray, reward: float = 0.0): + def learn(self, w: torch.Tensor, action_idx: int, reward: float = 0.0): """学习——更新基底神经节和小脑""" - w_np = w[0].cpu().numpy() - - # 基底神经节：更新动作价值 - action_idx = self.basal_ganglia.select_action(w_np, exploration=0.0) - self.basal_ganglia.update(w_np, action_idx, reward) + self.policy.learn(w, action_idx, reward) def step_once(self) -> dict: """执行一步完整的感知-思考-行动循环""" @@ -663,14 +401,14 @@ class EmbodiedAgent: w, modality = self.think(sensory_data) # 4. 回忆 - memories = self.recall_memories(w) + self.recall_memories(w) # 5. 决策和行动 action_info = self.decide_and_act(w, modality) # 6. 学习（自监督：预测误差作为 reward） reward = -action_info['risk'] # 简化：风险越低 reward 越高 - self.learn(w, np.array(action_info['action_params']), reward) + self.learn(w, action_info['action_idx'], reward) # 7. 记忆存储 self.remember(w, { @@ -686,6 +424,7 @@ class EmbodiedAgent: 'modality': modality, 'w_norm': w.norm().item(), 'action': action_info, + 'consensus': self.last_consensus.to_dict() if self.last_consensus else None, 'memories_count': self.hippocampus.size() if self.hippocampus else 0, } @@ -714,7 +453,7 @@ class EmbodiedAgent: elif info['step'] % 5 == 0: print(f" step {info['step']:3d} | mod {info['modality']:8s} | " f"||w|| {info['w_norm']:.3f} | " - f"action {info['action']['action_idx']} | " + f"action {info['action']['action_name']} | " f"executed {info['action']['executed']}") time.sleep(interval) diff --git a/jspaceai/events.py b/jspaceai/events.py new file mode 100644 index 0000000..3a275a1 --- /dev/null +++ b/jspaceai/events.py @@ -0,0 +1,128 @@ +""" +Serializable event types for the workspace-centered runtime. + +These data packets are intentionally small and plain. They give local modules, +future worker processes, and external transports the same language for passing +workspace state around without sharing live Python objects. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +import time +from typing import Any + +import numpy as np +import torch + +from .consensus import ConsensusSnapshot + + +def _workspace_vector(w: torch.Tensor | np.ndarray | list[float]) -> list[float]: + if isinstance(w, torch.Tensor): + data = w.detach().cpu() + if data.dim() > 1: + data = data[0] + return data.reshape(-1).tolist() + if isinstance(w, np.ndarray): + data = w + if data.ndim > 1: + data = data[0] + return data.reshape(-1).astype(float).tolist() + return [float(v) for v in w] + + +def _consensus_dict(consensus: ConsensusSnapshot | dict | None) -> dict | None: + if consensus is None: + return None + if isinstance(consensus, ConsensusSnapshot): + return consensus.to_dict() + return consensus + + +@dataclass +class WorkspaceEvent: + """A transport-friendly snapshot of workspace state.""" + + step: int + modality: str + workspace: list[float] + consensus: dict | None = None + context: dict[str, Any] = field(default_factory=dict) + timestamp: float = field(default_factory=time.time) + kind: str = "workspace" + + @classmethod + def from_tensor( + cls, + w: torch.Tensor | np.ndarray | list[float], + modality: str, + step: int = 0, + consensus: ConsensusSnapshot | dict | None = None, + context: dict | None = None, + ) -> "WorkspaceEvent": + return cls( + step=step, + modality=modality, + workspace=_workspace_vector(w), + consensus=_consensus_dict(consensus), + context=context or {}, + ) + + def workspace_array(self) -> np.ndarray: + return np.asarray(self.workspace, dtype=np.float32) + + def to_dict(self) -> dict: + return { + "kind": self.kind, + "step": self.step, + "modality": self.modality, + "workspace": list(self.workspace), + "consensus": self.consensus, + "context": self.context, + "timestamp": self.timestamp, + } + + +@dataclass +class ActionEvent: + """A transport-friendly action decision/result.""" + + step: int + action_idx: int + action_name: str + action_params: list[float] + executed: bool + risk: float = 0.0 + context: dict[str, Any] = field(default_factory=dict) + timestamp: float = field(default_factory=time.time) + kind: str = "action" + + @classmethod + def from_action_info( + cls, + action_info: dict, + step: int = 0, + context: dict | None = None, + ) -> "ActionEvent": + return cls( + step=step, + action_idx=int(action_info.get("action_idx", 0)), + action_name=str(action_info.get("action_name", "observe")), + action_params=[float(v) for v in action_info.get("action_params", [])], + executed=bool(action_info.get("executed", False)), + risk=float(action_info.get("risk", 0.0)), + context=context or {}, + ) + + def to_dict(self) -> dict: + return { + "kind": self.kind, + "step": self.step, + "action_idx": self.action_idx, + "action_name": self.action_name, + "action_params": list(self.action_params), + "executed": self.executed, + "risk": self.risk, + "context": self.context, + "timestamp": self.timestamp, + } diff --git a/jspaceai/memory.py b/jspaceai/memory.py new file mode 100644 index 0000000..72e7d46 --- /dev/null +++ b/jspaceai/memory.py @@ -0,0 +1,115 @@ +""" +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 diff --git a/jspaceai/multimodal.py b/jspaceai/multimodal.py index a36a6b8..5fe88bf 100644 --- a/jspaceai/multimodal.py +++ b/jspaceai/multimodal.py @@ -333,7 +333,7 @@ class MultimodalJSpaceModel(nn.Module): } def step(self, state: dict, x: torch.Tensor, - record_trajectory: bool = False) -> tuple[dict, list]: + record_trajectory: bool = False) -> tuple[dict, torch.Tensor, list]: """ 单步前向 @@ -343,7 +343,7 @@ class MultimodalJSpaceModel(nn.Module): record_trajectory: 是否记录 w 轨迹 Returns: - new_state, w_trajectory (list) + new_state, alpha, w_trajectory (list) """ w = state['w'] ms = state['m'] @@ -372,7 +372,7 @@ class MultimodalJSpaceModel(nn.Module): w_trajectory.append(w.detach()) new_state = {'w': w, 'm': ms} - return new_state, w_trajectory + return new_state, alpha, w_trajectory def encode_modality(self, modality: str, data: torch.Tensor) -> torch.Tensor: """编码任意模态到 input_dim @@ -437,14 +437,15 @@ class MultimodalJSpaceModel(nn.Module): # 处理序列或单步 w_traj_all = [] + alpha_last = None if x.dim() == 2: # 单步 (batch, input_dim) - state, w_traj = self.step(state, x, record_trajectory) + state, alpha_last, w_traj = self.step(state, x, record_trajectory) if record_trajectory: w_traj_all = w_traj w = state['w'] else: # 序列 (batch, T, input_dim) for t in range(x.shape[1]): - state, w_traj = self.step(state, x[:, t], record_trajectory) + state, alpha_last, w_traj = self.step(state, x[:, t], record_trajectory) if record_trajectory: w_traj_all.append(w_traj) w = state['w'] @@ -460,6 +461,8 @@ class MultimodalJSpaceModel(nn.Module): info = { 'w_norm': w.norm(dim=-1), } + if alpha_last is not None: + info['alpha'] = alpha_last if record_trajectory and w_traj_all: info['w_trajectory'] = w_traj_all diff --git a/jspaceai/policy.py b/jspaceai/policy.py new file mode 100644 index 0000000..c3ed5a2 --- /dev/null +++ b/jspaceai/policy.py @@ -0,0 +1,254 @@ +""" +Action policy utilities centered on the workspace state. + +This module keeps action selection, motor parameterization, and gating together +so the embodied runtime can stay focused on sensing and effectors. +""" +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +import time +from typing import Callable, Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +ACTION_LABELS = { + 0: "observe", + 1: "mouse_move", + 2: "left_click", + 3: "right_click", + 4: "scroll", +} + + +def compose_action_params(action_idx: int, raw_params: np.ndarray) -> np.ndarray: + """Map a discrete action choice onto concrete actuator parameters.""" + action = np.zeros(5, dtype=np.float32) + if action_idx == 1: + action[0] = float(raw_params[0]) + action[1] = float(raw_params[1]) + elif action_idx == 2: + action[2] = 1.0 if raw_params[2] >= 0 else 0.0 + elif action_idx == 3: + action[3] = 1.0 if raw_params[3] >= 0 else 0.0 + elif action_idx == 4: + action[4] = float(raw_params[4]) + return action + + +@dataclass +class ReflexRule: + """Fast path rule that can inhibit or redirect behavior before planning.""" + + trigger: str + condition: Callable + action: Callable + priority: int = 0 + + +class MotorController(nn.Module): + """ + Refines workspace intent into continuous motor parameters. + """ + + def __init__(self, workspace_dim: int, action_dim: int = 5): + super().__init__() + self.inverse_model = nn.Sequential( + nn.Linear(workspace_dim, 64), + nn.ReLU(), + nn.Linear(64, 32), + nn.ReLU(), + nn.Linear(32, action_dim), + nn.Tanh(), + ) + self.forward_model = nn.Sequential( + nn.Linear(action_dim + workspace_dim, 64), + nn.ReLU(), + nn.Linear(64, workspace_dim), + ) + self.action_dim = action_dim + + def compute_action(self, w: torch.Tensor) -> torch.Tensor: + return self.inverse_model(w) + + def predict_next(self, w: torch.Tensor, action: torch.Tensor) -> torch.Tensor: + return self.forward_model(torch.cat([action, w], dim=-1)) + + def compute_motor_error( + self, + w: torch.Tensor, + action: torch.Tensor, + w_actual_next: torch.Tensor, + ) -> torch.Tensor: + w_pred = self.predict_next(w, action) + return F.mse_loss(w_pred, w_actual_next) + + +class ActionGate: + """ + Lightweight action gate with reflex handling and execution thresholds. + """ + + def __init__(self, base_threshold: float = 0.3): + self.base_threshold = base_threshold + self.reflexes: list[ReflexRule] = [] + self.action_history: deque = deque(maxlen=100) + self.inhibit_score: float = 0.0 + + def add_reflex(self, reflex: ReflexRule): + self.reflexes.append(reflex) + self.reflexes.sort(key=lambda r: -r.priority) + + def check_reflexes(self, sensory_state: dict) -> Optional[Callable]: + for reflex in self.reflexes: + try: + if reflex.condition(sensory_state): + return reflex.action + except Exception: + continue + return None + + def should_execute(self, action_strength: float, risk: float = 0.0) -> bool: + threshold = self.base_threshold + risk * 0.5 + self.inhibit_score + return action_strength > threshold + + def record_action(self, action: np.ndarray, modality: str): + self.action_history.append({ + "time": time.time(), + "action": action.tolist() if hasattr(action, "tolist") else action, + "modality": modality, + }) + + +class ActionValueModel: + """ + Linear value function over workspace state for discrete action choice. + """ + + def __init__(self, workspace_dim: int = 64, n_actions: int = 5, learning_rate: float = 0.01): + self.workspace_dim = workspace_dim + self.n_actions = n_actions + self.lr = learning_rate + self.action_weights = np.zeros((n_actions, workspace_dim)) + self.habit_counts = np.zeros(n_actions) + self.action_labels = [ACTION_LABELS.get(i, f"action_{i}") for i in range(n_actions)] + + def compute_values(self, w: np.ndarray) -> np.ndarray: + return self.action_weights @ w + + def select_action(self, w: np.ndarray, exploration: float = 0.1) -> int: + values = self.compute_values(w) + if np.random.random() < exploration: + return np.random.randint(self.n_actions) + return int(np.argmax(values)) + + def update(self, w: np.ndarray, action: int, reward: float): + values = self.compute_values(w) + td_error = reward - values[action] + self.action_weights[action] += self.lr * td_error * w + self.habit_counts[action] += 1 + + def is_habitual(self, action: int, threshold: int = 10) -> bool: + return self.habit_counts[action] >= threshold + + +@dataclass +class ActionDecision: + action_idx: int + action_name: str + action_params: np.ndarray + raw_action_params: np.ndarray + action_strength: float + risk: float + should_execute: bool + + def to_dict(self, executed: bool) -> dict: + return { + "action_idx": self.action_idx, + "action_name": self.action_name, + "action_params": self.action_params.tolist(), + "raw_action_params": self.raw_action_params.tolist(), + "executed": executed, + "action_strength": self.action_strength, + "risk": self.risk, + } + + +class ActionPolicy: + """ + Unified policy around discrete action choice, motor refinement, and gating. + """ + + def __init__( + self, + workspace_dim: int, + action_dim: int = 5, + n_actions: int = 5, + learning_rate: float = 0.01, + exploration: float = 0.2, + base_threshold: float = 0.3, + device: str = "cpu", + ): + self.exploration = exploration + self.value_model = ActionValueModel( + workspace_dim=workspace_dim, + n_actions=n_actions, + learning_rate=learning_rate, + ) + self.motor_controller = MotorController( + workspace_dim=workspace_dim, + action_dim=action_dim, + ).to(device) + self.gate = ActionGate(base_threshold=base_threshold) + + @property + def action_labels(self) -> list[str]: + return self.value_model.action_labels + + def add_reflex(self, reflex: ReflexRule): + self.gate.add_reflex(reflex) + + def check_reflexes(self, sensory_state: dict) -> Optional[Callable]: + return self.gate.check_reflexes(sensory_state) + + def record_action(self, action: np.ndarray, modality: str): + self.gate.record_action(action, modality) + + def decide(self, w: torch.Tensor) -> ActionDecision: + w_np = w[0].detach().cpu().numpy() + action_idx = self.value_model.select_action(w_np, exploration=self.exploration) + action_name = self.value_model.action_labels[action_idx] + + with torch.no_grad(): + raw_action = self.motor_controller.compute_action(w)[0].detach().cpu().numpy() + action_params = compose_action_params(action_idx, raw_action) + + action_strength = float(np.abs(action_params).max()) + risk = 0.5 if action_params[2] > 0.5 or action_params[3] > 0.5 else 0.0 + should_execute = action_idx != 0 and self.gate.should_execute(action_strength, risk) + + return ActionDecision( + action_idx=action_idx, + action_name=action_name, + action_params=action_params, + raw_action_params=raw_action, + action_strength=action_strength, + risk=float(risk), + should_execute=should_execute, + ) + + def learn(self, w: torch.Tensor, action_idx: int, reward: float = 0.0): + w_np = w[0].detach().cpu().numpy() + self.value_model.update(w_np, action_idx, reward) + + +# Backward-compatible aliases for the old brain-region names. +ReflexArc = ReflexRule +CentralNervousSystem = ActionGate +BasalGanglia = ActionValueModel +Cerebellum = MotorController diff --git a/jspaceai/runtime.py b/jspaceai/runtime.py new file mode 100644 index 0000000..b0abcb5 --- /dev/null +++ b/jspaceai/runtime.py @@ -0,0 +1,454 @@ +""" +Unified runtime loop around a workspace-centered agent. + +This keeps the persistent/autonomous scaffolding separate from embodiment so we +can reuse one loop for live sensing, future task runners, and evaluation. +""" +from __future__ import annotations + +from collections import deque +import json +from pathlib import Path +import time +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .events import ActionEvent, WorkspaceEvent + + +class CuriosityDrive(nn.Module): + """Predictive world model that turns surprise into intrinsic reward.""" + + def __init__(self, workspace_dim: int, action_dim: int = 5, hidden_dim: int = 64): + super().__init__() + self.workspace_dim = workspace_dim + self.world_model = nn.Sequential( + nn.Linear(workspace_dim + action_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, workspace_dim), + ) + self.state_history: deque = deque(maxlen=500) + self.prediction_error_ema = 0.1 + + def predict_next(self, w, action): + return self.world_model(torch.cat([w, action], dim=-1)) + + def compute_curiosity(self, w_current, action, w_next): + with torch.no_grad(): + w_pred = self.predict_next(w_current, action) + pred_error = F.mse_loss(w_pred, w_next).item() + w_np = w_current[0].cpu().numpy() + novelty = self._compute_novelty(w_np) + progress = max(0, pred_error - self.prediction_error_ema * 0.9) + self.prediction_error_ema = 0.95 * self.prediction_error_ema + 0.05 * pred_error + curiosity = progress + 0.3 * novelty + self.state_history.append(w_np.copy()) + return curiosity + + def _compute_novelty(self, w): + if len(self.state_history) < 5: + return 1.0 + history = list(self.state_history)[-100:] + distances = [np.linalg.norm(w - h) for h in history] + return float(min(1.0, min(distances) / 2.0)) + + def train_world_model(self, w_current, action, w_next): + w_pred = self.predict_next(w_current, action.detach()) + return F.mse_loss(w_pred, w_next.detach()) + + +class RuntimeStateStore: + """Minimal persistence for workspace state and runtime-side learning state.""" + + def __init__(self, save_dir: Path): + self.save_dir = Path(save_dir) + self.save_dir.mkdir(parents=True, exist_ok=True) + self.state_file = self.save_dir / "mind_state.json" + self.tensors_file = self.save_dir / "mind_tensors.npz" + + def save(self, state: dict): + tensors = {} + if "w" in state: + tensors["w"] = state["w"].cpu().numpy() + if "m" in state: + for i, m in enumerate(state["m"]): + if m is not None: + tensors[f"m_{i}"] = m.cpu().numpy() + if "action_value_weights" in state and state["action_value_weights"] is not None: + tensors["action_value_weights"] = state["action_value_weights"] + if "curiosity_history" in state: + tensors["curiosity_history"] = np.array(state["curiosity_history"]) + if tensors: + np.savez(self.tensors_file, **tensors) + + json_state = { + "step_count": state.get("step_count", 0), + "total_runtime": state.get("total_runtime", 0.0), + "self_model": state.get("self_model", {}), + "saved_at": time.time(), + } + self.state_file.write_text(json.dumps(json_state, indent=2, default=str)) + + def load(self) -> Optional[dict]: + if not self.state_file.exists(): + return None + result = {} + if self.tensors_file.exists(): + data = np.load(self.tensors_file, allow_pickle=True) + if "w" in data: + result["w"] = torch.tensor(data["w"]) + ms = {} + for key in data.files: + if key.startswith("m_"): + idx = int(key.split("_")[1]) + ms[idx] = torch.tensor(data[key]) + if ms: + result["m"] = [ms[i] for i in sorted(ms.keys())] + if "action_value_weights" in data: + result["action_value_weights"] = data["action_value_weights"] + elif "bg_weights" in data: + result["action_value_weights"] = data["bg_weights"] + if "curiosity_history" in data: + result["curiosity_history"] = data["curiosity_history"].tolist() + json_state = json.loads(self.state_file.read_text()) + result.update(json_state) + return result + + +class SelfModel: + """Tracks which capability domains are currently working well.""" + + def __init__(self, capabilities=None): + if capabilities is None: + capabilities = [ + "visual", + "audio", + "text", + "motor_mouse", + "motor_keyboard", + "memory", + "prediction", + ] + self.capabilities = capabilities + self.confidence = {c: 0.0 for c in capabilities} + self.attempts = {c: 0 for c in capabilities} + self.recent_results = {c: deque(maxlen=20) for c in capabilities} + + def record_attempt(self, capability, success): + if capability not in self.confidence: + return + self.attempts[capability] += 1 + self.recent_results[capability].append(success) + recent = list(self.recent_results[capability]) + if recent: + weights = np.linspace(0.5, 1.0, len(recent)) + self.confidence[capability] = float(np.average(recent, weights=weights)) + + def get_weakness(self): + return min(self.confidence, key=self.confidence.get) + + def get_strength(self): + return max(self.confidence, key=self.confidence.get) + + def knows(self, capability, threshold=0.5): + return self.confidence.get(capability, 0.0) > threshold + + def summary(self): + return { + "capabilities": dict(self.confidence), + "attempts": dict(self.attempts), + "strength": self.get_strength(), + "weakness": self.get_weakness(), + } + + +class MetaLearner: + """Tiny adaptive layer for choosing how aggressively each domain learns.""" + + def __init__(self, capabilities=None): + if capabilities is None: + capabilities = [ + "visual", + "audio", + "text", + "motor_mouse", + "motor_keyboard", + "memory", + "prediction", + ] + self.learning_rates = {c: 1e-3 for c in capabilities} + self.loss_history = {c: deque(maxlen=20) for c in capabilities} + self.strategy_scores = { + "predict_next": 0.5, + "replay": 0.5, + "explore": 0.5, + "imitate": 0.5, + } + + def get_lr(self, capability): + return self.learning_rates.get(capability, 1e-3) + + def record_loss(self, capability, loss): + if capability not in self.loss_history: + return + self.loss_history[capability].append(loss) + history = list(self.loss_history[capability]) + if len(history) < 5: + return + recent_avg = np.mean(history[-5:]) + old_avg = np.mean(history[-10:-5]) if len(history) >= 10 else recent_avg + improvement = (old_avg - recent_avg) / max(old_avg, 1e-8) + lr = self.learning_rates[capability] + if improvement > 0.05: + lr *= 1.1 + elif improvement < 0.01: + lr *= 0.9 + self.learning_rates[capability] = max(1e-5, min(1e-2, lr)) + + def best_strategy(self): + return max(self.strategy_scores, key=self.strategy_scores.get) + + def reward_strategy(self, strategy, reward): + if strategy in self.strategy_scores: + self.strategy_scores[strategy] = 0.9 * self.strategy_scores[strategy] + 0.1 * reward + + +class WorkspaceRuntime: + """ + Canonical workspace loop: + observe -> think -> act -> predict -> learn -> persist + """ + + def __init__(self, agent, save_dir="outputs/mind", device="cpu"): + self.agent = agent + self.device = device + self.config = agent.config + + self.curiosity = CuriosityDrive( + workspace_dim=self.config.workspace_dim, + action_dim=5, + ).to(device) + self.persistence = RuntimeStateStore(Path(save_dir)) + self.self_model = SelfModel() + self.meta_learner = MetaLearner() + + self.step_count = 0 + self.total_runtime = 0.0 + self.start_time = time.time() + self.running = False + self.curiosity_history = deque(maxlen=1000) + self.event_log = deque(maxlen=1000) + self.last_workspace_event: WorkspaceEvent | None = None + self.last_action_event: ActionEvent | None = None + self.curiosity_optimizer = torch.optim.Adam( + self.curiosity.parameters(), + lr=1e-3, + ) + + self._load_state() + + def _load_state(self): + state = self.persistence.load() + if state is None: + print(" [runtime] 全新启动") + return + print(f" [runtime] 恢复状态: step={state.get('step_count', 0)}") + if "w" in state: + self.agent.state["w"] = state["w"].to(self.device) + if "m" in state: + for i, m in enumerate(state["m"]): + if m is not None and i < len(self.agent.state["m"]): + self.agent.state["m"][i] = m.to(self.device) + weights = state.get("action_value_weights") + if weights is not None: + if hasattr(self.agent, "policy"): + self.agent.policy.value_model.action_weights = weights + elif hasattr(self.agent, "basal_ganglia"): + self.agent.basal_ganglia.action_weights = weights + self.step_count = state.get("step_count", 0) + self.total_runtime = state.get("total_runtime", 0.0) + + def save_state(self): + action_value_weights = None + if hasattr(self.agent, "policy"): + action_value_weights = self.agent.policy.value_model.action_weights + elif hasattr(self.agent, "basal_ganglia"): + action_value_weights = self.agent.basal_ganglia.action_weights + + state = { + "w": self.agent.state["w"], + "m": self.agent.state["m"], + "action_value_weights": action_value_weights, + "curiosity_history": list(self.curiosity.state_history), + "step_count": self.step_count, + "total_runtime": self.total_runtime + (time.time() - self.start_time), + "self_model": {"confidence": self.self_model.confidence}, + } + self.persistence.save(state) + + def step(self) -> dict: + sensory_data = self.agent.perceive() + w_before = self.agent.state["w"].clone() + weakness = self.self_model.get_weakness() + strength = self.self_model.get_strength() + + w_after, modality = self.agent.think(sensory_data) + action_info = self.agent.decide_and_act(w_after, modality) + consensus = getattr(self.agent, "last_consensus", None) + event_step = self.step_count + 1 + workspace_event = WorkspaceEvent.from_tensor( + w_after, + modality=modality, + step=event_step, + consensus=consensus, + ) + action_event = ActionEvent.from_action_info( + action_info, + step=event_step, + context={"modality": modality}, + ) + self.last_workspace_event = workspace_event + self.last_action_event = action_event + self.event_log.extend([workspace_event, action_event]) + + action_tensor = torch.tensor( + action_info["action_params"], + dtype=torch.float32, + ).unsqueeze(0).to(self.device) + curiosity_reward = self.curiosity.compute_curiosity(w_before, action_tensor, w_after) + self.curiosity_history.append(curiosity_reward) + + world_loss = self.curiosity.train_world_model(w_before, action_tensor, w_after) + self.curiosity_optimizer.zero_grad() + world_loss.backward() + self.curiosity_optimizer.step() + + w_stability = 1.0 - min(1.0, abs(w_after.norm().item() - w_before.norm().item())) + success = ( + 0.3 * float(action_info["executed"]) + + 0.4 * min(1.0, curiosity_reward) + + 0.3 * w_stability + ) + + cap_map = { + "image": "visual", + "screen": "visual", + "audio": "audio", + "text": "text", + "keyboard": "text", + "mouse": "motor_mouse", + "idle": "prediction", + } + cap = cap_map.get(modality, "prediction") + self.self_model.record_attempt(cap, success) + self.self_model.record_attempt("prediction", 1.0 - min(1.0, world_loss.item())) + + self.meta_learner.record_loss(cap, world_loss.item()) + if curiosity_reward > 0.3: + self.meta_learner.reward_strategy("explore", curiosity_reward) + + workspace_event.context.update({ + "curiosity": curiosity_reward, + "success": success, + "action": action_event.to_dict(), + }) + self.agent.remember(w_after, { + "modality": modality, + "curiosity": curiosity_reward, + "success": success, + "step": event_step, + "consensus": workspace_event.consensus, + "action_event": action_event.to_dict(), + }) + self.agent.learn(w_after, action_info["action_idx"], reward=curiosity_reward) + + self.step_count += 1 + focus = consensus.primary_slot().label if consensus and consensus.primary_slot() else "none" + confidence = consensus.confidence if consensus else 0.0 + memory = getattr(self.agent, "hippocampus", None) + memory_count = memory.size() if memory else 0 + + return { + "step": self.step_count, + "modality": modality, + "w_norm": w_after.norm().item(), + "curiosity": curiosity_reward, + "world_loss": world_loss.item(), + "success": success, + "weakness": weakness, + "strength": strength, + "self_confidence": dict(self.self_model.confidence), + "best_strategy": self.meta_learner.best_strategy(), + "consensus_confidence": confidence, + "consensus_focus": focus, + "memory_count": memory_count, + "workspace_event": workspace_event.to_dict(), + "action_event": action_event.to_dict(), + } + + def step_once(self) -> dict: + return self.step() + + def run(self, n_steps=100, interval=0.2, save_every=50, on_step=None): + self.running = True + self.start_time = time.time() + self.agent.senses.start() + print(f"\nworkspace runtime 启动 | 总步数: {self.step_count} | 保存间隔: {save_every}") + print("=" * 60) + + log = [] + try: + for _ in range(n_steps): + if not self.running: + break + info = self.step() + log.append(info) + if on_step: + on_step(info) + elif info["step"] % 10 == 0: + print( + f" step {info['step']:4d} | mod {info['modality']:8s} | " + f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | " + f"success {info['success']:.2f} | weak={info['weakness']} | " + f"mem {info['memory_count']}" + ) + if info["step"] % save_every == 0: + self.save_state() + time.sleep(interval) + except KeyboardInterrupt: + print("\n用户中断") + finally: + self.running = False + self.agent.senses.stop() + if hasattr(self.agent, "audio_actuator"): + self.agent.audio_actuator.stop() + self.save_state() + self.total_runtime += time.time() - self.start_time + return log + + def introspect(self) -> str: + sm = self.self_model.summary() + avg_curio = np.mean(list(self.curiosity_history)) if self.curiosity_history else 0 + report = f"=== Workspace Runtime ===\n步数: {self.step_count}\n运行: {self.total_runtime:.0f}s\n" + report += f"记忆: {self.agent.hippocampus.size() if self.agent.hippocampus else 0}\n" + report += f"平均好奇心: {avg_curio:.3f}\n\n自我认知:\n" + for cap, conf in sm["capabilities"].items(): + bar = "█" * int(conf * 20) + report += f" {cap:15s}: {conf:.2f} {bar}\n" + report += f"\n最强: {sm['strength']}\n最弱: {sm['weakness']}\n" + report += f"最佳策略: {self.meta_learner.best_strategy()}\n" + return report + + +# Backward-compatible aliases. +PersistentState = RuntimeStateStore + + +class AutonomousMind(WorkspaceRuntime): + """Backward-compatible name for the unified runtime.""" diff --git a/main.py b/main.py index 2b6e1c9..83a1ca8 100644 --- a/main.py +++ b/main.py @@ -28,7 +28,7 @@ import time from jspaceai import ( MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent, - AutonomousMind, PLATFORM, + WorkspaceRuntime, PLATFORM, get_screen_size, print_permission_guide, check_camera_permission, check_microphone_permission, check_input_monitoring_permission, @@ -85,7 +85,7 @@ def test_subsystems(): for _ in range(5): info = agent.step_once() print(f" step {info['step']:2d} | mod {info['modality']:8s} | " - f"||w|| {info['w_norm']:.3f} | action {info['action']['action_idx']} | " + f"||w|| {info['w_norm']:.3f} | action {info['action']['action_name']} | " f"executed {info['action']['executed']} | mem {info['memories_count']}") time.sleep(0.5) agent.senses.stop() @@ -136,10 +136,10 @@ def live(n_steps: int, device: str, safe_mode: bool = False, unsafe: bool = Fals enable_screen_output=not safe_mode, risk_threshold=0.5 if safe_mode else 0.3, ) - mind = AutonomousMind(agent, save_dir='outputs/mind', device=device) + runtime = WorkspaceRuntime(agent, save_dir='outputs/mind', device=device) print("\n" + "=" * 60) - print("自主心智 - 全感官具身循环") + print("Workspace Runtime - 全感官具身循环") print("=" * 60) print("好奇心驱动 + 状态持久化 + 自我模型 + 元学习") print(f"运行 {n_steps} 步（Ctrl+C 中断，状态自动保存）\n") @@ -151,12 +151,13 @@ def live(n_steps: int, device: str, safe_mode: bool = False, unsafe: bool = Fals if info['step'] % 10 == 0: print(f" step {info['step']:4d} | mod {info['modality']:8s} | " f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | " - f"success {info['success']:.2f} | weak={info['weakness']} | " + f"success {info['success']:.2f} | focus={info['consensus_focus']} | " + f"weak={info['weakness']} | " f"mem {info['memory_count']}") - mind.run(n_steps=n_steps, interval=0.2, save_every=30, on_step=on_step) + runtime.run(n_steps=n_steps, interval=0.2, save_every=30, on_step=on_step) - print("\n" + mind.introspect()) + print("\n" + runtime.introspect()) # 可视化 if log: diff --git a/main_chat.py b/main_chat.py index 175ca0e..18b346f 100644 --- a/main_chat.py +++ b/main_chat.py @@ -22,6 +22,7 @@ from pathlib import Path from jspaceai import ( LanguageConfig, JSpaceLanguageModel, EvolutionTrainer, + OnlineLanguageLearner, CharTokenizer, load_chinese_corpus, ) from train_chat import clean_corpus @@ -168,43 +169,6 @@ def generate_response(model, tokenizer, prompt: str, n_new: int = 60, return tokenizer.decode(generated) -def online_learn(model, config, tokenizer, text: str | None, user_input: str, device: str): - """在线学习用户输入 + 回放一段 Shakespeare 防遗忘""" - import torch.nn.functional as F - # 把用户输入作为新语料学习 - user_tokens = tokenizer.encode(user_input) - if len(user_tokens) < 4: - return # 太短不学 - - # 重复用户输入凑够 seq_len - seq_len = 48 - if len(user_tokens) < seq_len: - user_tokens = user_tokens * (seq_len // len(user_tokens) + 1) - user_seq = user_tokens[:seq_len] - token_tensor = torch.tensor([user_seq], dtype=torch.long).to(device) - - # forward + next-token loss - model.train() - model.zero_grad() - logits, _ = model(token_tensor) - pred = logits[:, :-1] - target = token_tensor[:, 1:] - loss = F.cross_entropy( - pred.reshape(-1, config.vocab_size), - target.reshape(-1), - ) - loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) - - # 手动 SGD step（EvolutionTrainer 内部有 EWC，这里简化用直接 step） - with torch.no_grad(): - for p in model.parameters(): - if p.grad is not None: - p -= 5e-3 * p.grad - model.eval() - return loss.item() - - def chat(model, config, tokenizer, text, device: str): """交互对话循环""" print("\n" + "=" * 60) @@ -215,6 +179,7 @@ def chat(model, config, tokenizer, text, device: str): print("=" * 60 + "\n") model.eval() + learner = OnlineLanguageLearner(model, config, tokenizer, device=device) while True: try: user_input = input("你: ").strip() @@ -238,6 +203,7 @@ def chat(model, config, tokenizer, text, device: str): continue elif cmd == '/reset': model = JSpaceLanguageModel(config).to(device) + learner = OnlineLanguageLearner(model, config, tokenizer, device=device) print("(模型已重置为随机初始化)") continue elif cmd.startswith('/train'): @@ -245,6 +211,7 @@ def chat(model, config, tokenizer, text, device: str): n = int(parts[1]) if len(parts) > 1 else 50 text = ensure_corpus(text) train(model, tokenizer, text, n, device) + learner = OnlineLanguageLearner(model, config, tokenizer, device=device) model.eval() continue else: @@ -252,7 +219,7 @@ def chat(model, config, tokenizer, text, device: str): continue # 在线学习用户输入 - loss = online_learn(model, config, tokenizer, text, user_input, device) + learn_stats = learner.learn_text(user_input) # 生成回复 response = generate_response( @@ -260,8 +227,8 @@ def chat(model, config, tokenizer, text, device: str): n_new=40, temperature=0.8, top_k=5, ) print(f"AI: {response}") - if loss is not None: - print(f" (学习 loss={loss:.3f})") + if learn_stats is not None: + print(f" (学习 loss={learn_stats['loss']:.3f} replay={learn_stats['replay_loss']:.3f} step={learn_stats['step']})") def generate_once(model, tokenizer, prompt: str, n_new: int = 80, diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 76ed3b0..a5ae40a 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -1,17 +1,27 @@ import unittest +from pathlib import Path +import tempfile import torch from jspaceai import ( + ActionEvent, + ActionPolicy, CharTokenizer, + ConsensusSnapshot, + InMemoryVectorMemoryStore, JSpaceConfig, JSpaceLanguageModel, JSpaceModel, LanguageConfig, MultimodalConfig, MultimodalJSpaceModel, + OnlineLanguageLearner, + WorkspaceEvent, + WorkspaceRuntime, ) from main_chat import generate_response +from jspaceai import compose_action_params class SmokeTests(unittest.TestCase): @@ -87,6 +97,7 @@ class SmokeTests(unittest.TestCase): ) self.assertEqual(tuple(outputs["w"].shape), (1, 16)) + self.assertIn("alpha", info) self.assertEqual(len(info["w_trajectory"]), 2) def test_tokenizer_unknown_maps_to_zero(self): @@ -95,6 +106,169 @@ class SmokeTests(unittest.TestCase): self.assertEqual(tokenizer.encode("?"), [0]) self.assertEqual(tokenizer.decode([0]), "") + def test_consensus_snapshot_extracts_primary_slot(self): + w = torch.ones(1, 4) + alpha = torch.tensor([[0.1, 0.7, 0.2]]) + + snapshot = ConsensusSnapshot.from_workspace( + w, + alpha, + modality="text", + labels=["vision", "language", "memory"], + ) + + self.assertEqual(snapshot.primary_slot().label, "language") + self.assertGreater(snapshot.confidence, 0.0) + + def test_online_language_learner_tracks_session_state(self): + tokenizer = CharTokenizer.from_text("学而时习之学而时习之") + config = LanguageConfig( + vocab_size=tokenizer.vocab_size, + embed_dim=8, + input_dim=8, + workspace_dim=16, + expert_dim=8, + num_experts=3, + ode_steps=1, + noise_std=0.0, + ) + model = JSpaceLanguageModel(config) + learner = OnlineLanguageLearner( + model, + config, + tokenizer, + seq_len=8, + replay_batch_size=1, + consolidate_every=2, + ) + + first = learner.learn_text("学而时习之") + second = learner.learn_text("学而时习之") + + self.assertEqual(first["step"], 1) + self.assertEqual(second["step"], 2) + self.assertGreaterEqual(second["replay_loss"], 0.0) + + def test_compose_action_params_respects_discrete_mode(self): + raw = torch.tensor([0.5, -0.25, 0.3, -0.4, 0.8]).numpy() + + move = compose_action_params(1, raw) + left_click = compose_action_params(2, raw) + + self.assertEqual(move.tolist(), [0.5, -0.25, 0.0, 0.0, 0.0]) + self.assertEqual(left_click.tolist(), [0.0, 0.0, 1.0, 0.0, 0.0]) + + def test_action_policy_decides_from_workspace(self): + policy = ActionPolicy(workspace_dim=4, exploration=0.0) + policy.value_model.action_weights[4, 0] = 2.0 + w = torch.tensor([[1.0, 0.0, 0.0, 0.0]]) + + decision = policy.decide(w) + + self.assertEqual(decision.action_idx, 4) + self.assertEqual(decision.action_name, "scroll") + + def test_workspace_event_and_memory_store_round_trip(self): + event = WorkspaceEvent.from_tensor( + torch.tensor([[1.0, 0.0, 0.0, 0.0]]), + modality="text", + step=7, + context={"source": "test"}, + ) + action = ActionEvent.from_action_info({ + "action_idx": 1, + "action_name": "mouse_move", + "action_params": [0.1, 0.0, 0.0, 0.0, 0.0], + "executed": True, + "risk": 0.0, + }, step=7) + memory = InMemoryVectorMemoryStore(capacity=4, workspace_dim=4) + + memory.put(event) + results = memory.query(torch.tensor([[0.9, 0.0, 0.0, 0.0]]), top_k=1) + + self.assertEqual(event.to_dict()["step"], 7) + self.assertEqual(action.to_dict()["action_name"], "mouse_move") + self.assertEqual(results[0].event.modality, "text") + self.assertGreater(results[0].similarity, 0.9) + + def test_workspace_runtime_runs_minimal_loop(self): + class DummySenses: + def start(self): + return None + + def stop(self): + return None + + class DummyAudio: + def stop(self): + return None + + class DummyMemory: + def __init__(self): + self.items = [] + + def store(self, w, context=None): + self.items.append((w, context or {})) + + def size(self): + return len(self.items) + + class DummyConfig: + workspace_dim = 4 + + class DummyAgent: + def __init__(self): + self.config = DummyConfig() + self.state = { + "w": torch.zeros(1, 4), + "m": [torch.zeros(1, 2)], + } + self.senses = DummySenses() + self.audio_actuator = DummyAudio() + self.hippocampus = DummyMemory() + self.last_consensus = None + self.learned = [] + + def perceive(self): + return {} + + def think(self, sensory_data): + del sensory_data + self.state["w"] = self.state["w"] + 0.25 + return self.state["w"], "idle" + + def decide_and_act(self, w, modality): + del w, modality + return { + "action_idx": 1, + "action_name": "mouse_move", + "action_params": [0.0, 0.0, 0.0, 0.0, 0.0], + "raw_action_params": [0.0, 0.0, 0.0, 0.0, 0.0], + "executed": False, + "action_strength": 0.0, + "risk": 0.0, + } + + def remember(self, w, context): + self.hippocampus.store(w[0].numpy(), context) + + def learn(self, w, action_idx, reward=0.0): + del w + self.learned.append((action_idx, reward)) + + agent = DummyAgent() + with tempfile.TemporaryDirectory() as tmpdir: + runtime = WorkspaceRuntime(agent, save_dir=Path(tmpdir), device="cpu") + info = runtime.step() + + self.assertEqual(info["step"], 1) + self.assertEqual(info["modality"], "idle") + self.assertEqual(info["memory_count"], 1) + self.assertEqual(agent.learned[0][0], 1) + self.assertEqual(info["workspace_event"]["step"], 1) + self.assertEqual(info["action_event"]["action_name"], "mouse_move") + if __name__ == "__main__": unittest.main()