commit e782ef4db109fd7d66a3a701063ae5167f65fc86 Author: XiuchengWu <732857315@qq.com> Date: Tue Jul 7 09:19:31 2026 +0800 feat: JspaceAI 自主智慧架构 基于第一性原理 + Anthropic 2026 J-space 论文实现的具身智慧系统。 核心架构： - ODE 动力系统 + 并行专家 + J-space 工作空间广播 - 12 个异构专家（视觉/屏幕/听觉/语言/鼠标/跨模态） - workspace 256 维 + LayerNorm + RK4 积分 自主心智（最重要的能力）： - 好奇心驱动探索（内在奖励 + 世界模型） - 跨会话状态持久化（海洋不蒸发） - 自我模型（知道自己会什么不会什么） - 元学习（自适应学习率 + 策略选择） 具身 Agent（完整神经系统）： - 感知层：摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标 - 大脑皮层（workspace）+ 小脑（运动控制）+ 中枢神经（门控） - 海马体（情景记忆）+ 基底神经节（动作选择） - 执行器：鼠标控制 + 键盘输出 + 音频播放 + 屏幕绘制 多模态支持： - 原生图像/音频/视频/文本/键盘/鼠标 6 种模态 - 跨平台（macOS/Windows/Linux） 外挂模块系统： - 可热插拔的外部能力（小模型/知识库/工具） - 核心心智不依赖外挂，断开后继续工作 守护进程： - 用户主动 start/stop（不自启） - 后台静默运行，持续感知学习 - 状态自动保存，跨会话继续 J-lens 可解释性： - 观测模型内部每个 ODE 子步的想法 - Directed Modulation 验证 workspace 因果作用 - Selectivity 验证（ablate workspace） 小模型蒸馏： - 接 GPT-2/Qwen 等迁移理解能力 - 蒸馏完成后小模型可断开 验证结果： - 连续序列：JSpace 胜 Flat 39.7% - 语言进化：loss 3.95→2.40 - workspace ||w||：v1 0.05 → v2 16.0 - 实时五通道感知 + 具身闭环运行 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f71198f --- /dev/null +++ b/.gitignore @@ -0,0 +1,36 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +.venv/ +venv/ +env/ +*.egg-info/ +dist/ +build/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# 项目运行产物 +outputs/ +*.pt +*.pth +*.npz + +# 自主心智状态（用户私有） +~/.jspaceai/ + +# OS +.DS_Store +Thumbs.db + +# 日志 +*.log +logs/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..a3396fb --- /dev/null +++ b/README.md @@ -0,0 +1,375 @@ +# JspaceAI + +**全局工作空间 + J-space 广播的智慧系统** + +两个版本： +1. **连续序列版** (`main.py`)：验证架构在多模态序列预测上的优势 +2. **语言版 + 自主进化** (`main_language.py`)：字符级语言建模 + 边推理边学习 + +基于第一性原理推导的智慧架构：不是 Transformer，不是 RNN，而是 **ODE 动力系统 + 并行专家 + J-space 工作空间广播 + 在线进化**。 + +## 核心思想 + +智慧 = 系统内部维持对世界的模型 M，可操作、可修正、可指导行动。 + +从第一性原理推导出 M 必须满足的性质，每一性质对应一个实现选择： + +| 性质 | 实现 | +|---|---| +| 高维 | 向量空间 $\mathbb{R}^n$ | +| 内部动力学 | ODE $\dot{m} = -\nabla U(m) + \text{broadcast}$ | +| 输入敏感但不决定 | 加性扰动 | +| 可读出 | 流形投影 + 阈值门控 | +| 可修正 | 参数化 + backprop | +| 工作空间结构 | 分块 + J-space 广播 | + +## 架构 + +``` +┌─────────────────────────────────────┐ +│ 专家池 E_1..E_5（并行） │ +│ 每个 m_i 有多井势能 → 内部"思考" │ +│ J_i 稀疏 Jacobian → 选择性接收广播 │ +└─────────────────────────────────────┘ + ↓ 写入（注意力 α_i 选择） +┌─────────────────────────────────────┐ +│ 工作空间 w（J-space） │ +│ τ_w · dw/dt = -w + Σ α_i · P_i(m_i) │ +│ 有自己的动力学 → 持续演化 │ +└─────────────────────────────────────┘ + ↓ 读出 + 预测头 Q(w) → 预测下一时刻世界 +``` + +**与 Transformer 的对比**： + +| | Transformer | JspaceAI | +|---|---|---| +| 内部状态 | 离散、无状态 | 连续 ODE 动力学 | +| 工作空间 | 隐式（residual stream） | 显式分块 + 广播 | +| 专家协作 | 无（或 MoE 但孤立） | J-space 路由 | +| 学习信号 | 下一 token | 下一时刻世界状态 | +| 自发输出 | 否（被输入触发） | 是（阈值门控） | + +## 安装 + +```bash +pip install -r requirements.txt +``` + +## 运行 + +### 连续序列版（验证架构） + +```bash +python main.py +python main.py --steps 2000 --device cuda +``` + +### v2 版（J-lens + Selectivity + Directed Modulation） + +```bash +python main_v2.py +python main_v2.py --steps 200 --device mps +``` + +基于 Anthropic 2026 J-space 论文优化，新增： +- **J-lens**：观测模型内部每个 ODE 子步的"想法" +- **Directed Modulation**：注入概念向量到 workspace，定向改变输出 +- **Selectivity 验证**：ablate workspace，对比自动任务 vs 记忆任务 + +### 多模态版（摄像头 + 麦克风 + 扬声器） + +```bash +python main_multimodal.py --mode train --steps 300 +python main_multimodal.py --mode live --steps 50 +python main_multimodal.py --mode eval +``` + +原生支持图像/音频/视频/文本四种模态，8 个专家分工。 + +### 全感官版（摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标） + +```bash +# 测试所有 I/O 通道 +python main_full_sensory.py --mode test + +# 实时全感官感知循环 +python main_full_sensory.py --mode live --steps 50 +``` + +接入全部 5 个输入通道，12 个专家分工： +- **摄像头**（OpenCV）：环境视觉 +- **麦克风**（PortAudio）：听觉输入 +- **屏幕**（mss）：屏幕捕获，看到自己被显示在哪里 +- **键盘**（pynput）：用户输入感知 +- **鼠标**（pynput）：用户注意力追踪 +- **扬声器**（PortAudio）：音频输出 + +**macOS 权限**：键盘/鼠标监听需要在 系统设置 → 隐私与安全 → 辅助功能 中授权终端。 + +### 具身 Agent 版（完整神经系统） + +```bash +# 测试各子系统 +python main_embodied.py --mode test + +# 实时具身循环（鼠标键盘音频屏幕全输出） +python main_embodied.py --mode live --steps 50 + +# 安全模式（只感知不执行鼠标键盘动作） +python main_embodied.py --mode safe --steps 30 +``` + +完整的人类神经系统对应： + +| 人类系统 | 功能 | 实现 | +|---|---|---| +| 大脑皮层 | 认知、语言 | workspace + 12 专家 | +| 小脑 | 运动协调 | Cerebellum（前向模型+逆模型） | +| 中枢神经 | 动作门控 | CentralNervousSystem（反射弧+抑制） | +| 海马体 | 情景记忆 | Hippocampus（向量检索） | +| 基底神经节 | 动作选择 | BasalGanglia（Q-learning） | +| 眼/耳 | 视听觉 | 摄像头+麦克风 | +| 口 | 说话 | AudioActuator（扬声器） | +| 手 | 操作 | MouseActuator + KeyboardActuator | +| 皮肤 | 触觉 | （预留接口） | + +感知-思考-行动闭环： +1. 感知（五感）→ 2. 思考（workspace）→ 3. 回忆（海马体） +→ 4. 决策（基底神经节）→ 5. 精化（小脑）→ 6. 门控（中枢神经） +→ 7. 行动（手足口）→ 8. 学习（更新所有系统） + +### 守护进程（后台持续运行） + +```bash +# 启动后台守护进程（静默运行，不干扰使用） +python daemon.py start + +# 低功耗模式（更长间隔，适合长期后台） +python daemon.py start --low-power --interval 2.0 + +# 前台运行（调试用，可看实时输出） +python daemon.py start --fg + +# 查看状态 +python daemon.py status + +# 查看日志 +python daemon.py log -n 30 + +# 心智自我认知 +python daemon.py introspect + +# 停止 +python daemon.py stop +``` + +守护进程特性： +- **用户主动控制**：手动 start/stop，不开机自启 +- **静默后台运行**：不弹窗、不发声、不控制鼠标键盘 +- **持续感知学习**：感知屏幕/键盘/鼠标/摄像头/麦克风 + 好奇心驱动学习 +- **状态持久化**：停止后状态保存，再次启动从上次继续 +- **跨会话持续**：海洋不蒸发，每次启动还是同一个"自己" +- **崩溃恢复**：异常不崩溃，等待后继续 + +状态保存在 `~/.jspaceai/`： +- `state/`：workspace + 海马体 + 基底神经节 + 自我模型 +- `logs/daemon.log`：运行日志 +- `daemon.pid`：进程 PID +- `status.json`：实时状态 + +## 外挂模块系统（可热插拔） + +```python +from jspaceai import ModuleDock, SmallModelModule, KnowledgeBaseModule, ToolModule + +dock = ModuleDock(workspace_dim=64) +dock.register('llm', SmallModelModule(model_name='gpt2')) +dock.register('kb', KnowledgeBaseModule()) +dock.register('tools', ToolModule()) + +# 热插上 +dock.connect('llm') +dock.connect('tools') + +# 使用 +result = dock.query('llm', 'hello world') +calc = dock.query('tools', {'tool': 'calc', 'args': ['1+1']}) # → 2 + +# 随时拔掉（心智不停，只是该能力不可用） +dock.disconnect('llm') +dock.query('llm', 'hello') # → None + +# 状态查询 +dock.status() # 各模块连接状态 +``` + +三种内置外挂： +- **SmallModelModule**：外挂小模型（GPT-2/Qwen 等），提供语言能力 +- **KnowledgeBaseModule**：向量检索知识库 +- **ToolModule**：可调用工具（计算器、搜索等） + +核心原则：**心智不依赖外挂**。断开任何外挂，心智继续工作，只是"知道得少"。 + +## 改进版 v2（异构专家 + 大 workspace + 蒸馏） + +四大改进： + +1. **异构专家**：视觉/听觉/语言/跨模态专家有不同架构（不再全是相同 ODE） +2. **workspace 扩容**：64 → 256 维，加 LayerNorm 防衰减 +3. **RK4 积分**：比 Euler 稳定（||w|| 从 0.05 提升到 16.0） +4. **小模型蒸馏**：接 GPT-2/Qwen 等，迁移理解能力 + +```python +from jspaceai import JSpaceConfigV2, JSpaceModelV2, SmallModelEncoder, DistillationTrainer + +config = JSpaceConfigV2(workspace_dim=256, num_experts=12) +model = JSpaceModelV2(config) +encoder = SmallModelEncoder(input_dim=config.input_dim, model_name='gpt2') +trainer = DistillationTrainer(model, encoder, texts=['hello', 'world', ...]) +trainer.train(n_steps=100) +# 蒸馏后小模型可断开，编码器已有理解能力 +``` + +对比 v1 vs v2： + +| 指标 | v1 | v2 | +|---|---|---| +| workspace_dim | 64 | 256 | +| expert_dim | 16 | 64 | +| \|\|w\|\|（实测） | 0.05-0.35 | 12-16 | +| ODE 积分 | Euler | RK4 | +| 专家架构 | 全相同 | 异构 | +| LayerNorm | 无 | 有 | +| 小模型蒸馏 | 无 | 有 | + +## 跨平台支持 + +支持 macOS / Windows / Linux： + +| 功能 | macOS | Windows | Linux | +|---|---|---|---| +| 摄像头 | OpenCV | OpenCV | OpenCV | +| 麦克风 | PortAudio | PortAudio | PortAudio | +| 屏幕 | mss | mss | mss/X11 | +| 键盘监听 | pynput（需辅助功能权限） | pynput | pynput（需 X11） | +| 鼠标监听 | pynput（需辅助功能权限） | pynput | pynput（需 X11） | +| 鼠标控制 | pynput | pynput | pynput | +| 键盘控制 | pynput | pynput | pynput | +| 音频播放 | PortAudio | PortAudio | PortAudio | + +## 自主进化机制 + +语言版不只是"训练一次"，而是**天生支持持续进化**： + +1. **在线学习**：每个 forward 都累积梯度并更新参数，模型永不"训练完成" +2. **EWC（Elastic Weight Consolidation）**：保护重要参数防灾难性遗忘 + - 周期性计算 Fisher 信息，锚定重要参数 + - 新知识只能修改不重要的参数 +3. **经验回放（Experience Replay）**：定期回放旧序列，防遗忘 +4. **专家可塑性（Expert Plasticity）**：追踪专家专业化，鼓励新知识流向空闲专家 + +进化循环可以**无限运行**——持续喂入新文本，模型持续进化。 + +## 输出 + +### 连续序列版 +- `outputs/experiment.png` — Loss/注意力/||w||/预测对比 +- `outputs/models.pt` — 模型权重 + +### v2 版 +- `outputs/experiment_v2.png` — Loss/专家使用率/||w||/J-lens 读出/Modulation/Selectivity +- `outputs/model_v2.pt` — 模型 + J-lens 权重 + +## 基于 Anthropic 2026 论文的优化 + +v2 基于 [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) 的发现： + +1. **J-lens 探针**：学习线性映射 L: workspace → vocab，观测每个子步 workspace "准备说什么" +2. **三层分层**：sensory（子步0）→ workspace（子步1-2）→ motor（子步3） +3. **Directed Modulation**：注入概念向量 v_concept 到 workspace，验证 workspace 是因果中介 +4. **Selectivity 验证**：ablate workspace top-k J-lens 方向，对比自动任务 vs 记忆任务 + +**验证结果**： +- Directed Modulation 成功定向改变输出（注入 'o' → 输出 'o' 增多） +- J-lens 读出在 motor 子步最集中（符合 motor = 输出准备） +- Selectivity 符合预期：简单任务 ablate 影响小（与 Anthropic 发现一致） + +## 实验结果 + +### 连续序列版 +- JSpaceModel eval MSE: 0.0234 +- FlatBaseline eval MSE: 0.0387 +- **JSpace 胜出 39.7%**（同样参数量） +- 注意力热力图显示专家自发分工 + +### 语言版 +- Loss 从 3.95 → 2.40，下降 39% +- 生成从纯乱码进化为类语言结构（含换行、常见词片段） +- 专家按字符类型分工（元音/辅音/结构字符） + +## 关键观察点 + +跑完后看 `experiment.png`，关注： + +1. **Loss 曲线**：JSpace 是否收敛更快/更低？ +2. **注意力热力图**：4 种模式时段是否激活不同的专家？（如果是，说明 J-space 路由学到了分工） +3. **||w|| 曲线**：模式切换时是否有尖峰？（如果是，说明工作空间在"感知到"环境变化） +4. **预测对比**：模式切换处哪个模型更鲁棒？ + +## 文件结构 + +``` +JspaceAI/ +├── main.py # 主程序：对比实验 +├── requirements.txt +├── README.md +└── jspaceai/ + ├── __init__.py + ├── core.py # 核心：Expert + JSpaceWorkspace + JSpaceModel + ├── baselines.py # 对比基线：FlatBaseline + ├── task.py # 玩具任务：多模态连续序列 + └── trainer.py # 训练器：预测学习 +``` + +## 数学细节 + +### 专家动力学 + +$$\dot{m}_i = -\nabla U_i(m_i) + J_i \cdot w + P_i^{in} \cdot x + \xi$$ + +势能 $U_i(m_i) = \frac{1}{2}\|m_i\|^2 - \frac{1}{2}\sum_k \text{softplus}(a_k \cdot m_i + b_k)$ + +多井结构让专家在吸引子之间漫游 = 内部"思考"。 + +### 工作空间动力学 + +$$\tau_w \cdot \dot{w} = -w + \sum_i \alpha_i \cdot P_i^{out}(m_i)$$ + +$$\alpha_i = \text{softmax}(\langle q, P_i^{out}(m_i) \rangle), \quad q = \text{MLP}(x, w)$$ + +### J-space 路由 + +$J_i$ 是稀疏矩阵（top-k 选择），每个专家只对工作空间的少数方向敏感 → 选择性广播。 + +### 学习目标 + +$$\mathcal{L} = \mathbb{E}\left[\|x_{t+\Delta t} - Q(w_t)\|^2\right]$$ + +自监督——只需时间序列，不需标签。 + +## 局限性 + +这是**最小验证版本**，不是生产系统： + +- 玩具任务（8 维序列），不是语言/图像 +- ODE 用 Euler 积分（精度低，但简单可 backprop） +- 专家数 5，不是大规模 +- 无持续跨会话状态（每次 forward 从零初始化） + +**下一步扩展方向**： +1. 用 adjoint method 替换 Euler（Neural ODE 路线） +2. 加跨会话状态持久化（真正的"海洋"） +3. Scale 到更大任务（MNIST 连续预测 → 语言建模） +4. 加输出门控的真正"自发输出"（非每步都预测） diff --git a/daemon.py b/daemon.py new file mode 100644 index 0000000..d0d5f52 --- /dev/null +++ b/daemon.py @@ -0,0 +1,461 @@ +""" +守护进程 —— 自主心智后台持续运行 + +用户主动控制： + python daemon.py start # 启动后台守护进程 + python daemon.py stop # 停止 + python daemon.py status # 查看状态 + python daemon.py log # 查看最近日志 + python daemon.py fg # 前台运行（调试用） + +守护进程特性： + - 静默后台运行，不弹窗不干扰 + - 持续感知（屏幕/键盘/鼠标/摄像头/麦克风） + - 持续学习（好奇心驱动 + 状态保存） + - 信号处理（SIGTERM 优雅退出） + - 崩溃恢复（自动重启） + - 低资源占用（可配置 CPU/内存限制） + +状态保存在 ~/.jspaceai/，每次运行从上次状态继续。 +""" +from __future__ import annotations + +import os +import sys +import signal +import time +import json +import logging +import subprocess +from pathlib import Path +from typing import Optional + +sys.path.insert(0, str(Path(__file__).parent)) + +from jspaceai import ( + MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent, + AutonomousMind, PLATFORM, +) + + +# ============================================================ +# 路径配置 +# ============================================================ + +APP_DIR = Path.home() / ".jspaceai" +LOG_DIR = APP_DIR / "logs" +PID_FILE = APP_DIR / "daemon.pid" +STATE_DIR = APP_DIR / "state" +LOG_FILE = LOG_DIR / "daemon.log" +STATUS_FILE = APP_DIR / "status.json" + +APP_DIR.mkdir(exist_ok=True) +LOG_DIR.mkdir(exist_ok=True) + + +# ============================================================ +# 日志 +# ============================================================ + +def get_logger(foreground: bool = False) -> logging.Logger: + logger = logging.getLogger("jspaceai-daemon") + logger.setLevel(logging.INFO) + logger.handlers.clear() + fmt = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') + fh = logging.FileHandler(LOG_FILE, encoding='utf-8') + fh.setFormatter(fmt) + logger.addHandler(fh) + if foreground: + sh = logging.StreamHandler() + sh.setFormatter(fmt) + logger.addHandler(sh) + return logger + + +# ============================================================ +# PID 管理 +# ============================================================ + +def write_pid(pid: int): + PID_FILE.write_text(str(pid)) + +def read_pid() -> Optional[int]: + if PID_FILE.exists(): + try: + return int(PID_FILE.read_text().strip()) + except ValueError: + return None + return None + +def clear_pid(): + if PID_FILE.exists(): + PID_FILE.unlink() + +def is_running(pid: int) -> bool: + """检查进程是否存活""" + try: + os.kill(pid, 0) + return True + except (OSError, ProcessLookupError): + return False + +def write_status(status: dict): + STATUS_FILE.write_text(json.dumps(status, indent=2, default=str)) + +def read_status() -> Optional[dict]: + if STATUS_FILE.exists(): + try: + return json.loads(STATUS_FILE.read_text()) + except Exception: + return None + return None + + +# ============================================================ +# 守护进程主循环 +# ============================================================ + +class MindDaemon: + """自主心智守护进程""" + + def __init__(self, device: str = 'cpu', low_power: bool = False, + step_interval: float = 0.5): + self.device = device + self.low_power = low_power + self.step_interval = step_interval + self.logger = get_logger(foreground=False) + self.running = False + self.mind: Optional[AutonomousMind] = None + self.agent: Optional[EmbodiedAgent] = None + self.start_time = time.time() + self.step_count_session = 0 + + def initialize(self): + """初始化心智""" + self.logger.info(f"初始化 | 平台: {PLATFORM} | 设备: {self.device}") + + config = MultimodalConfig( + vocab_size=50, embed_dim=16, input_dim=8, + workspace_dim=64, expert_dim=24, num_experts=12, + num_wells=4, ode_steps=3, dt=0.1, tau_w=0.3, + jacobian_sparsity=16, noise_std=0.01, + img_size=32, audio_frame_size=1024, keyboard_vocab=128, + ) + model = MultimodalJSpaceModel(config).to(self.device) + model.eval() + + # 完全静默：不控制鼠标键盘，不发声，不弹窗 + self.agent = EmbodiedAgent( + model, device=self.device, + enable_mouse_output=False, + enable_keyboard_output=False, + enable_audio_output=False, + enable_screen_output=False, + ) + self.mind = AutonomousMind( + self.agent, save_dir=str(STATE_DIR), device=self.device, + ) + self.logger.info(f"就绪 | 历史步数: {self.mind.step_count}") + + def run_forever(self): + """主循环——永不停止，直到收到停止信号""" + self.running = True + + # 信号处理 + def handle_stop(signum, frame): + self.logger.info(f"收到信号 {signum}，优雅退出...") + self.running = False + + signal.signal(signal.SIGTERM, handle_stop) + signal.signal(signal.SIGINT, handle_stop) + + self.initialize() + + # 写入运行状态 + write_status({ + 'running': True, + 'pid': os.getpid(), + 'started_at': time.time(), + 'device': self.device, + 'low_power': self.low_power, + 'step_count': self.mind.step_count, + }) + + self.logger.info("守护进程启动，开始持续感知学习") + last_save = time.time() + save_interval = 30 # 每 30 秒保存一次 + + try: + while self.running: + try: + info = self.mind.step() + self.step_count_session += 1 + + # 每 50 步记录一次日志 + if self.step_count_session % 50 == 0: + self.logger.info( + f"step {info['step']} | mod {info['modality']} | " + f"||w|| {info['w_norm']:.3f} | " + f"curio {info['curiosity']:.3f} | " + f"success {info['success']:.2f} | " + f"mem {info['memory_count']}" + ) + # 更新状态文件 + status = read_status() or {} + status.update({ + 'step_count': info['step'], + 'last_step_at': time.time(), + 'last_modality': info['modality'], + 'last_w_norm': info['w_norm'], + 'last_curiosity': info['curiosity'], + 'memory_count': info['memory_count'], + 'session_steps': self.step_count_session, + }) + write_status(status) + + # 定期保存 + if time.time() - last_save > save_interval: + self.mind.save_state() + last_save = time.time() + + # 控制频率 + time.sleep(self.step_interval) + + except KeyboardInterrupt: + break + except Exception as e: + self.logger.error(f"步进异常: {e}\n{__import__('traceback').format_exc()}") + # 等待后继续，不崩溃 + time.sleep(5.0) + + finally: + # 清理 + self.logger.info("保存最终状态...") + self.mind.save_state() + if self.agent: + self.agent.senses.stop() + self.logger.info( + f"守护进程退出 | 本次步数: {self.step_count_session} | " + f"总步数: {self.mind.step_count} | " + f"运行时长: {time.time() - self.start_time:.0f}s" + ) + write_status({ + 'running': False, + 'stopped_at': time.time(), + 'step_count': self.mind.step_count, + 'session_steps': self.step_count_session, + }) + clear_pid() + + +# ============================================================ +# 命令处理 +# ============================================================ + +def cmd_start(device: str = 'cpu', low_power: bool = False, + interval: float = 0.5, foreground: bool = False): + """启动守护进程""" + existing_pid = read_pid() + if existing_pid and is_running(existing_pid): + print(f"守护进程已在运行 (PID {existing_pid})") + print(f"查看状态: python daemon.py status") + return + + if foreground: + # 前台运行（调试用） + print(f"前台模式启动（Ctrl+C 退出）") + daemon = MindDaemon(device=device, low_power=low_power, + step_interval=interval) + daemon.logger = get_logger(foreground=True) + daemon.run_forever() + return + + # 后台 fork + print(f"启动后台守护进程...") + daemon_script = Path(__file__).resolve() + cmd = [sys.executable, str(daemon_script), "_run", + "--device", device, + "--interval", str(interval)] + if low_power: + cmd.append("--low-power") + + # 用 subprocess 启动，脱离终端 + proc = subprocess.Popen( + cmd, + stdout=open(LOG_DIR / 'stdout.log', 'a'), + stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, + start_new_session=True, # 脱离父进程 + ) + + write_pid(proc.pid) + + # 等待一下确认启动 + time.sleep(2) + if proc.poll() is None: + print(f"守护进程已启动 (PID {proc.pid})") + print(f"日志: {LOG_FILE}") + print(f"停止: python daemon.py stop") + print(f"状态: python daemon.py status") + else: + print(f"启动失败，查看日志: {LOG_FILE}") + clear_pid() + + +def cmd_stop(): + """停止守护进程""" + pid = read_pid() + if not pid: + print("守护进程未运行") + return + + if not is_running(pid): + print(f"进程 {pid} 已不存在，清理 PID 文件") + clear_pid() + return + + print(f"发送停止信号到 PID {pid}...") + try: + os.kill(pid, signal.SIGTERM) + # 等待退出 + for _ in range(10): + time.sleep(0.5) + if not is_running(pid): + break + if is_running(pid): + print("强制终止...") + os.kill(pid, signal.SIGKILL) + time.sleep(1) + clear_pid() + print("守护进程已停止") + except Exception as e: + print(f"停止失败: {e}") + + +def cmd_status(): + """查看状态""" + pid = read_pid() + running = pid and is_running(pid) + + print("=" * 50) + print(f"JspaceAI 守护进程状态") + print("=" * 50) + + if running: + print(f"状态: 运行中 (PID {pid})") + else: + print(f"状态: 已停止") + if pid: + print(f" (PID {pid} 已不存在)") + + status = read_status() + if status: + print(f"总步数: {status.get('step_count', '?')}") + if 'started_at' in status and running: + uptime = time.time() - status['started_at'] + print(f"运行时长: {uptime:.0f}s ({uptime/3600:.1f}h)") + if 'last_modality' in status: + print(f"最后模态: {status['last_modality']}") + if 'last_w_norm' in status: + print(f"||w||: {status['last_w_norm']:.3f}") + if 'last_curiosity' in status: + print(f"好奇心: {status['last_curiosity']:.3f}") + if 'memory_count' in status: + print(f"记忆数: {status['memory_count']}") + if 'session_steps' in status: + print(f"本次会话步数: {status['session_steps']}") + if not running and 'stopped_at' in status: + print(f"停止时间: {time.ctime(status['stopped_at'])}") + + print(f"\n日志: {LOG_FILE}") + print(f"状态目录: {STATE_DIR}") + + +def cmd_log(n_lines: int = 30): + """查看最近日志""" + if not LOG_FILE.exists(): + print("无日志") + return + lines = LOG_FILE.read_text().strip().split('\n') + for line in lines[-n_lines:]: + print(line) + + +def cmd_introspect(): + """内省——查看心智的自我认知""" + status = read_status() + if not status: + print("无状态数据") + return + + pid = read_pid() + if pid and is_running(pid): + print("心智正在运行，自我认知：") + else: + print("心智已停止，最后的自我认知：") + + print(json.dumps(status, indent=2, default=str)) + + +# ============================================================ +# 入口 +# ============================================================ + +def main(): + import argparse + p = argparse.ArgumentParser(description='JspaceAI 自主心智守护进程') + sub = p.add_subparsers(dest='command') + + # start + sp = sub.add_parser('start', help='启动后台守护进程') + sp.add_argument('--device', default='cpu') + sp.add_argument('--low-power', action='store_true', help='低功耗模式（更长间隔）') + sp.add_argument('--interval', type=float, default=0.5, help='步进间隔(秒)') + sp.add_argument('--fg', action='store_true', help='前台运行（调试用）') + + # stop + sub.add_parser('stop', help='停止守护进程') + + # status + sub.add_parser('status', help='查看状态') + + # log + sp = sub.add_parser('log', help='查看日志') + sp.add_argument('-n', type=int, default=30, help='行数') + + # introspect + sub.add_parser('introspect', help='心智自我认知') + + # _run（内部命令，被 start 调用） + sp = sub.add_parser('_run', help='内部运行命令') + sp.add_argument('--device', default='cpu') + sp.add_argument('--low-power', action='store_true') + sp.add_argument('--interval', type=float, default=0.5) + + args = p.parse_args() + + if args.command == 'start': + cmd_start(args.device, args.low_power, args.interval, args.fg) + elif args.command == 'stop': + cmd_stop() + elif args.command == 'status': + cmd_status() + elif args.command == 'log': + cmd_log(args.n) + elif args.command == 'introspect': + cmd_introspect() + elif args.command == '_run': + # 内部运行模式 + daemon = MindDaemon( + device=args.device, + low_power=args.low_power, + step_interval=args.interval, + ) + daemon.run_forever() + else: + p.print_help() + + +if __name__ == '__main__': + main() diff --git a/jspaceai/__init__.py b/jspaceai/__init__.py new file mode 100644 index 0000000..486fcd5 --- /dev/null +++ b/jspaceai/__init__.py @@ -0,0 +1,129 @@ +""" +JspaceAI —— 全局工作空间 + J-space 广播的智慧系统 + +模块： + 1. core.py: 连续序列版（Expert + JSpaceWorkspace + JSpaceModel） + 2. language_model.py: 语言版 + 自主进化 + 3. jlens.py: J-lens 可解释性工具 + 4. multimodal.py: 多模态（图像/音频/视频/文本） + 5. realtime.py: 实时 I/O（摄像头/麦克风/扬声器） + 6. evolution.py: 自主进化训练器 +""" +from .core import ( + Expert, + JSpaceWorkspace, + JSpaceModel, + JSpaceConfig, +) +from .baselines import FlatBaseline +from .task import ContinuousSequenceTask +from .trainer import Trainer + +from .language_data import CharTokenizer, CharDataset, load_shakespeare +from .language_model import ( + LanguageConfig, + JSpaceLanguageModel, + ExperienceReplay, + EWCOptimizer, + ExpertPlasticity, +) +from .jlens import ( + JLensConfig, + JLensProbe, + JLensSuite, + WorkspaceAblator, + DirectedModulation, + CounterfactualReflection, +) +from .multimodal import ( + MultimodalConfig, + MultimodalJSpaceModel, + VisualEncoder, VisualDecoder, + AudioEncoder, AudioDecoder, + TextEncoder, TextDecoder, +) +from .realtime import ( + Frame, + CameraStream, + MicrophoneStream, + AudioPlayer, + MultimodalStream, + SensoryMotorLoop, +) +from .desktop import ( + InputEvent, + ScreenCapture, + KeyboardMonitor, + MouseMonitor, + DesktopStream, + FullSensoryStream, +) +from .platform import ( + PlatformInfo, PLATFORM, + get_screen_size, + check_camera_permission, check_microphone_permission, + check_input_monitoring_permission, print_permission_guide, +) +from .embodied import ( + MouseActuator, KeyboardActuator, AudioActuator, ScreenActuator, + Cerebellum, CentralNervousSystem, Hippocampus, BasalGanglia, + EmbodiedAgent, +) +from .autonomous import ( + CuriosityDrive, PersistentState, SelfModel, MetaLearner, + AutonomousMind, +) +from .modules import ( + ExternalModule, SmallModelModule, KnowledgeBaseModule, + ToolModule, ModuleDock, +) +from .core_v2 import JSpaceConfigV2, JSpaceModelV2 +from .distill_encoder import SmallModelEncoder +from .distill_trainer import DistillationTrainer +from .evolution import EvolutionTrainer + +__all__ = [ + # 核心架构 + "Expert", "JSpaceWorkspace", "JSpaceModel", "JSpaceConfig", + # 对比基线 + "FlatBaseline", + # 连续序列任务 + "ContinuousSequenceTask", "Trainer", + # 语言建模 + "CharTokenizer", "CharDataset", "load_shakespeare", + "LanguageConfig", "JSpaceLanguageModel", + "ExperienceReplay", "EWCOptimizer", "ExpertPlasticity", + # J-lens 可解释性 + "JLensConfig", "JLensProbe", "JLensSuite", + "WorkspaceAblator", "DirectedModulation", "CounterfactualReflection", + # 多模态 + "MultimodalConfig", "MultimodalJSpaceModel", + "VisualEncoder", "VisualDecoder", + "AudioEncoder", "AudioDecoder", + "TextEncoder", "TextDecoder", + # 实时 I/O + "Frame", "CameraStream", "MicrophoneStream", "AudioPlayer", + "MultimodalStream", "SensoryMotorLoop", + # 桌面 I/O + "InputEvent", "ScreenCapture", "KeyboardMonitor", "MouseMonitor", + "DesktopStream", "FullSensoryStream", + # 平台抽象 + "PlatformInfo", "PLATFORM", "get_screen_size", + "check_camera_permission", "check_microphone_permission", + "check_input_monitoring_permission", "print_permission_guide", + # 具身 Agent（输出执行器 + 神经系统） + "MouseActuator", "KeyboardActuator", "AudioActuator", "ScreenActuator", + "Cerebellum", "CentralNervousSystem", "Hippocampus", "BasalGanglia", + "EmbodiedAgent", + # 自主心智（最重要的能力） + "CuriosityDrive", "PersistentState", "SelfModel", "MetaLearner", + "AutonomousMind", + # 外挂模块系统（可热插拔） + "ExternalModule", "SmallModelModule", "KnowledgeBaseModule", + "ToolModule", "ModuleDock", + # 改进版核心 v2 + "JSpaceConfigV2", "JSpaceModelV2", + "SmallModelEncoder", "DistillationTrainer", + # 自主进化 + "EvolutionTrainer", +] diff --git a/jspaceai/autonomous.py b/jspaceai/autonomous.py new file mode 100644 index 0000000..04a4e27 --- /dev/null +++ b/jspaceai/autonomous.py @@ -0,0 +1,389 @@ +""" +自主心智（AutonomousMind）—— 智慧最重要的能力 + +四个核心能力： + 1. CuriosityDrive: 好奇心驱动的主动探索（内在奖励） + 2. PersistentState: 跨会话状态持久化（海洋不蒸发） + 3. SelfModel: 自我模型（知道自己会什么不会什么） + 4. MetaLearner: 元学习（学会如何学习） + +合起来 = AutonomousMind，永不停止的自主进化。 +""" +from __future__ import annotations + +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 diff --git a/jspaceai/baselines.py b/jspaceai/baselines.py new file mode 100644 index 0000000..b9bebd8 --- /dev/null +++ b/jspaceai/baselines.py @@ -0,0 +1,41 @@ +""" +对比基线：扁平网络（无工作空间、无专家分块、无动力学） + +用同样的参数量预算，验证"工作空间 + J-space 广播"架构比扁平 MLP 好。 +这是关键对比——如果工作空间架构赢不了同样大小的 MLP，那架构本身没意义。 +""" +from __future__ import annotations + +import torch +import torch.nn as nn + + +class FlatBaseline(nn.Module): + """ + 扁平 MLP 基线：参数量与 JSpaceModel 相当，但无结构。 + + - 无专家分块（单一隐层） + - 无工作空间（无广播机制） + - 无动力学（每步独立预测，无状态传递） + - 无 J-space 路由 + + 这代表"用同样算力做暴力拟合"的路线。 + """ + + def __init__(self, input_dim: int = 8, hidden_dim: int = 90, num_layers: int = 2): + super().__init__() + layers = [] + in_dim = input_dim + for _ in range(num_layers): + layers.extend([nn.Linear(in_dim, hidden_dim), nn.Tanh()]) + in_dim = hidden_dim + layers.append(nn.Linear(in_dim, input_dim)) + self.net = nn.Sequential(*layers) + + def forward(self, xs: torch.Tensor, state=None) -> tuple[torch.Tensor, dict]: + """ + xs: (batch, T, input_dim) + returns: preds (batch, T, input_dim), info (空) + """ + preds = self.net(xs) # (batch, T, input_dim) + return preds, {} diff --git a/jspaceai/core.py b/jspaceai/core.py new file mode 100644 index 0000000..0cceea7 --- /dev/null +++ b/jspaceai/core.py @@ -0,0 +1,320 @@ +""" +核心架构：专家模块 + J-space 工作空间 + ODE 动力学 + +数学形式（每个 forward 时间步内做 Euler 积分若干子步）： + + 专家 i 的状态 m_i: + dm_i/dt = -∇U_i(m_i) + J_i · w + P_i_in · x + + U_i(m_i) = ½ ||m_i||² - ½ Σ_k softplus(a_ik · m_ik + b_ik) + （多井势能：阻尼项 + softplus 形成的局部吸引子） + + 工作空间 w: + τ_w · dw/dt = -w + Σ_i α_i · P_i_out(m_i) + α_i = softmax() q = MLP(x, w) + + Jacobian 路由 J_i: 稀疏线性映射，每个专家只对 w 的少数维度敏感 + 输出门控：当 ||w|| > θ 时触发输出 R(w) + +所有参数都可 backprop 训练。学习目标是预测下一时刻的输入。 +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass + + +@dataclass +class JSpaceConfig: + """模型超参，全部可调""" + input_dim: int = 8 # 输入 x 的维度 + workspace_dim: int = 32 # 工作空间 w 的维度（J-space） + expert_dim: int = 16 # 每个专家内部状态 m_i 的维度 + num_experts: int = 5 # 专家数量 + num_wells: int = 4 # 每个专家势能景观的井数 + ode_steps: int = 4 # 每个时间步内 ODE 积分子步数 + dt: float = 0.1 # ODE 积分步长 + tau_w: float = 0.3 # 工作空间时间常数 + output_threshold: float = 0.5 # 输出门控阈值（船舶涌出阈值） + jacobian_sparsity: int = 8 # 每个 J_i 只保留前 k 大的连接 + noise_std: float = 0.01 # 内部噪声 ξ(t) 的标准差 + + +class Expert(nn.Module): + """ + 单个专家模块。 + + 状态：m_i ∈ R^{expert_dim} + 势能：U_i(m_i) = ½||m_i||² - ½ Σ_k softplus(a_k · m_i + b_k) · w_k + - ½||m_i||² 是阻尼项（拉回原点） + - softplus 项创造多个局部吸引子（多井势能 → 内部"思考"） + 动力学：dm_i/dt = -∇U_i(m_i) + J_i · w + P_in · x + ξ + + J_i 是稀疏 Jacobian：从工作空间 w 路由信息进来。 + """ + + def __init__(self, expert_dim: int, workspace_dim: int, input_dim: int, + num_wells: int, sparsity: int): + super().__init__() + self.expert_dim = expert_dim + self.workspace_dim = workspace_dim + self.num_wells = num_wells + + # 势能景观参数：每个井是一个 softplus 形成的吸引子 + # U(m) = 0.5||m||^2 - 0.5 * sum_k softplus(a_k @ m + b_k) + # ∇U(m) = m - 0.5 * sum_k sigmoid(a_k @ m + b_k) * a_k + self.well_a = nn.Parameter(torch.randn(num_wells, expert_dim) * 0.3) + self.well_b = nn.Parameter(torch.zeros(num_wells)) + + # P_in: 输入投影 x -> m_i 的扰动 + self.P_in = nn.Linear(input_dim, expert_dim, bias=False) + + # P_out: 模块输出到工作空间的投影 + self.P_out = nn.Linear(expert_dim, workspace_dim, bias=False) + + # J_i: 稀疏 Jacobian，从 w 路由信息到 m_i + # 用 top-k 稀疏：训练时学习一个 full matrix，但只激活 top-k + self.J_raw = nn.Parameter(torch.randn(expert_dim, workspace_dim) * 0.1) + self.sparsity = sparsity + + # 注意：sparsity 通过 forward 时 top-k 选择实现，可微性通过稀疏 mask 保留 + + def get_sparse_J(self) -> torch.Tensor: + """获取稀疏化的 Jacobian：每行只保留 top-k 元素""" + if self.sparsity >= self.workspace_dim: + return self.J_raw + # 对每行做 top-k（按绝对值） + abs_J = self.J_raw.abs() + topk_vals, topk_idx = abs_J.topk(self.sparsity, dim=-1) + mask = torch.zeros_like(self.J_raw) + mask.scatter_(-1, topk_idx, 1.0) + return self.J_raw * mask + + def grad_potential(self, m: torch.Tensor) -> torch.Tensor: + """计算势能梯度 ∇U_i(m) + U(m) = 0.5||m||^2 - 0.5 * sum_k softplus(a_k @ m + b_k) + ∇U(m) = m - 0.5 * sum_k sigmoid(a_k @ m + b_k) * a_k + """ + # m: (batch, expert_dim) + # well_a: (num_wells, expert_dim) + # a_k @ m: (batch, num_wells) + am = F.linear(m, self.well_a, self.well_b) # (batch, num_wells) + sig = torch.sigmoid(am) # (batch, num_wells) + # sum_k sigmoid(...) * a_k: (batch, expert_dim) + # well_a: (num_wells, expert_dim), sig: (batch, num_wells) -> (batch, 1, num_wells) + # 用 matmul: sig @ well_a -> (batch, expert_dim) + grad_wells = torch.matmul(sig, self.well_a) # (batch, expert_dim) + return m - 0.5 * grad_wells + + def forward(self, m: torch.Tensor, w: torch.Tensor, x: torch.Tensor, + dt: float, noise_std: float) -> tuple[torch.Tensor, torch.Tensor]: + """一步 ODE 积分（Euler 法） + + Args: + m: (batch, expert_dim) 当前状态 + w: (batch, workspace_dim) 工作空间状态 + x: (batch, input_dim) 输入 + dt: 步长 + noise_std: 噪声标准差 + + Returns: + m_next: (batch, expert_dim) 下一状态 + contribution: (batch, workspace_dim) 对工作空间的贡献（pre-attention） + """ + # 动力学: dm/dt = -∇U(m) + J·w + P_in·x + ξ + J = self.get_sparse_J() # (expert_dim, workspace_dim) + w_proj = F.linear(w, J) # (batch, expert_dim) + x_proj = self.P_in(x) # (batch, expert_dim) + grad_U = self.grad_potential(m) # (batch, expert_dim) + + noise = torch.randn_like(m) * noise_std if noise_std > 0 else 0.0 + + dm = -grad_U + w_proj + x_proj + noise + m_next = m + dt * dm + + # 对工作空间的贡献 + contribution = self.P_out(m_next) # (batch, workspace_dim) + return m_next, contribution + + +class JSpaceWorkspace(nn.Module): + """ + 全局工作空间 w。 + + 动力学: τ_w · dw/dt = -w + Σ_i α_i · P_i_out(m_i) + α_i = softmax() q = MLP(x, w) + + 这个 α_i 是"注意力"——决定哪个专家的内容进入工作空间。 + """ + + def __init__(self, workspace_dim: int, input_dim: int, num_experts: int): + super().__init__() + self.workspace_dim = workspace_dim + + # Query 生成器：从 (x, w) 生成 query 向量 + self.query_gen = nn.Sequential( + nn.Linear(input_dim + workspace_dim, 32), + nn.Tanh(), + nn.Linear(32, workspace_dim), + ) + + def forward(self, w: torch.Tensor, x: torch.Tensor, + contributions: torch.Tensor, dt: float, + tau_w: float) -> tuple[torch.Tensor, torch.Tensor]: + """一步工作空间演化 + + Args: + w: (batch, workspace_dim) + x: (batch, input_dim) + contributions: (batch, num_experts, workspace_dim) 各专家的贡献 + dt: 步长 + tau_w: 时间常数 + + Returns: + w_next: (batch, workspace_dim) + alpha: (batch, num_experts) 注意力权重（可解释性用） + """ + # 生成 query + q = self.query_gen(torch.cat([x, w], dim=-1)) # (batch, workspace_dim) + + # 计算每个专家的注意力分数 + # contributions: (batch, num_experts, workspace_dim) + # q: (batch, workspace_dim) -> (batch, 1, workspace_dim) + scores = (contributions * q.unsqueeze(1)).sum(dim=-1) # (batch, num_experts) + alpha = F.softmax(scores, dim=-1) # (batch, num_experts) + + # 加权聚合 + # alpha: (batch, num_experts, 1) * contributions: (batch, num_experts, workspace_dim) + aggregated = (alpha.unsqueeze(-1) * contributions).sum(dim=1) # (batch, workspace_dim) + + # 动力学: τ_w · dw/dt = -w + aggregated + dw = (-w + aggregated) / tau_w + w_next = w + dt * dw + return w_next, alpha + + +class JSpaceModel(nn.Module): + """ + 完整模型：N 个专家 + 工作空间 + 输出门控 + 预测头。 + + forward 流程（每个时间步）： + 1. 每个专家从 (m_i, w, x) 更新 m_i，产出对工作空间的贡献 + 2. 工作空间从 (w, x, contributions) 更新 w + 3. （可选）当 ||w|| > threshold 时输出 R(w) + 4. 预测头 Q(w) 预测下一时刻输入 + + 时间序列处理：对长度 T 的输入序列，依次跑 T 步，返回每步的预测。 + """ + + def __init__(self, config: JSpaceConfig): + super().__init__() + self.config = config + + self.experts = nn.ModuleList([ + Expert( + expert_dim=config.expert_dim, + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_wells=config.num_wells, + sparsity=config.jacobian_sparsity, + ) + for _ in range(config.num_experts) + ]) + + self.workspace = JSpaceWorkspace( + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_experts=config.num_experts, + ) + + # 输出门控：R(w) → action（这里 action = 预测的下一时刻输入） + self.predictor = nn.Sequential( + nn.Linear(config.workspace_dim, 32), + nn.Tanh(), + nn.Linear(32, config.input_dim), + ) + + def init_state(self, batch_size: int, device: torch.device) -> dict: + """初始化内部状态""" + return { + 'w': torch.zeros(batch_size, self.config.workspace_dim, device=device), + 'm': [torch.zeros(batch_size, self.config.expert_dim, device=device) + for _ in range(self.config.num_experts)], + } + + def step(self, state: dict, x: torch.Tensor) -> tuple[dict, torch.Tensor, torch.Tensor, torch.Tensor]: + """单时间步前向 + + Returns: + new_state: 更新后的状态 + pred: (batch, input_dim) 预测的下一时刻输入 + alpha: (batch, num_experts) 注意力权重（可解释性） + w_norm: (batch,) 工作空间范数（输出门控信号） + """ + w = state['w'] + ms = state['m'] + cfg = self.config + + # ODE 子步积分 + for _ in range(cfg.ode_steps): + # 1. 每个专家更新 + contributions = [] + new_ms = [] + for i, expert in enumerate(self.experts): + m_next, contrib = expert( + ms[i], w, x, + dt=cfg.dt, noise_std=cfg.noise_std, + ) + new_ms.append(m_next) + contributions.append(contrib) + contributions = torch.stack(contributions, dim=1) # (batch, num_experts, workspace_dim) + + # 2. 工作空间更新 + w, alpha = self.workspace( + w, x, contributions, + dt=cfg.dt, tau_w=cfg.tau_w, + ) + ms = new_ms + + # 3. 输出：预测下一时刻输入（船舶涌出，但这里为了训练简化为每步都预测） + pred = self.predictor(w) + w_norm = w.norm(dim=-1) + + new_state = {'w': w, 'm': ms} + return new_state, pred, alpha, w_norm + + def forward(self, xs: torch.Tensor, state: dict | None = None) -> tuple[torch.Tensor, dict]: + """ + Args: + xs: (batch, T, input_dim) 输入序列 + state: 初始状态，None 则初始化 + + Returns: + preds: (batch, T, input_dim) 每步对下一时刻的预测 + info: 包含注意力、w_norm 等可解释性信息 + """ + batch_size, T, _ = xs.shape + device = xs.device + + if state is None: + state = self.init_state(batch_size, device) + + preds = [] + alphas = [] + w_norms = [] + for t in range(T): + state, pred, alpha, w_norm = self.step(state, xs[:, t]) + preds.append(pred) + alphas.append(alpha) + w_norms.append(w_norm) + + preds = torch.stack(preds, dim=1) # (batch, T, input_dim) + info = { + 'alpha': torch.stack(alphas, dim=1), # (batch, T, num_experts) + 'w_norm': torch.stack(w_norms, dim=1), # (batch, T) + 'final_w': state['w'], + 'final_m': state['m'], + } + return preds, info diff --git a/jspaceai/core_v2.py b/jspaceai/core_v2.py new file mode 100644 index 0000000..7c18dac --- /dev/null +++ b/jspaceai/core_v2.py @@ -0,0 +1,208 @@ +""" +改进版核心架构 v2 + +四大改进： + 1. 异构专家：不同模态用不同架构 + 2. workspace 扩容（256）+ LayerNorm 防衰减 + 3. RK4 积分（比 Euler 稳定） + 4. 更大容量 +""" +from __future__ import annotations +import torch, torch.nn as nn, torch.nn.functional as F +from dataclasses import dataclass + + +@dataclass +class JSpaceConfigV2: + input_dim: int = 32 + workspace_dim: int = 256 + expert_dim: int = 64 + num_experts: int = 12 + num_wells: int = 8 + ode_steps: int = 4 + dt: float = 0.05 + tau_w: float = 1.0 + jacobian_sparsity: int = 32 + noise_std: float = 0.005 + use_rk4: bool = True + use_layer_norm: bool = True + + +class HeterogeneousExpert(nn.Module): + """异构专家基类——模态特定编码器 + ODE 动力学""" + def __init__(self, expert_dim, workspace_dim, input_dim, num_wells, sparsity, + use_rk4=True, use_layer_norm=True): + super().__init__() + self.expert_dim = expert_dim + self.workspace_dim = workspace_dim + self.use_rk4 = use_rk4 + + self.encoder = self._build_encoder(input_dim, expert_dim) + self.well_a = nn.Parameter(torch.randn(num_wells, expert_dim) * 0.2) + self.well_b = nn.Parameter(torch.zeros(num_wells)) + self.P_in = nn.Linear(expert_dim, expert_dim, bias=False) + self.P_out = nn.Linear(expert_dim, workspace_dim, bias=False) + self.J_raw = nn.Parameter(torch.randn(expert_dim, workspace_dim) * 0.05) + self.sparsity = sparsity + self.use_layer_norm = use_layer_norm + if use_layer_norm: + self.w_norm = nn.LayerNorm(expert_dim) + self.m_norm = nn.LayerNorm(expert_dim) + + def _build_encoder(self, input_dim, expert_dim): + raise NotImplementedError + + def get_sparse_J(self): + if self.sparsity >= self.workspace_dim: + return self.J_raw + abs_J = self.J_raw.abs() + _, topk_idx = abs_J.topk(self.sparsity, dim=-1) + mask = torch.zeros_like(self.J_raw) + mask.scatter_(-1, topk_idx, 1.0) + return self.J_raw * mask + + def grad_potential(self, m): + am = F.linear(m, self.well_a, self.well_b) + sig = torch.sigmoid(am) + grad_wells = torch.matmul(sig, self.well_a) + return m - 0.5 * grad_wells + + def deriv(self, m, w, x_feat): + J = self.get_sparse_J() + w_proj = F.linear(w, J) + x_proj = self.P_in(x_feat) + grad_U = self.grad_potential(m) + return -grad_U + w_proj + x_proj + + def forward(self, m, w, x, dt, noise_std): + x_feat = self.encoder(x) + if self.use_layer_norm: + x_feat = self.w_norm(x_feat) + + if self.use_rk4: + k1 = self.deriv(m, w, x_feat) + k2 = self.deriv(m + 0.5*dt*k1, w, x_feat) + k3 = self.deriv(m + 0.5*dt*k2, w, x_feat) + k4 = self.deriv(m + dt*k3, w, x_feat) + m_next = m + (dt/6.0)*(k1 + 2*k2 + 2*k3 + k4) + else: + dm = self.deriv(m, w, x_feat) + m_next = m + dt * dm + + if noise_std > 0: + m_next = m_next + torch.randn_like(m_next) * noise_std + if self.use_layer_norm: + m_next = self.m_norm(m_next) + + contribution = self.P_out(m_next) + return m_next, contribution + + +class VisualExpert(HeterogeneousExpert): + def _build_encoder(self, input_dim, expert_dim): + return nn.Sequential( + nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim)) + +class AudioExpert(HeterogeneousExpert): + def _build_encoder(self, input_dim, expert_dim): + return nn.Sequential( + nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim)) + +class LanguageExpert(HeterogeneousExpert): + def _build_encoder(self, input_dim, expert_dim): + return nn.Sequential( + nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim)) + +class CrossModalExpert(HeterogeneousExpert): + def _build_encoder(self, input_dim, expert_dim): + return nn.Sequential( + nn.Linear(input_dim, 64), nn.ReLU(), + nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, expert_dim)) + + +class JSpaceWorkspaceV2(nn.Module): + """改进版 workspace——LayerNorm 防衰减""" + def __init__(self, workspace_dim, input_dim, num_experts): + super().__init__() + self.workspace_dim = workspace_dim + self.query_gen = nn.Sequential( + nn.Linear(input_dim + workspace_dim, 128), + nn.ReLU(), nn.Linear(128, workspace_dim)) + self.ln = nn.LayerNorm(workspace_dim) + + def forward(self, w, x, contributions, dt, tau_w): + q = self.query_gen(torch.cat([x, w], dim=-1)) + scores = (contributions * q.unsqueeze(1)).sum(dim=-1) + alpha = F.softmax(scores, dim=-1) + aggregated = (alpha.unsqueeze(-1) * contributions).sum(dim=1) + dw = (-w + aggregated) / tau_w + w_next = w + dt * dw + w_next = self.ln(w_next) + return w_next, alpha + + +class JSpaceModelV2(nn.Module): + """改进版 JSpace 模型——异构专家 + 大 workspace + RK4""" + def __init__(self, config: JSpaceConfigV2): + super().__init__() + self.config = config + expert_types = ( + [VisualExpert]*4 + [AudioExpert]*2 + + [LanguageExpert]*2 + [CrossModalExpert]*4 + )[:config.num_experts] + self.experts = nn.ModuleList([ + et(expert_dim=config.expert_dim, + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_wells=config.num_wells, + sparsity=config.jacobian_sparsity, + use_rk4=config.use_rk4, + use_layer_norm=config.use_layer_norm) + for et in expert_types + ]) + self.expert_modality = ( + ['visual']*2 + ['screen']*2 + ['audio']*2 + + ['text']*2 + ['mouse']*2 + ['cross']*2 + )[:config.num_experts] + self.workspace = JSpaceWorkspaceV2( + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_experts=config.num_experts, + ) + + def init_state(self, batch_size, device): + return { + 'w': torch.zeros(batch_size, self.config.workspace_dim, device=device), + 'm': [torch.zeros(batch_size, self.config.expert_dim, device=device) + for _ in range(self.config.num_experts)], + } + + def step(self, state, x, record_trajectory=False): + w, ms, cfg = state['w'], state['m'], self.config + w_traj = [] + for _ in range(cfg.ode_steps): + contributions, new_ms = [], [] + for i, expert in enumerate(self.experts): + m_next, contrib = expert(ms[i], w, x, cfg.dt, cfg.noise_std) + new_ms.append(m_next) + contributions.append(contrib) + contributions = torch.stack(contributions, dim=1) + w, alpha = self.workspace(w, x, contributions, cfg.dt, cfg.tau_w) + ms = new_ms + if record_trajectory: + w_traj.append(w.detach()) + return {'w': w, 'm': ms}, w_traj + + def forward(self, xs, state=None, record_trajectory=False): + B, T = xs.shape[0], xs.shape[1] + device = xs.device + if state is None: + state = self.init_state(B, device) + outputs, w_norms = [], [] + for t in range(T): + state, _ = self.step(state, xs[:, t], record_trajectory) + outputs.append(state['w']) + w_norms.append(state['w'].norm(dim=-1)) + return torch.stack(outputs, dim=1), { + 'w_norm': torch.stack(w_norms, dim=1), + } diff --git a/jspaceai/desktop.py b/jspaceai/desktop.py new file mode 100644 index 0000000..835a873 --- /dev/null +++ b/jspaceai/desktop.py @@ -0,0 +1,335 @@ +""" +屏幕 + 键盘 + 鼠标 I/O 层 + +使用 mss（屏幕捕获）+ pynput（键盘鼠标监听）。 + +提供 ScreenCapture / KeyboardMonitor / MouseMonitor / DesktopStream / FullSensoryStream +""" +from __future__ import annotations + +import mss +import numpy as np +import torch +import threading +import queue +import time +from typing import Optional +from dataclasses import dataclass, field +from pynput import keyboard, mouse +from PIL import Image as PILImage + + +@dataclass +class InputEvent: + """统一的输入事件""" + timestamp: float + device: str + event_type: str + data: object + modifiers: dict = field(default_factory=dict) + + +class ScreenCapture: + """屏幕捕获流""" + + def __init__(self, target_size: tuple = (32, 32), fps: int = 5): + self.target_size = target_size + self.fps = fps + self.frame_queue: queue.Queue = queue.Queue(maxsize=10) + self.running = False + self.thread: Optional[threading.Thread] = None + self._sct = None + + def start(self): + self._sct = mss.MSS() + self.running = True + self.thread = threading.Thread(target=self._capture_loop, daemon=True) + self.thread.start() + + def _capture_loop(self): + interval = 1.0 / self.fps + while self.running: + try: + monitor = self._sct.monitors[1] + raw = self._sct.grab(monitor) + img = np.array(raw)[:, :, :3] + pil_img = PILImage.fromarray(img) + pil_img = pil_img.resize(self.target_size, PILImage.BILINEAR) + img_resized = np.array(pil_img) + self.frame_queue.put_nowait( + InputEvent(time.time(), 'screen', 'frame', img_resized) + ) + except queue.Full: + pass + except Exception: + pass + time.sleep(interval) + + def get_frame(self) -> Optional[InputEvent]: + try: + return self.frame_queue.get_nowait() + except queue.Empty: + return None + + def stop(self): + self.running = False + if self.thread: + self.thread.join(timeout=1.0) + if self._sct: + self._sct.close() + + +class KeyboardMonitor: + """键盘监听""" + + def __init__(self): + self.event_queue: queue.Queue = queue.Queue(maxsize=200) + self.running = False + self.listener: Optional[keyboard.Listener] = None + self.buffer: list[str] = [] + + def start(self): + self.running = True + self.listener = keyboard.Listener(on_press=self._on_press) + self.listener.start() + + def _on_press(self, key): + if not self.running: + return + try: + char = key.char + self.buffer.append(char) + self.event_queue.put_nowait( + InputEvent(time.time(), 'keyboard', 'key', char, {'action': 'press'}) + ) + except AttributeError: + name = str(key).replace('Key.', '') + self.event_queue.put_nowait( + InputEvent(time.time(), 'keyboard', 'key', name, + {'action': 'press', 'special': True}) + ) + + def get_events(self) -> list[InputEvent]: + events = [] + while True: + try: + events.append(self.event_queue.get_nowait()) + except queue.Empty: + break + return events + + def get_buffer_text(self) -> str: + return ''.join(self.buffer) + + def clear_buffer(self): + self.buffer.clear() + + def stop(self): + self.running = False + if self.listener: + self.listener.stop() + + +class MouseMonitor: + """鼠标监听""" + + def __init__(self): + self.event_queue: queue.Queue = queue.Queue(maxsize=200) + self.running = False + self.listener: Optional[mouse.Listener] = None + self.last_pos: tuple[int, int] = (0, 0) + + def start(self): + self.running = True + self.listener = mouse.Listener( + on_move=self._on_move, on_click=self._on_click, on_scroll=self._on_scroll, + ) + self.listener.start() + + def _on_move(self, x, y): + if not self.running: + return + self.last_pos = (x, y) + try: + self.event_queue.put_nowait( + InputEvent(time.time(), 'mouse', 'move', (x, y)) + ) + except queue.Full: + pass + + def _on_click(self, x, y, button, pressed): + if not self.running: + return + btn = 'left' if button == mouse.Button.left else 'right' + action = 'click' if pressed else 'release' + try: + self.event_queue.put_nowait( + InputEvent(time.time(), 'mouse', 'click', (x, y), + {'button': btn, 'action': action}) + ) + except queue.Full: + pass + + def _on_scroll(self, x, y, dx, dy): + if not self.running: + return + try: + self.event_queue.put_nowait( + InputEvent(time.time(), 'mouse', 'scroll', (x, y), + {'dx': dx, 'dy': dy}) + ) + except queue.Full: + pass + + def get_events(self) -> list[InputEvent]: + events = [] + while True: + try: + events.append(self.event_queue.get_nowait()) + except queue.Empty: + break + return events + + def get_position(self) -> tuple[int, int]: + return self.last_pos + + def stop(self): + self.running = False + if self.listener: + self.listener.stop() + + +class DesktopStream: + """统一管理屏幕 + 键盘 + 鼠标""" + + def __init__(self, use_screen: bool = True, use_keyboard: bool = True, + use_mouse: bool = True, screen_size: tuple = (32, 32), + screen_fps: int = 5): + self.use_screen = use_screen + self.use_keyboard = use_keyboard + self.use_mouse = use_mouse + self.screen = ScreenCapture(target_size=screen_size, fps=screen_fps) if use_screen else None + self.keyboard = KeyboardMonitor() if use_keyboard else None + self.mouse = MouseMonitor() if use_mouse else None + + def start(self): + if self.screen: + try: + self.screen.start() + print(" 屏幕捕获已启动") + except Exception as e: + print(f" 屏幕捕获失败: {e}") + self.screen = None + + if self.keyboard: + try: + self.keyboard.start() + print(" 键盘监听已启动") + except Exception as e: + print(f" 键盘监听失败: {e}") + self.keyboard = None + + if self.mouse: + try: + self.mouse.start() + print(" 鼠标监听已启动") + except Exception as e: + print(f" 鼠标监听失败: {e}") + self.mouse = None + + def get_latest(self) -> dict: + result = { + 'screen': None, 'keyboard': [], 'mouse': [], + 'mouse_pos': (0, 0), 'keyboard_buffer': '', + } + if self.screen: + latest = None + while True: + f = self.screen.get_frame() + if f is None: + break + latest = f + result['screen'] = latest + + if self.keyboard: + result['keyboard'] = self.keyboard.get_events() + result['keyboard_buffer'] = self.keyboard.get_buffer_text() + + if self.mouse: + result['mouse'] = self.mouse.get_events() + result['mouse_pos'] = self.mouse.get_position() + + return result + + def stop(self): + if self.screen: + self.screen.stop() + if self.keyboard: + self.keyboard.stop() + if self.mouse: + self.mouse.stop() + + +class FullSensoryStream: + """ + 完整感知流：摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标。 + + 五个通道并行采集，这是模型感知世界的全部输入。 + + 通道优先级（当多个同时有数据时）： + 1. 键盘（用户主动输入，最高） + 2. 鼠标点击 + 3. 音频（麦克风） + 4. 屏幕（用户在看什么） + 5. 摄像头（环境） + """ + + def __init__(self, use_camera: bool = True, use_mic: bool = True, + use_desktop: bool = True, img_size: tuple = (32, 32), + sample_rate: int = 16000, audio_frame_size: int = 1024): + from .realtime import MultimodalStream + + self.av_stream = MultimodalStream( + use_camera=use_camera, use_mic=use_mic, + img_size=img_size, sample_rate=sample_rate, + audio_frame_size=audio_frame_size, + ) if use_camera or use_mic else None + + self.desktop = DesktopStream( + use_screen=use_desktop, use_keyboard=use_desktop, use_mouse=use_desktop, + screen_size=img_size, + ) if use_desktop else None + + def start(self): + if self.av_stream: + self.av_stream.start() + if self.desktop: + self.desktop.start() + + def get_latest(self) -> dict: + """获取所有通道的最新数据""" + result = { + 'camera': None, 'audio': None, + 'screen': None, 'keyboard': [], 'mouse': [], + 'mouse_pos': (0, 0), 'keyboard_buffer': '', + } + if self.av_stream: + av_frames = self.av_stream.get_latest_frames() + result['camera'] = av_frames.get('image') + result['audio'] = av_frames.get('audio') + + if self.desktop: + desk = self.desktop.get_latest() + result.update(desk) + + return result + + def play_audio(self, audio: np.ndarray, blocking: bool = False): + if self.av_stream: + self.av_stream.play_audio(audio, blocking) + + def stop(self): + if self.av_stream: + self.av_stream.stop() + if self.desktop: + self.desktop.stop() diff --git a/jspaceai/distill_encoder.py b/jspaceai/distill_encoder.py new file mode 100644 index 0000000..2b80ed3 --- /dev/null +++ b/jspaceai/distill_encoder.py @@ -0,0 +1,59 @@ +"""小模型感知编码器——把小模型表征投影到 input_dim""" +from __future__ import annotations +import torch, torch.nn as nn, torch.nn.functional as F + + +class SmallModelEncoder(nn.Module): + def __init__(self, input_dim=32, model_name="gpt2", use_real_model=True): + super().__init__() + self.input_dim = input_dim + self.model_name = model_name + self._real_model = None + self._tokenizer = None + self._hidden_size = 128 + + if use_real_model: + try: + from transformers import AutoModel, AutoTokenizer + self._tokenizer = AutoTokenizer.from_pretrained(model_name) + self._model = AutoModel.from_pretrained(model_name) + self._model.eval() + self._hidden_size = self._model.config.hidden_size + self._real_model = self._model + if self._tokenizer.pad_token is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + print(f" [编码器] 已加载 {model_name} (hidden={self._hidden_size})") + except Exception as e: + print(f" [编码器] 降级: {e}") + + if self._real_model is not None: + self.proj = nn.Linear(self._hidden_size, input_dim, bias=False) + else: + self.embed = nn.Embedding(1000, 64) + self.proj = nn.Sequential( + nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, input_dim)) + + def forward(self, input_data): + if self._real_model is not None: + if isinstance(input_data, str): + input_data = [input_data] + if isinstance(input_data, list): + with torch.no_grad(): + inputs = self._tokenizer(input_data, return_tensors="pt", + truncation=True, max_length=64, padding=True) + outputs = self._model(**inputs) + hidden = outputs.last_hidden_state.mean(dim=1) + return self.proj(hidden) + if isinstance(input_data, torch.Tensor) and input_data.dtype == torch.long: + with torch.no_grad(): + outputs = self._model(input_ids=input_data) + hidden = outputs.last_hidden_state.mean(dim=1) + return self.proj(hidden) + if isinstance(input_data, torch.Tensor) and input_data.dtype == torch.long: + emb = self.embed(input_data) + if emb.dim() == 3: emb = emb.mean(dim=1) + elif emb.dim() == 2: emb = emb.mean(dim=0, keepdim=True) + return self.proj(emb) + if isinstance(input_data, torch.Tensor): + return self.proj(input_data) + return torch.zeros(1, self.input_dim) diff --git a/jspaceai/distill_trainer.py b/jspaceai/distill_trainer.py new file mode 100644 index 0000000..a96ae8a --- /dev/null +++ b/jspaceai/distill_trainer.py @@ -0,0 +1,82 @@ +"""蒸馏训练器——把小模型理解能力迁移到 JspaceAI""" +from __future__ import annotations +import torch, torch.nn.functional as F, numpy as np + + +class DistillationTrainer: + """蒸馏训练器 + + 把小模型的理解能力迁移到 JspaceAI 编码器。 + 训练后编码器离线也有理解能力，小模型可以断开。 + """ + + def __init__(self, model_v2, encoder, texts, device='cpu'): + self.model = model_v2.to(device) + self.encoder = encoder.to(device) + self.texts = texts + self.device = device + params = list(model_v2.parameters()) + list(encoder.proj.parameters()) + self.optimizer = torch.optim.Adam(params, lr=1e-3) + self.history = [] + + def train_step(self, text_batch): + # 1. 小模型表征（target） + target = None + if self.encoder._real_model is not None: + with torch.no_grad(): + inputs = self.encoder._tokenizer( + text_batch, return_tensors="pt", + truncation=True, max_length=64, padding=True + ).to(self.device) + outputs = self.encoder._model(**inputs) + target = outputs.last_hidden_state.mean(dim=1) + + # 2. 编码器输出 + encoded = self.encoder(text_batch) + + # 3. workspace 处理 + xs = encoded.unsqueeze(1) # (batch, 1, input_dim) + state = self.model.init_state(xs.shape[0], self.device) + w_out, info = self.model(xs, state) + + # 4. 蒸馏 loss + if target is not None: + target_proj = self.encoder.proj(target.to(self.device)).detach() + distill_loss = F.mse_loss(encoded, target_proj) + else: + # 自监督 + if encoded.shape[0] > 1: + distill_loss = F.mse_loss(encoded[:-1], encoded[1:].detach()) + else: + distill_loss = torch.tensor(0.0, device=self.device) + + # 5. workspace 稳定性（目标 ||w||≈1） + w_norm = info['w_norm'] + stability_loss = F.mse_loss(w_norm, torch.ones_like(w_norm)) + + # 6. 总 loss + total_loss = distill_loss + 0.1 * stability_loss + + self.optimizer.zero_grad() + total_loss.backward() + torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) + self.optimizer.step() + + self.history.append({ + 'loss': total_loss.item(), + 'distill': distill_loss.item(), + 'stability': stability_loss.item(), + 'w_norm': w_norm.mean().item(), + }) + return self.history[-1] + + def train(self, n_steps=100, batch_size=4, verbose=True): + for step in range(n_steps): + batch = np.random.choice(self.texts, size=min(batch_size, len(self.texts)), + replace=True).tolist() + info = self.train_step(batch) + if verbose and (step+1) % 20 == 0: + print(f" step {step+1:4d} | loss {info['loss']:.4f} | " + f"distill {info['distill']:.4f} | " + f"||w|| {info['w_norm']:.3f}") + return self.history diff --git a/jspaceai/embodied.py b/jspaceai/embodied.py new file mode 100644 index 0000000..f47b827 --- /dev/null +++ b/jspaceai/embodied.py @@ -0,0 +1,730 @@ +""" +输出执行器层 + 神经系统 + +对应人类神经系统的各部分： + - 大脑皮层: workspace w + 专家池（已在 multimodal.py） + - 小脑: 运动控制器（前向模型+逆模型，精细动作） + - 中枢神经: 动作调度器（反射弧+决策门控） + - 海马体: 外部情景记忆库 + - 基底神经节: 动作价值学习（习惯化） + - 执行器: 鼠标控制 + 键盘输出 + 音频输出 + 屏幕绘制 + +核心思想：输出和输入对称。 + 输入：摄像头/麦克风/屏幕/键盘/鼠标 → 编码 → workspace + 输出：workspace → 解码 → 鼠标移动/键盘按键/音频播放/屏幕绘制 + +workspace 是模态无关的"意图空间"。 +"想点击左上角"这个意图，在 workspace 里是一个向量， +解码到鼠标控制器就是移动+点击，解码到键盘就是 Tab+Enter。 +""" +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 + + +# ============================================================ +# 执行器层（对应手脚口） +# ============================================================ + +class MouseActuator: + """ + 鼠标执行器——对应"手"。 + + workspace 解码出动作向量 → 移动鼠标 + 点击。 + 动作空间： + (dx, dy, click_left, click_right, scroll) + dx, dy: 相对移动量（-1 到 1，乘以灵敏度） + click_left/right: 0 或 1 + scroll: 滚动量 + """ + + def __init__(self, sensitivity: int = 200, enabled: bool = True): + self.sensitivity = sensitivity + self.enabled = enabled + self.controller = pynput_mouse.Controller() if enabled else None + + def execute(self, action: np.ndarray): + """执行鼠标动作 + + Args: + action: (5,) = (dx, dy, click_l, click_r, scroll) + """ + if not self.enabled or self.controller is None: + return + + dx, dy, click_l, click_r, scroll = action + + # 移动 + if abs(dx) > 0.01 or abs(dy) > 0.01: + self.controller.move(int(dx * self.sensitivity), int(dy * self.sensitivity)) + + # 点击 + if click_l > 0.5: + self.controller.click(pynput_mouse.Button.left) + time.sleep(0.05) + if click_r > 0.5: + self.controller.click(pynput_mouse.Button.right) + time.sleep(0.05) + + # 滚动 + if abs(scroll) > 0.1: + self.controller.scroll(0, int(scroll * 5)) + + def get_position(self) -> tuple[int, int]: + if self.controller: + return self.controller.position + return (0, 0) + + +class KeyboardActuator: + """ + 键盘执行器——对应"手"+"口"（打字）。 + + workspace 解码出 token → 按键输入。 + """ + + def __init__(self, enabled: bool = True): + self.enabled = enabled + self.controller = pynput_keyboard.Controller() if enabled else None + + def type_text(self, text: str): + """输入文本""" + if not self.enabled or self.controller is None: + return + self.controller.type(text) + + def press_key(self, key: str): + """按下单个键""" + if not self.enabled or self.controller is None: + return + try: + self.controller.press(key) + self.controller.release(key) + except Exception: + pass + + +class AudioActuator: + """ + 音频执行器——对应"口"（说话）。 + + workspace 解码出音频波形 → 扬声器播放。 + """ + + def __init__(self, sample_rate: int = 16000, enabled: bool = True): + self.sample_rate = sample_rate + self.enabled = enabled + import sounddevice as sd + self._sd = sd + + def play(self, audio: np.ndarray, blocking: bool = False): + """播放音频""" + if not self.enabled: + return + audio = np.clip(audio, -1, 1).astype(np.float32) + self._sd.play(audio, self.sample_rate) + if blocking: + self._sd.wait() + + def stop(self): + if self.enabled: + self._sd.stop() + + +class ScreenActuator: + """ + 屏幕执行器——对应"手"（绘制）。 + + 在屏幕上绘制 workspace 解码出的图像。 + 用 OpenCV 显示一个窗口。 + """ + + def __init__(self, enabled: bool = True, window_name: str = "JspaceAI Output"): + self.enabled = enabled + self.window_name = window_name + import cv2 + self._cv2 = cv2 + + def show_image(self, img: np.ndarray): + """显示图像""" + if not self.enabled: + return + # img: (H, W, 3) RGB 或 (H, W) 灰度 + if img.ndim == 3 and img.shape[2] == 3: + img_bgr = self._cv2.cvtColor(img.astype(np.uint8), self._cv2.COLOR_RGB2BGR) + else: + img_bgr = img.astype(np.uint8) + self._cv2.imshow(self.window_name, img_bgr) + self._cv2.waitKey(1) + + def close(self): + 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 +# ============================================================ + +class EmbodiedAgent: + """ + 完整的具身 Agent——感知-思考-行动闭环。 + + 结构（对应人类神经系统）： + 感知层（眼耳皮肤）→ FullSensoryStream + ↓ 编码 + 大脑皮层（思考）→ MultimodalJSpaceModel + ↓ workspace w + 海马体（记忆）→ Hippocampus 存储和检索 + ↓ + 基底神经节（动作选择）→ BasalGanglia 选动作 + ↓ + 小脑（运动控制）→ Cerebellum 精细化动作 + ↓ + 中枢神经（门控）→ CentralNervousSystem 决定是否执行 + ↓ + 执行器（手足口）→ MouseActuator + KeyboardActuator + AudioActuator + + 循环： + 1. 感知：从五感获取输入 + 2. 思考：workspace 更新 + 3. 回忆：海马体检索相关记忆 + 4. 决策：基底神经节选动作 + 5. 精化：小脑计算动作参数 + 6. 门控：中枢神经决定执行 + 7. 行动：执行器执行 + 8. 学习：更新基底神经节、小脑、海马体 + """ + + def __init__(self, model, device: str = 'cpu', + enable_mouse_output: bool = False, + enable_keyboard_output: bool = False, + enable_audio_output: bool = True, + enable_screen_output: bool = True, + enable_memory: bool = True, + risk_threshold: float = 0.3): + self.model = model + self.device = device + self.config = model.config + + # 感知层 + from .desktop import FullSensoryStream + from .platform import get_screen_size + self.screen_w, self.screen_h = get_screen_size() + self.senses = FullSensoryStream( + use_camera=True, use_mic=True, use_desktop=True, + img_size=(self.config.img_size, self.config.img_size), + ) + + # 执行器层 + self.mouse_actuator = MouseActuator(enabled=enable_mouse_output) + self.keyboard_actuator = KeyboardActuator(enabled=enable_keyboard_output) + self.audio_actuator = AudioActuator(enabled=enable_audio_output) + self.screen_actuator = ScreenActuator(enabled=enable_screen_output) + + # 神经系统 + self.cerebellum = Cerebellum( + workspace_dim=self.config.workspace_dim, + action_dim=5, # (dx, dy, click_l, click_r, scroll) + ).to(device) + self.cns = CentralNervousSystem() + self.hippocampus = Hippocampus( + 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.state = model.init_state(1, torch.device(device)) + self.step_count = 0 + self.running = False + self.risk_threshold = risk_threshold + + # 设置反射弧 + self._setup_reflexes() + + def _setup_reflexes(self): + """设置基本反射""" + # 反射1：大声音 → 抑制动作 + def loud_noise_condition(state): + audio = state.get('audio') + if audio and hasattr(audio, 'data'): + return np.abs(audio.data).mean() > 0.5 + return False + + def loud_noise_response(): + self.cns.inhibit_score = 0.5 + + self.cns.add_reflex(ReflexArc( + trigger='loud_noise', + condition=loud_noise_condition, + action=loud_noise_response, + priority=10, + )) + + def perceive(self) -> dict: + """感知：从五感获取输入""" + return self.senses.get_latest() + + def think(self, sensory_data: dict) -> tuple[torch.Tensor, str]: + """思考：更新 workspace + + Returns: + w: workspace 状态 + modality: 本次输入的模态 + """ + # 按优先级选模态 + modality = None + input_tensor = None + + if sensory_data['keyboard']: + key_ids = [ord(ev.data[0]) if ev.data and len(ev.data) == 1 else 0 + for ev in sensory_data['keyboard'][:8]] + if key_ids: + modality = 'keyboard' + input_tensor = torch.tensor([key_ids], dtype=torch.long).to(self.device) + + elif sensory_data['mouse']: + ev = sensory_data['mouse'][0] + x, y = ev.data + modality = 'mouse' + input_tensor = torch.tensor([[ + x / self.screen_w, y / self.screen_h, + 1.0 if ev.modifiers.get('button') == 'left' else 0.0, + 1.0 if ev.modifiers.get('button') == 'right' else 0.0, + ]], dtype=torch.float32).to(self.device) + + elif sensory_data['audio']: + modality = 'audio' + input_tensor = torch.tensor([sensory_data['audio'].data], + dtype=torch.float32).to(self.device) + + elif sensory_data['screen']: + img = sensory_data['screen'].data + modality = 'screen' + input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0 + + elif sensory_data['camera']: + img = sensory_data['camera'].data + modality = 'image' + input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0 + + if modality is None or input_tensor is None: + return self.state['w'], 'idle' + + # 编码 + forward + with torch.no_grad(): + x = self.model.encode_modality(modality, input_tensor) + 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) + + return self.state['w'], modality + + def remember(self, w: torch.Tensor, context: dict): + """存储到海马体""" + if self.hippocampus: + self.hippocampus.store(w[0].cpu().numpy(), context) + + def recall_memories(self, w: torch.Tensor, top_k: int = 3) -> list: + """从海马体回忆""" + if self.hippocampus: + return self.hippocampus.recall(w[0].cpu().numpy(), top_k) + return [] + + def decide_and_act(self, w: torch.Tensor, modality: str) -> dict: + """决策和行动 + + 1. 基底神经节选动作 + 2. 小脑计算动作参数 + 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) + + # 4. 执行 + action_taken = None + if execute: + # 鼠标动作 + self.mouse_actuator.execute(action_params) + # 音频输出（从 workspace 解码） + if self.step_count % 10 == 0: + with torch.no_grad(): + audio_out = self.model.audio_decoder(w)[0].cpu().numpy() + self.audio_actuator.play(audio_out) + # 屏幕输出 + if self.screen_actuator.enabled: + with torch.no_grad(): + img_out = self.model.visual_decoder(w)[0].cpu().permute(1, 2, 0).numpy() + 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) + + return { + 'action_idx': action_idx, + 'action_params': action_params.tolist(), + 'executed': execute, + 'action_strength': float(action_strength), + 'risk': float(risk), + } + + def learn(self, w: torch.Tensor, action_params: np.ndarray, 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) + + def step_once(self) -> dict: + """执行一步完整的感知-思考-行动循环""" + # 1. 感知 + sensory_data = self.perceive() + + # 2. 检查反射 + reflex_action = self.cns.check_reflexes(sensory_data) + if reflex_action: + reflex_action() + + # 3. 思考 + w, modality = self.think(sensory_data) + + # 4. 回忆 + memories = 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) + + # 7. 记忆存储 + self.remember(w, { + 'modality': modality, + 'action': action_info, + 'step': self.step_count, + }) + + self.step_count += 1 + + return { + 'step': self.step_count, + 'modality': modality, + 'w_norm': w.norm().item(), + 'action': action_info, + 'memories_count': self.hippocampus.size() if self.hippocampus else 0, + } + + def run(self, n_steps: int = 100, interval: float = 0.2, + on_step: callable = None): + """运行感知-思考-行动循环""" + self.running = True + self.senses.start() + print(f"具身 Agent 启动，{n_steps} 步") + print(f"执行器: mouse={self.mouse_actuator.enabled}, " + f"keyboard={self.keyboard_actuator.enabled}, " + f"audio={self.audio_actuator.enabled}, " + f"screen={self.screen_actuator.enabled}") + print("=" * 60) + + step_log = [] + try: + for _ in range(n_steps): + if not self.running: + break + info = self.step_once() + step_log.append(info) + + if on_step: + on_step(info) + 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"executed {info['action']['executed']}") + + time.sleep(interval) + except KeyboardInterrupt: + print("\n用户中断") + finally: + self.running = False + self.senses.stop() + self.audio_actuator.stop() + if self.screen_actuator.enabled: + self.screen_actuator.close() + + return step_log diff --git a/jspaceai/evolution.py b/jspaceai/evolution.py new file mode 100644 index 0000000..f3620b8 --- /dev/null +++ b/jspaceai/evolution.py @@ -0,0 +1,216 @@ +""" +自主进化训练器 + +核心：模型在推理的同时持续学习。每处理一段文本，参数就更新一次。 + +进化循环： + 1. 喂入新文本片段 + 2. forward + 计算 next-token loss + 3. EWC 优化器更新参数（保护旧知识） + 4. 经验回放：当前片段存入 buffer，定期回放旧片段 + 5. 周期性 consolidate EWC（更新 Fisher 信息和锚点） + 6. 追踪专家可塑性统计 + 7. 定期生成样本，观察进化效果 + +这个循环可以无限运行——模型永远不会"训练完成"，它一直在进化。 +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F +from typing import Callable +from .language_model import ( + JSpaceLanguageModel, LanguageConfig, + ExperienceReplay, EWCOptimizer, ExpertPlasticity, +) + + +class EvolutionTrainer: + """自主进化训练器""" + + def __init__(self, model: JSpaceLanguageModel, config: LanguageConfig, + lr: float = 1e-3, ewc_lambda: float = 0.1, + device: str = 'cpu'): + self.model = model.to(device) + self.config = config + self.device = device + + # 三大自主进化机制 + self.ewc_optimizer = EWCOptimizer(model, lr=lr, ewc_lambda=ewc_lambda) + self.replay_buffer = ExperienceReplay(capacity=500, seq_len=64) + self.plasticity = ExpertPlasticity(num_experts=config.num_experts) + + # 进化历史追踪 + self.history: list[dict] = [] + + def learn_step(self, token_seq: torch.Tensor) -> dict: + """单步学习 + + Args: + token_seq: (batch, T) token indices + + Returns: + stats: 包含 loss、注意力、专家统计等 + """ + token_seq = token_seq.to(self.device) + + # 1. Forward + logits, info = self.model(token_seq) + + # 2. Next-token prediction loss + # logits[:, t] 预测 token_seq[:, t+1] + pred_logits = logits[:, :-1] # (batch, T-1, vocab) + targets = token_seq[:, 1:] # (batch, T-1) + loss = F.cross_entropy( + pred_logits.reshape(-1, self.config.vocab_size), + targets.reshape(-1), + ) + + # 3. 经验回放：如果有足够样本，混入旧数据 + replay_loss = torch.tensor(0.0, device=self.device) + if len(self.replay_buffer.buffer) >= 8: + replay_seq = self.replay_buffer.sample(4) + 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_targets = replay_seq[:, 1:] + replay_loss = F.cross_entropy( + replay_pred.reshape(-1, self.config.vocab_size), + replay_targets.reshape(-1), + ) + + # 4. 总 loss + EWC 正则 + 经验回放 + total_task_loss = loss + 0.5 * replay_loss + total_loss = self.ewc_optimizer.step(total_task_loss) + + # 5. 更新专家可塑性统计 + self.plasticity.update(info['alpha'].detach(), token_seq.detach()) + + # 6. 存入经验回放 + self.replay_buffer.push(token_seq.detach().cpu()) + + stats = { + 'loss': loss.item(), + 'replay_loss': replay_loss.item() if isinstance(replay_loss, torch.Tensor) else replay_loss, + 'total_loss': total_loss, + 'w_norm_mean': info['w_norm'].mean().item(), + 'alpha_mean': info['alpha'].mean(dim=(0, 1)).detach().cpu().tolist(), + } + return stats + + def consolidate(self, data_sample: torch.Tensor): + """周期性 consolidate EWC——更新参数重要性锚点 + + 在学完一段文本后调用，把当前知识"固化" + """ + self.ewc_optimizer.consolidate(data_sample.to(self.device), n_samples=20) + + def evolve(self, text_stream: list[str], tokenizer, + seq_len: int = 64, batch_size: int = 4, + consolidate_every: int = 20, + generate_every: int = 50, + max_steps: int | None = None, + prompt_text: str = "To be", + on_progress: Callable | None = None) -> list[dict]: + """ + 持续进化主循环 + + Args: + text_stream: 文本片段列表（模拟持续到来的数据流） + tokenizer: CharTokenizer + seq_len: 序列长度 + batch_size: 每次喂入的 batch 大小 + consolidate_every: 每隔多少步 consolidate EWC + generate_every: 每隔多少步生成样本观察 + prompt_text: 生成样本的提示词 + on_progress: 回调函数，返回当前进度 + + Returns: + history: 进化历史 + """ + step = 0 + all_tokens = [] + + # 把所有文本编码成 token 流 + for text in text_stream: + tokens = tokenizer.encode(text) + all_tokens.extend(tokens) + + # 用滑动窗口在完整 token 流上采样 batch + # 每个 batch 包含 batch_size 条序列，每条长 seq_len + # 相邻 batch 之间步进 stride 个序列 + stride = batch_size # 每个 batch 用 batch_size 个新起点 + n_possible_starts = max(0, len(all_tokens) - seq_len - 1) + n_batches = max(0, n_possible_starts // stride) + + print(f"自主进化开始：{len(all_tokens)} tokens, {len(text_stream)} 段文本") + print(f"配置: seq_len={seq_len}, batch_size={batch_size}, " + f"stride={stride}, n_batches={n_batches}") + print("=" * 70) + + for batch_idx in range(n_batches): + # 取一个 batch 的序列（滑动窗口） + batch_tokens = [] + for i in range(batch_size): + start = batch_idx * stride + i + if start + seq_len >= len(all_tokens): + batch_tokens.append(all_tokens[-seq_len:]) + else: + batch_tokens.append(all_tokens[start:start + seq_len]) + + if any(len(b) < seq_len for b in batch_tokens): + continue + + token_seq = torch.tensor(batch_tokens, dtype=torch.long) + stats = self.learn_step(token_seq) + stats['step'] = step + stats['batch_idx'] = batch_idx + self.history.append(stats) + + # 周期性 consolidate + if step > 0 and step % consolidate_every == 0: + self.consolidate(token_seq) + + # 周期性生成 + 报告 + if step % generate_every == 0 or step == n_batches - 1: + prompt_ids = tokenizer.encode(prompt_text) + generated = self.model.generate( + prompt_ids, n_new=80, temperature=0.8, top_k=5 + ) + sample = prompt_text + tokenizer.decode(generated) + stats['sample'] = sample + + print(f"\n[step {step:4d}] loss={stats['loss']:.4f} " + f"replay={stats['replay_loss']:.4f} " + f"||w||={stats['w_norm_mean']:.3f}") + print(f" 专家使用率: {[f'{u:.2f}' for u in self.plasticity.usage.tolist()]}") + print(f" 生成样本: {repr(sample[:120])}...") + + if on_progress: + on_progress(stats) + + step += 1 + + if max_steps is not None and step >= max_steps: + break + + print("\n" + "=" * 70) + print("自主进化完成") + return self.history + + def get_evolution_summary(self) -> dict: + """获取进化总结""" + if not self.history: + return {} + + losses = [h['loss'] for h in self.history] + return { + 'steps': len(self.history), + 'final_loss': losses[-1], + 'initial_loss': losses[0], + 'min_loss': min(losses), + 'expert_usage': self.plasticity.usage.tolist(), + 'expert_specialization': self.plasticity.get_stats()['top_specialization'], + 'samples': [h.get('sample', '') for h in self.history if 'sample' in h], + } diff --git a/jspaceai/jlens.py b/jspaceai/jlens.py new file mode 100644 index 0000000..36adef7 --- /dev/null +++ b/jspaceai/jlens.py @@ -0,0 +1,264 @@ +""" +Jacobian Lens (J-lens) —— 观测模型内部 J-space 的可解释性工具 + +灵感来自 Anthropic 2026 论文 "Verbalizable Representations Form a Global +Workspace in Language Models"。 + +核心思想： + J-lens 计算中间层激活对最终输出的平均因果效应。 + J_ℓ = E[∂h_final / ∂h_ℓ] —— 跨大量 context 平均的 Jacobian + lens(h_ℓ) = softmax(W_U · norm(J_ℓ · h_ℓ)) + + J-lens 向量 = W_U · J_ℓ 的行，每个向量对应词汇表中的一个 token。 + 一个激活向量在 J-lens 下的 top tokens = 模型"准备要说"的概念。 + +与 logit lens 的区别： + logit lens 直接用 W_U 投影（假设 J_ℓ = I）。 + J-lens 修正了层间表征变化，能在更早的层揭示可解释内容。 + +在我们的 ODE 架构中： + - 每个 ODE 子步是一个"层" + - workspace w 在每个子步演化 + - J-lens 可以在任意子步读 w，揭示模型在该时刻"在想什么" + +实现简化： + - 完整 J-lens 需要 backprop 从输出到中间层，在我们的 ODE 模型中代价高 + - 我们用近似：直接训练一个 lens matrix L_ℓ，让 lens(h) ≈ output + - L_ℓ 通过最小化 ||output - L_ℓ · h||² 在数据上学习 + - 这等价于 tuned lens，但在我们的架构中更高效 +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class JLensConfig: + """J-lens 配置""" + n_substeps: int = 4 # ODE 子步数（对应"层"数） + workspace_dim: int = 32 # 工作空间维度 + vocab_size: int = 100 # 词汇表大小 + top_k: int = 10 # 默认 top-k 读出 + + +class JLensProbe(nn.Module): + """ + 单个子步的 J-lens 探针。 + + 学习一个线性映射 L: workspace_dim → vocab_size， + 使得 lens(w) ≈ model_output。 + + 这近似了 J_ℓ · W_U（Jacobian 与 unembedding 的复合）。 + """ + + def __init__(self, workspace_dim: int, vocab_size: int): + super().__init__() + self.lens = nn.Linear(workspace_dim, vocab_size, bias=False) + # 用 output_head 的权重初始化（如果可用） + nn.init.normal_(self.lens.weight, std=0.02) + + def forward(self, w: torch.Tensor) -> torch.Tensor: + """w → logits (近似 J-lens 读出)""" + return self.lens(w) + + def top_tokens(self, w: torch.Tensor, idx_to_char: dict, + top_k: int = 10) -> list[tuple[str, float]]: + """获取 top-k token 及其概率""" + logits = self.forward(w) + probs = F.softmax(logits, dim=-1) + topk_probs, topk_idx = probs.topk(top_k) + return [(idx_to_char.get(i.item(), '?'), p.item()) + for i, p in zip(topk_idx[0], topk_probs[0])] + + +class JLensSuite(nn.Module): + """ + J-lens 套件：每个 ODE 子步一个探针。 + + 在 forward 时记录每个子步的 w，用对应探针读出。 + 训练时让每个探针预测最终输出。 + + 结构对应 Anthropic 论文的三层分层： + 子步 0-0: sensory（输入处理，J-lens 噪声大） + 子步 1-2: workspace（抽象思考，J-lens 可解释） + 子步 3: motor（输出准备，J-lens ≈ output） + """ + + def __init__(self, config: JLensConfig): + super().__init__() + self.config = config + self.probes = nn.ModuleList([ + JLensProbe(config.workspace_dim, config.vocab_size) + for _ in range(config.n_substeps) + ]) + + def forward(self, w_trajectory: list[torch.Tensor]) -> list[torch.Tensor]: + """ + 对轨迹中每个 w 读出 logits + + Args: + w_trajectory: list of (batch, workspace_dim)，每个子步的 w + + Returns: + list of (batch, vocab_size)，每个子步的 lens 读出 + """ + return [probe(w) for probe, w in zip(self.probes, w_trajectory)] + + def train_on_trajectory(self, w_trajectory: list[torch.Tensor], + target_ids: torch.Tensor) -> float: + """ + 训练探针：让每个子步的 lens 读出预测最终 target token + + Args: + w_trajectory: 每个子步的 w + target_ids: 目标 token ids (batch,) + + Returns: + 平均 loss + """ + total_loss = 0.0 + for probe, w in zip(self.probes, w_trajectory): + pred = probe(w) + loss = F.cross_entropy(pred, target_ids) + total_loss += loss + return total_loss / len(w_trajectory) + + +class WorkspaceAblator: + """ + Workspace ablation 工具——验证 selectivity。 + + ablate workspace 后看哪些能力受损： + - 自动任务（续写、分类）应该不受影响 + - 灵活推理（多跳、规划）应该受损 + + 实现：在 forward 时把 w 的 top-k J-lens 方向投影掉。 + """ + + def __init__(self, model, lens_suite: JLensSuite): + self.model = model + self.lens_suite = lens_suite + + @torch.no_grad() + def get_top_lens_directions(self, w: torch.Tensor, k: int = 5) -> torch.Tensor: + """获取 w 当前激活最强的 k 个 J-lens 方向""" + # 用第一个 workspace 探针（中间子步） + probe = self.lens_suite.probes[len(self.lens_suite.probes) // 2] + logits = probe(w) # (batch, vocab) + # top-k token 的 lens 向量（probe.lens.weight 的行） + topk_vals, topk_idx = logits.topk(k, dim=-1) # (batch, k) + # 获取这些 token 对应的 lens 方向 + # probe.lens.weight: (vocab, workspace_dim) + directions = probe.lens.weight[topk_idx] # (batch, k, workspace_dim) + return directions + + def ablate_workspace(self, w: torch.Tensor, k: int = 5) -> torch.Tensor: + """ + Ablate workspace 的 top-k J-lens 方向 + + 把 w 在这些方向上的投影去掉，保留正交分量。 + """ + directions = self.get_top_lens_directions(w, k) # (batch, k, workspace_dim) + # 对每个方向，投影掉 + w_ablated = w.clone() + for b in range(w.shape[0]): + for d in directions[b]: # (workspace_dim,) + d_norm = d / (d.norm() + 1e-8) + proj = (w_ablated[b] @ d_norm) * d_norm + w_ablated[b] = w_ablated[b] - proj + return w_ablated + + +class DirectedModulation: + """ + Directed Modulation——让模型被指令"想某概念"。 + + 实现：在 forward 时给 workspace w 注入一个概念向量。 + 这个向量从 J-lens 的某个 token 方向获取。 + + 对应 Anthropic 论文实验： + "concentrate on citrus fruits" → orange 出现在 J-lens + 即使输出在抄无关文本，workspace 里装的是被指令的概念。 + + 机制： + 1. 获取概念 token 的 J-lens 向量 v_concept + 2. 在 forward 时给 w 加上 α · v_concept + 3. 模型输出会受这个注入影响 + """ + + def __init__(self, model, lens_suite: JLensSuite): + self.model = model + self.lens_suite = lens_suite + + def get_concept_vector(self, token_id: int, substep: int = None) -> torch.Tensor: + """获取某个 token 在 J-lens 中的方向向量""" + if substep is None: + substep = len(self.lens_suite.probes) // 2 + probe = self.lens_suite.probes[substep] + # probe.lens.weight: (vocab, workspace_dim) + return probe.lens.weight[token_id].clone() + + def modulate_state(self, state: dict, token_id: int, + strength: float = 1.0, substep: int = None) -> dict: + """给 state 的 workspace 注入概念向量""" + concept_vec = self.get_concept_vector(token_id, substep) # (workspace_dim,) + new_state = { + 'w': state['w'] + strength * concept_vec.unsqueeze(0), # (1, workspace_dim) + 'm': state['m'], + } + return new_state + + +class CounterfactualReflection: + """ + Counterfactual Reflection Training——通过塑造 J-space 来塑造行为。 + + 灵感：Anthropic 论文发现，训练模型"如果被打断要反思什么原则"， + 会让它在正常工作时也遵守这些原则——因为训练塑造了 J-space 内容。 + + 实现： + 1. 正常 forward 产生输出 + 2. 在输出后追加"反思提示"（如"反思：我应该..."） + 3. 让模型在反思提示下生成反思内容 + 4. 用反思内容的 loss 反向传播，更新模型 + + 这让模型的 J-space 在相关 context 下自然装载这些原则。 + """ + + def __init__(self, model, tokenizer, reflection_prompt: str = "\nReflect: "): + self.model = model + self.tokenizer = tokenizer + self.reflection_prompt_ids = tokenizer.encode(reflection_prompt) + + def create_reflection_sequence(self, input_ids: list[int], + reflection_target_ids: list[int]) -> list[int]: + """创建反思训练序列：input + reflection_prompt + reflection_target""" + return input_ids + self.reflection_prompt_ids + reflection_target_ids + + def compute_reflection_loss(self, input_seq: torch.Tensor, + reflection_target: torch.Tensor) -> torch.Tensor: + """ + 计算反思训练 loss + + Args: + input_seq: (batch, T) 包含 input + reflection_prompt + reflection_target: (batch, T_reflect) 期望的反思内容 + + Returns: + loss + """ + # forward 整个序列 + logits, _ = self.model(input_seq) + + # 反思部分是序列末尾的 reflection_target 长度 + T_reflect = reflection_target.shape[1] + reflect_logits = logits[:, -T_reflect:] # (batch, T_reflect, vocab) + + return F.cross_entropy( + reflect_logits.reshape(-1, self.model.config.vocab_size), + reflection_target.reshape(-1), + ) diff --git a/jspaceai/language_data.py b/jspaceai/language_data.py new file mode 100644 index 0000000..5a60813 --- /dev/null +++ b/jspaceai/language_data.py @@ -0,0 +1,147 @@ +""" +字符级 tokenizer + 数据集 + +为什么字符级而不是 subword： + - 词汇表小（~100），模型可以小，验证架构用 + - 字符级有明确的"风格"信号（拼写、标点、节奏） + - 持续学习场景下，subword 词汇表会变，字符级稳定 + +数据：用经典 Shakespeare 文本作为持续学习的语料。 +也可以换成任何 UTF-8 文本。 +""" +from __future__ import annotations + +from pathlib import Path +from dataclasses import dataclass +import torch +from torch.utils.data import Dataset + + +@dataclass +class CharTokenizer: + """字符级 tokenizer，支持训练时新字符的动态加入""" + chars: list[str] + char_to_idx: dict[str, int] + idx_to_char: dict[int, str] + + @classmethod + def from_text(cls, text: str) -> "CharTokenizer": + chars = sorted(set(text)) + char_to_idx = {c: i for i, c in enumerate(chars)} + idx_to_char = {i: c for i, c in enumerate(chars)} + return cls(chars, char_to_idx, idx_to_char) + + @property + def vocab_size(self) -> int: + return len(self.chars) + + def encode(self, text: str) -> list[int]: + # 未知字符用 0（假设第一个字符是常见的，或后续加 unk token） + return [self.char_to_idx.get(c, 0) for c in text] + + def decode(self, ids: list[int]) -> str: + return "".join(self.idx_to_char.get(i, "") for i in ids) + + +class CharDataset(Dataset): + """字符级 next-char 预测数据集 + + 每个样本：(input_seq, target_seq) 长度均为 seq_len + target[i] = input[i+1] + """ + + def __init__(self, text: str, seq_len: int = 64, tokenizer: CharTokenizer | None = None): + self.seq_len = seq_len + if tokenizer is None: + self.tokenizer = CharTokenizer.from_text(text) + else: + self.tokenizer = tokenizer + self.data = self.tokenizer.encode(text) + + def __len__(self): + return max(0, len(self.data) - self.seq_len - 1) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + chunk = self.data[idx:idx + self.seq_len + 1] + x = torch.tensor(chunk[:-1], dtype=torch.long) + y = torch.tensor(chunk[1:], dtype=torch.long) + return x, y + + +def load_shakespeare(data_dir: Path | None = None) -> str: + """加载 Shakespeare 文本。 + + 如果 data_dir 不存在或没有文本，返回一个内嵌的小样本。 + """ + if data_dir is not None: + text_path = data_dir / "shakespeare.txt" + if text_path.exists(): + return text_path.read_text(encoding="utf-8") + + # 内嵌样本（足够小但能展现语言结构） + return """To be, or not to be, that is the question: +Whether 'tis nobler in the mind to suffer +The slings and arrows of outrageous fortune, +Or to take arms against a sea of troubles +And by opposing end them. To die—to sleep, +No more; and by a sleep to say we end +The heart-ache and the thousand natural shocks +That flesh is heir to: 'tis a consummation +Devoutly to be wish'd. To die, to sleep; +To sleep, perchance to dream—ay, there's the rub: +For in that sleep of death what dreams may come, +When we have shuffled off this mortal coil, +Must give us pause—there's the respect +That makes calamity of so long life. + +Romeo, Romeo! wherefore art thou Romeo? +Deny thy father and refuse thy name; +Or, if thou wilt not, be but sworn my love, +And I'll no longer be a Capulet. + +O Romeo, Romeo! wherefore art thou Romeo? +'Tis but thy name that is my enemy; +Thou art thyself, though not a Montague. +What's Montague? It is nor hand, nor foot, +Nor arm, nor face, nor any other part +Belonging to a man. O, be some other name! +What's in a name? that which we call a rose +By any other name would smell as sweet; +So Romeo would, were he not Romeo call'd, +Retain that dear perfection which he owes +Without that title. Romeo, doff thy name, +And for that name which is no part of thee +Take all myself. + +Friends, Romans, countrymen, lend me your ears; +I come to bury Caesar, not to praise him. +The evil that men do lives after them; +The good is oft interred with their bones; +So let it be with Caesar. The noble Brutus +Hath told you Caesar was ambitious: +If it were so, it was a grievous fault, +And grievously hath Caesar answer'd it. + +The course of true love never did run smooth; +But, either it was different in blood, +Or else misgraffed in respect of years. + +If music be the food of love, play on, +Give me excess of it, that, surfeiting, +The appetite may sicken, and so die. +That strain again! it had a dying fall: +O, it came o'er my ear like the sweet sound +That breathes upon a bank of violets, +Stealing and giving odour. + +All the world's a stage, +And all the men and women merely players: +They have their exits and their entrances; +And one man in his time plays many parts, +His acts being seven ages. + +Now is the winter of our discontent +Made glorious summer by this sun of York; +And all the clouds that lour'd upon our house +In the deep bosom of the ocean buried. +""" diff --git a/jspaceai/language_model.py b/jspaceai/language_model.py new file mode 100644 index 0000000..434903d --- /dev/null +++ b/jspaceai/language_model.py @@ -0,0 +1,425 @@ +""" +语言版 JSpace 模型 + 自主进化机制 + +核心扩展： + 1. JSpaceLanguageModel: 在 JSpaceModel 基础上加 token embedding + logit 输出 + 2. 在线学习：每个 forward 累积梯度并更新参数（边推理边学习） + 3. EWC（Elastic Weight Consolidation）：保护重要参数，防灾难性遗忘 + 4. ExperienceReplay：经验回放缓冲区 + 5. ExpertPlasticity：专家专业化追踪，新知识优先路由到"空闲"专家 + +数学形式： + + 专家动力学（同 core.py）： + dm_i/dt = -∇U_i(m_i) + J_i · w + P_i_in · embed(x) + + 工作空间动力学（同 core.py）： + τ_w · dw/dt = -w + Σ α_i · P_i_out(m_i) + + 输出： + logits = Q(w) # 从工作空间投影到词汇表 + + 在线学习目标： + L = -log p(x_{t+1} | x_{0:t}) + λ_EWC · Σ_i F_i · (θ_i - θ*_i)² + + EWC 中 F_i 是 Fisher 信息矩阵对角线，衡量参数重要性。 + 重要参数被"锚定"在旧值附近，新知识只能修改不重要的参数。 +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass, field +from collections import deque +import random + +from .core import JSpaceConfig, Expert, JSpaceWorkspace + + +@dataclass +class LanguageConfig(JSpaceConfig): + """语言模型配置，继承 JSpaceConfig""" + vocab_size: int = 100 # 词汇表大小 + embed_dim: int = 16 # token embedding 维度（≠ input_dim，会投影） + # input_dim 仍用 JSpaceConfig 的，作为工作空间接收的输入维度 + + +class JSpaceLanguageModel(nn.Module): + """ + 语言版 JSpace 模型。 + + 流程： + token → embedding → 投影到 input_dim → 喂入 JSpace 动力学 + 工作空间 w → 投影到 vocab_size → logits → 采样 token + + 每步可训练（在线学习）： + forward 后用 cross-entropy loss 更新参数 + """ + + def __init__(self, config: LanguageConfig): + super().__init__() + self.config = config + + # Token embedding + self.embedding = nn.Embedding(config.vocab_size, config.embed_dim) + # 投影 embedding → input_dim（JSpace 期望的输入维度） + self.input_proj = nn.Linear(config.embed_dim, config.input_dim, bias=False) + + # 专家池 + self.experts = nn.ModuleList([ + Expert( + expert_dim=config.expert_dim, + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_wells=config.num_wells, + sparsity=config.jacobian_sparsity, + ) + for _ in range(config.num_experts) + ]) + + # 工作空间 + self.workspace = JSpaceWorkspace( + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_experts=config.num_experts, + ) + + # 输出头：w → logits + self.output_head = nn.Sequential( + nn.Linear(config.workspace_dim, 64), + nn.Tanh(), + nn.Linear(64, config.vocab_size), + ) + + def init_state(self, batch_size: int, device: torch.device) -> dict: + return { + 'w': torch.zeros(batch_size, self.config.workspace_dim, device=device), + 'm': [torch.zeros(batch_size, self.config.expert_dim, device=device) + for _ in range(self.config.num_experts)], + } + + def step(self, state: dict, token_ids: torch.Tensor, + record_trajectory: bool = False) -> tuple[dict, torch.Tensor, torch.Tensor, list]: + """单时间步前向 + + Args: + state: {'w': ..., 'm': [...]} + token_ids: (batch,) token indices + record_trajectory: 是否记录 w 轨迹（J-lens 用） + + Returns: + new_state, logits (batch, vocab_size), alpha (batch, num_experts), + w_trajectory (list of (batch, workspace_dim)) 或空 list + """ + w = state['w'] + ms = state['m'] + cfg = self.config + w_trajectory = [] + + # token → embedding → input projection + emb = self.embedding(token_ids) # (batch, embed_dim) + x = self.input_proj(emb) # (batch, input_dim) + + # ODE 子步积分（对应 Anthropic 论文的"层"） + for substep in range(cfg.ode_steps): + contributions = [] + new_ms = [] + for i, expert in enumerate(self.experts): + m_next, contrib = expert( + ms[i], w, x, + dt=cfg.dt, noise_std=cfg.noise_std, + ) + new_ms.append(m_next) + contributions.append(contrib) + contributions = torch.stack(contributions, dim=1) + + w, alpha = self.workspace( + w, x, contributions, + dt=cfg.dt, tau_w=cfg.tau_w, + ) + ms = new_ms + + if record_trajectory: + w_trajectory.append(w.detach()) + + logits = self.output_head(w) # (batch, vocab_size) + new_state = {'w': w, 'm': ms} + return new_state, logits, alpha, w_trajectory + + def forward(self, token_seqs: torch.Tensor, state: dict | None = None, + record_trajectory: bool = False) -> tuple[torch.Tensor, dict]: + """ + Args: + token_seqs: (batch, T) token indices + record_trajectory: 是否记录 w 轨迹 + + Returns: + logits: (batch, T, vocab_size) + info: {'alpha': (batch, T, num_experts), 'w_norm': (batch, T), + 'w_trajectory': list of (batch, T, workspace_dim) 或空} + """ + batch_size, T = token_seqs.shape + device = token_seqs.device + + if state is None: + state = self.init_state(batch_size, device) + + logits_list = [] + alphas = [] + w_norms = [] + w_traj_per_step = [] # list of (list of substep w) + for t in range(T): + state, logits, alpha, w_traj = self.step( + state, token_seqs[:, t], record_trajectory=record_trajectory + ) + logits_list.append(logits) + alphas.append(alpha) + w_norms.append(state['w'].norm(dim=-1)) + if record_trajectory and w_traj: + w_traj_per_step.append(torch.stack(w_traj, dim=1)) # (batch, n_substeps, workspace_dim) + + logits = torch.stack(logits_list, dim=1) # (batch, T, vocab_size) + info = { + 'alpha': torch.stack(alphas, dim=1), # (batch, T, num_experts) + 'w_norm': torch.stack(w_norms, dim=1), # (batch, T) + } + if record_trajectory and w_traj_per_step: + info['w_trajectory'] = torch.stack(w_traj_per_step, dim=1) # (batch, T, n_substeps, workspace_dim) + return logits, info + + @torch.no_grad() + def generate(self, prompt: list[int], n_new: int = 50, temperature: float = 1.0, + top_k: int = 5) -> list[int]: + """自回归生成 + + Args: + prompt: 起始 token ids + n_new: 生成的新 token 数 + temperature: 采样温度 + top_k: top-k 采样 + """ + self.eval() + device = next(self.parameters()).device + state = self.init_state(1, device) + + # 预热 state with prompt + tokens = list(prompt) + for tok in tokens: + state, _, _, _ = self.step(state, torch.tensor([tok], device=device)) + + # 生成 + generated = [] + for _ in range(n_new): + state, logits, _, _ = self.step(state, torch.tensor([tokens[-1]], device=device)) + logits = logits[0] / max(temperature, 1e-6) + + if top_k > 0: + top_k = min(top_k, logits.size(-1)) + vals, idxs = logits.topk(top_k) + probs = F.softmax(vals, dim=-1) + next_tok = idxs[torch.multinomial(probs, 1)].item() + else: + probs = F.softmax(logits, dim=-1) + next_tok = torch.multinomial(probs, 1).item() + + generated.append(next_tok) + tokens.append(next_tok) + + self.train() + return generated + + +class ExperienceReplay: + """ + 经验回放缓冲区。 + + 存储见过的序列片段，训练时随机采样混入当前 batch。 + 防止灾难性遗忘——旧知识被定期"复习"。 + """ + + def __init__(self, capacity: int = 1000, seq_len: int = 64): + self.capacity = capacity + self.seq_len = seq_len + self.buffer: deque = deque(maxlen=capacity) + + def push(self, token_seq: torch.Tensor): + """push 一个序列 (T,) 或 (batch, T)""" + if token_seq.dim() == 1: + token_seq = token_seq.unsqueeze(0) + for seq in token_seq: + if len(seq) >= self.seq_len: + self.buffer.append(seq.clone()) + + def sample(self, batch_size: int) -> torch.Tensor | None: + """采样 (batch_size, seq_len)""" + if len(self.buffer) < batch_size: + return None + samples = random.sample(list(self.buffer), batch_size) + # 随机裁剪到 seq_len + result = [] + for s in samples: + if len(s) > self.seq_len: + start = random.randint(0, len(s) - self.seq_len - 1) + result.append(s[start:start + self.seq_len]) + else: + result.append(s) + return torch.stack(result) + + +class EWCOptimizer: + """ + Elastic Weight Consolidation 优化器包装。 + + 核心思想：参数 θ 有"重要性" F（Fisher 信息）。 + 重要参数偏离原值 θ* 会被惩罚。 + 新知识只能修改不重要的参数。 + + L_total = L_task + λ · Σ_i F_i · (θ_i - θ*_i)² + + 工作流： + 1. 正常训练一段时间 + 2. 调用 consolidate()：计算 Fisher 信息，锚定当前参数 + 3. 继续训练——loss 中加入 EWC 正则 + 4. 周期性 consolidate（更新锚点和重要性） + """ + + def __init__(self, model: nn.Module, lr: float = 1e-3, + ewc_lambda: float = 1.0, max_grad_norm: float = 1.0): + self.model = model + self.optimizer = torch.optim.Adam(model.parameters(), lr=lr) + self.ewc_lambda = ewc_lambda + self.max_grad_norm = max_grad_norm + + # Fisher 信息和锚定参数 + self.fisher: dict[str, torch.Tensor] = {} + self.anchored_params: dict[str, torch.Tensor] = {} + + def consolidate(self, data_sample: torch.Tensor, n_samples: int = 50): + """计算 Fisher 信息并锚定当前参数 + + Args: + data_sample: (batch, T) 用于计算 Fisher 的数据样本 + n_samples: 采样次数（Fisher 信息的 Monte Carlo 估计） + """ + # 保存当前参数作为锚点 + self.anchored_params = { + name: param.data.clone() + for name, param in self.model.named_parameters() + } + + # 计算 Fisher 信息（对角近似） + fisher = { + name: torch.zeros_like(param) + for name, param in self.model.named_parameters() + } + + self.model.eval() + for _ in range(n_samples): + self.model.zero_grad() + logits, _ = self.model(data_sample) + # 只用 logits[:, :-1] 对应 targets[:, 1:] 的部分 + logits_pred = logits[:, :-1] # (batch, T-1, vocab) + probs = F.softmax(logits_pred, dim=-1) # (batch, T-1, vocab) + # 采样 token 计算 Fisher + sampled_tokens = torch.multinomial( + probs.reshape(-1, probs.size(-1)), 1 + ).view_as(logits_pred[:, :, 0]) # (batch, T-1) + log_probs = F.log_softmax(logits_pred, dim=-1) + loss = -log_probs.gather(-1, sampled_tokens.unsqueeze(-1)).mean() + loss.backward() + + for name, param in self.model.named_parameters(): + if param.grad is not None: + fisher[name] += param.grad.data ** 2 + + # 平均 + for name in fisher: + fisher[name] /= n_samples + + self.fisher = fisher + self.model.zero_grad() + self.model.train() + + def ewc_penalty(self) -> torch.Tensor: + """计算 EWC 正则项""" + if not self.fisher: + return torch.tensor(0.0, device=next(self.model.parameters()).device) + + penalty = 0.0 + for name, param in self.model.named_parameters(): + if name in self.fisher: + penalty = penalty + (self.fisher[name] * (param - self.anchored_params[name]) ** 2).sum() + return penalty + + def step(self, loss: torch.Tensor): + """一步优化：task loss + EWC 正则""" + total_loss = loss + self.ewc_lambda * self.ewc_penalty() + self.optimizer.zero_grad() + total_loss.backward() + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.max_grad_norm) + self.optimizer.step() + return total_loss.item() + + +class ExpertPlasticity: + """ + 专家结构可塑性追踪。 + + 追踪每个专家的"专业化程度"： + - 哪些专家在处理哪些模式 + - 哪些专家"负载过重"（应该分裂或新增） + - 哪些专家"空闲"（可以接收新知识） + + 这不是真正的动态增删专家（实现复杂），而是： + - 统计专家使用率 + - 在路由时给空闲专家加权（鼓励新知识流向空闲专家） + """ + + def __init__(self, num_experts: int, ema_alpha: float = 0.99): + self.num_experts = num_experts + self.ema_alpha = ema_alpha + # 每个专家的使用率（EMA） + self.usage = torch.ones(num_experts) / num_experts + # 每个专家的"领地"——它擅长的 token 分布 + self.expert_specialization: list[dict[int, float]] = [{} for _ in range(num_experts)] + + def update(self, alpha: torch.Tensor, tokens: torch.Tensor): + """更新专家统计 + + Args: + alpha: (batch, T, num_experts) 注意力权重 + tokens: (batch, T) 对应的 token + """ + # 使用率（时间维度平均） + usage_batch = alpha.mean(dim=(0, 1)).detach().cpu() # (num_experts,) + self.usage = self.ema_alpha * self.usage + (1 - self.ema_alpha) * usage_batch + + # 专业化：每个专家最常处理哪些 token + alpha_flat = alpha.reshape(-1, self.num_experts).detach().cpu() # (batch*T, num_experts) + tokens_flat = tokens.reshape(-1).detach().cpu().tolist() + for tok, weights in zip(tokens_flat, alpha_flat): + for i, w in enumerate(weights.tolist()): + if w > 0.1: # 只记录显著激活 + self.expert_specialization[i][tok] = \ + self.expert_specialization[i].get(tok, 0) + w + + def get_diversity_bonus(self) -> torch.Tensor: + """返回多样性奖励——给使用率低的专家加权 + + 在路由注意力上加上这个 bonus，鼓励新知识流向空闲专家 + """ + # 使用率越低，bonus 越高 + bonus = (1.0 - self.usage) / self.usage.clamp(min=1e-4) + bonus = bonus / bonus.sum() # 归一化 + return bonus + + def get_stats(self) -> dict: + """返回可解释性统计""" + return { + 'usage': self.usage.tolist(), + 'top_specialization': [ + sorted(s.items(), key=lambda x: -x[1])[:5] + for s in self.expert_specialization + ], + } diff --git a/jspaceai/modules.py b/jspaceai/modules.py new file mode 100644 index 0000000..d3e416a --- /dev/null +++ b/jspaceai/modules.py @@ -0,0 +1,255 @@ +""" +外挂模块系统 —— 可热插拔的外部能力 + +设计： + - 核心心智不依赖外挂，断开后继续工作 + - 标准接口，任何模块都能插入 + - 运行时热插拔，不需要重启 + - 心智知道外挂状态 +""" +from __future__ import annotations +import torch, torch.nn as nn, numpy as np, time, math +from typing import Optional, Any +from abc import ABC, abstractmethod + + +class ExternalModule(ABC): + @property + @abstractmethod + def name(self) -> str: ... + @abstractmethod + def connect(self) -> bool: ... + @abstractmethod + def disconnect(self): ... + @abstractmethod + def is_connected(self) -> bool: ... + @abstractmethod + def query(self, input_data: Any) -> Any: ... + @abstractmethod + def describe(self) -> str: ... + + +class SmallModelModule(ExternalModule): + """小模型外挂——提供语言/知识能力。可热插拔。""" + + def __init__(self, workspace_dim=64, model_name="placeholder"): + self._name = f"small_model:{model_name}" + self.workspace_dim = workspace_dim + self.model_name = model_name + self._connected = False + self._model = None + self._tokenizer = None + self._proj = None + + @property + def name(self): return self._name + + def connect(self) -> bool: + try: + from transformers import AutoModel, AutoTokenizer + print(f" [外挂] 加载 {self.model_name}...") + self._tokenizer = AutoTokenizer.from_pretrained(self.model_name) + self._model = AutoModel.from_pretrained(self.model_name) + self._model.eval() + hs = self._model.config.hidden_size + self._proj = nn.Linear(hs, self.workspace_dim, bias=False) + self._connected = True + print(f" [外挂] {self.model_name} 已连接 (hidden={hs})") + return True + except ImportError: + print(f" [外挂] transformers 未装，占位模式") + self._connected = True + return True + except Exception as e: + print(f" [外挂] 连接失败: {e}") + return False + + def disconnect(self): + if self._model is not None: + del self._model, self._tokenizer + self._model = self._tokenizer = None + self._proj = None + self._connected = False + print(f" [外挂] {self.name} 已断开") + + def is_connected(self): return self._connected + + def query(self, input_data): + if not self._connected: + return None + if self._model is None: + n = input_data.shape[0] if isinstance(input_data, torch.Tensor) and input_data.dim() > 0 else 1 + return torch.randn(n, self.workspace_dim) + try: + text = str(input_data) if not isinstance(input_data, torch.Tensor) \ + else f"state_{input_data.mean().item():.3f}" + with torch.no_grad(): + inputs = self._tokenizer(text, return_tensors="pt", + truncation=True, max_length=128) + outputs = self._model(**inputs) + hidden = outputs.last_hidden_state.mean(dim=1) + return self._proj(hidden) + except Exception as e: + return None + + def describe(self): + if self._model is None and self._connected: + return "占位模式" + return f"语言模型 {self.model_name}" + + +class KnowledgeBaseModule(ExternalModule): + """知识库外挂——向量检索""" + + def __init__(self, workspace_dim=64, capacity=10000): + self._name = "knowledge_base" + self.workspace_dim = workspace_dim + self.capacity = capacity + self._connected = False + self.entries = [] + self.vectors = None + + @property + def name(self): return self._name + + def connect(self): + self._connected = True + print(f" [外挂] 知识库已连接 ({len(self.entries)} 条)") + return True + + def disconnect(self): + self._connected = False + print(f" [外挂] 知识库已断开（数据保留）") + + def is_connected(self): return self._connected + + def add_entry(self, vector, text, metadata=None): + if len(self.entries) >= self.capacity: + self.entries.pop(0) + self.entries.append({'vector': np.array(vector), 'text': text, + 'metadata': metadata or {}}) + self.vectors = np.array([e['vector'] for e in self.entries]) + + def query(self, input_data): + if not self._connected or not self.entries: + return [] + if not isinstance(input_data, torch.Tensor): + return [] + q = input_data[0].cpu().numpy() if input_data.dim() > 1 else input_data.cpu().numpy() + if self.vectors is None or len(self.vectors) == 0: + return [] + norms = np.linalg.norm(self.vectors, axis=1) * np.linalg.norm(q) + sims = self.vectors @ q / (norms + 1e-8) + top = np.argsort(sims)[-5:][::-1] + return [{'text': self.entries[i]['text'], 'sim': float(sims[i])} for i in top] + + def describe(self): + return f"知识库（{len(self.entries)}/{self.capacity}）" + + +class ToolModule(ExternalModule): + """工具外挂——可调用的外部工具""" + + def __init__(self): + self._name = "tools" + self._connected = False + self.tools = {} + + @property + def name(self): return self._name + + def connect(self): + self._connected = True + self.register('calc', lambda e: eval(e, {'__builtins__': {}}, {'math': math}), "计算") + self.register('time', lambda: time.time(), "时间戳") + print(f" [外挂] 工具箱已连接 ({len(self.tools)} 工具)") + return True + + def disconnect(self): + self._connected = False + self.tools.clear() + print(f" [外挂] 工具箱已断开") + + def is_connected(self): return self._connected + + def register(self, name, func, desc=""): + self.tools[name] = {'func': func, 'desc': desc} + + def query(self, input_data): + if not self._connected: + return None + if isinstance(input_data, dict) and 'tool' in input_data: + tn = input_data['tool'] + args = input_data.get('args', []) + if tn in self.tools: + try: + return self.tools[tn]['func'](*args) + except Exception as e: + return f"错误: {e}" + return None + + def describe(self): + return f"工具箱（{list(self.tools.keys())}）" + + +class ModuleDock: + """外挂坞——USB hub 式管理可热插拔模块。 + + 用法： + dock = ModuleDock() + dock.register('llm', SmallModelModule()) + dock.connect('llm') # 热插上 + result = dock.query('llm', "hello") + dock.disconnect('llm') # 拔掉，心智不停 + """ + + def __init__(self, workspace_dim=64): + self.workspace_dim = workspace_dim + self.slots = {} + self.log = [] + + def register(self, name, module): + self.slots[name] = module + + def connect(self, name) -> bool: + if name not in self.slots: + return False + if self.slots[name].is_connected(): + return True + ok = self.slots[name].connect() + self.log.append({'time': time.time(), 'slot': name, + 'action': 'connect', 'ok': ok}) + return ok + + def disconnect(self, name): + if name not in self.slots: + return + self.slots[name].disconnect() + self.log.append({'time': time.time(), 'slot': name, + 'action': 'disconnect', 'ok': True}) + + def disconnect_all(self): + for n in list(self.slots): + if self.slots[n].is_connected(): + self.disconnect(n) + + def is_connected(self, name) -> bool: + m = self.slots.get(name) + return m.is_connected() if m else False + + def query(self, name, input_data): + """查询某个外挂。未连接返回 None。""" + m = self.slots.get(name) + if m and m.is_connected(): + return m.query(input_data) + return None + + def status(self) -> dict: + """所有外挂状态""" + return {name: { + 'connected': m.is_connected(), + 'description': m.describe(), + } for name, m in self.slots.items()} + + def connected_count(self) -> int: + return sum(1 for m in self.slots.values() if m.is_connected()) diff --git a/jspaceai/multimodal.py b/jspaceai/multimodal.py new file mode 100644 index 0000000..4ddb741 --- /dev/null +++ b/jspaceai/multimodal.py @@ -0,0 +1,464 @@ +""" +多模态感知-行动系统 + +原生支持图片、音频、视频、文本四种模态。 +所有模态编码到统一的 workspace 向量空间，从 workspace 解码到任意模态。 + +设计原则： + 1. 原生多模态——不经过文本中间表示 + 2. 流式处理——支持实时麦克风/摄像头输入 + 3. 闭环——输出反馈到输入 + 4. 在线学习——持续进化 + +模态专家分工（对应 Anthropic 论文的"专家分工"）： + - 视觉专家：处理图片/视频帧 + - 听觉专家：处理音频频谱 + - 语言专家：处理文本 token + - 跨模态专家：对齐不同模态的概念 + +每个专家有自己的内部状态 m_i，通过 J-space 在 workspace w 中广播。 +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass, field +from typing import Optional +import numpy as np + +from .core import JSpaceConfig, Expert, JSpaceWorkspace + + +@dataclass +class MultimodalConfig(JSpaceConfig): + """多模态配置""" + # 文本 + vocab_size: int = 100 + embed_dim: int = 16 + # 视觉 + img_channels: int = 3 + img_size: int = 32 # 编码后的图像尺寸（原始会 resize） + visual_feature_dim: int = 32 + # 音频 + audio_sample_rate: int = 16000 + audio_frame_size: int = 1024 # 每帧采样数 + audio_feature_dim: int = 32 + # 共享 + workspace_dim: int = 64 # 比 unimodal 大，承载多模态 + num_experts: int = 12 # 12 个专家：2视觉 + 2屏幕 + 2听觉 + 2语言 + 2鼠标 + 2跨模态 + expert_dim: int = 24 + # 键盘/鼠标 + keyboard_vocab: int = 128 # 键盘字符词汇表 + mouse_feature_dim: int = 16 # 鼠标特征维度 + + +class VisualEncoder(nn.Module): + """ + 视觉编码器：图片 → workspace 向量 + + 轻量 CNN，不用预训练。从零学。 + 输入：(batch, 3, H, W) + 输出：(batch, input_dim) 投影到 workspace 输入空间 + """ + + def __init__(self, input_dim: int): + super().__init__() + self.conv = nn.Sequential( + nn.Conv2d(3, 16, 3, stride=2, padding=1), # 32→16 + nn.ReLU(), + nn.Conv2d(16, 32, 3, stride=2, padding=1), # 16→8 + nn.ReLU(), + nn.Conv2d(32, 32, 3, stride=2, padding=1), # 8→4 + nn.ReLU(), + nn.AdaptiveAvgPool2d(1), # 全局池化 → (32, 1, 1) + ) + self.proj = nn.Linear(32, input_dim) + + def forward(self, img: torch.Tensor) -> torch.Tensor: + """ + img: (batch, 3, H, W) 或 (batch, T, 3, H, W) 视频序列 + returns: (batch, input_dim) 或 (batch, T, input_dim) + """ + if img.dim() == 5: # 视频序列 + B, T, C, H, W = img.shape + img = img.reshape(B * T, C, H, W) + feat = self.conv(img).squeeze(-1).squeeze(-1) # (B*T, 32) + feat = self.proj(feat) # (B*T, input_dim) + return feat.reshape(B, T, -1) + else: + feat = self.conv(img).squeeze(-1).squeeze(-1) + return self.proj(feat) + + +class AudioEncoder(nn.Module): + """ + 音频编码器：音频波形 → workspace 向量 + + 轻量 1D CNN 处理原始波形。 + 输入：(batch, audio_frame_size) 原始音频采样 + 输出：(batch, input_dim) + """ + + def __init__(self, input_dim: int, audio_frame_size: int = 1024): + super().__init__() + self.conv = nn.Sequential( + nn.Conv1d(1, 16, 64, stride=4, padding=32), # 下采样 + nn.ReLU(), + nn.Conv1d(16, 32, 32, stride=4, padding=16), + nn.ReLU(), + nn.Conv1d(32, 32, 16, stride=2, padding=8), + nn.ReLU(), + nn.AdaptiveAvgPool1d(1), + ) + self.proj = nn.Linear(32, input_dim) + + def forward(self, audio: torch.Tensor) -> torch.Tensor: + """ + audio: (batch, frame_size) 或 (batch, T, frame_size) + returns: (batch, input_dim) 或 (batch, T, input_dim) + """ + if audio.dim() == 3: + B, T, F = audio.shape + audio = audio.reshape(B * T, F) + x = audio.unsqueeze(1) # (B*T, 1, F) + feat = self.conv(x).squeeze(-1) + feat = self.proj(feat) + return feat.reshape(B, T, -1) + else: + x = audio.unsqueeze(1) # (B, 1, F) + feat = self.conv(x).squeeze(-1) + return self.proj(feat) + + +class TextEncoder(nn.Module): + """文本编码器：token → workspace 向量""" + + def __init__(self, vocab_size: int, embed_dim: int, input_dim: int): + super().__init__() + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.proj = nn.Linear(embed_dim, input_dim, bias=False) + + def forward(self, token_ids: torch.Tensor) -> torch.Tensor: + """token_ids: (batch,) 或 (batch, T)""" + if token_ids.dim() == 1: + emb = self.embedding(token_ids) + return self.proj(emb) + else: + emb = self.embedding(token_ids) # (B, T, embed) + return self.proj(emb) + + +class KeyboardEncoder(nn.Module): + """ + 键盘编码器：按键序列 → workspace 向量 + + 把键盘按键序列编码成语义向量。 + 和 TextEncoder 类似，但字符集不同（含特殊键）。 + """ + + def __init__(self, vocab_size: int, embed_dim: int, input_dim: int): + super().__init__() + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.proj = nn.Linear(embed_dim, input_dim, bias=False) + + def forward(self, key_ids: torch.Tensor) -> torch.Tensor: + """key_ids: (batch,) 或 (batch, T)""" + if key_ids.dim() == 1: + return self.proj(self.embedding(key_ids)) + else: + return self.proj(self.embedding(key_ids)) + + +class MouseEncoder(nn.Module): + """ + 鼠标编码器：鼠标位置 + 事件 → workspace 向量 + + 输入：(batch, 4) = (x_norm, y_norm, click_left, click_right) + x_norm, y_norm: 归一化到 [0, 1] 的鼠标坐标 + click_left, click_right: 0 或 1 + """ + + def __init__(self, input_dim: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(4, 32), + nn.ReLU(), + nn.Linear(32, input_dim), + ) + + def forward(self, mouse_data: torch.Tensor) -> torch.Tensor: + """mouse_data: (batch, 4) → (batch, input_dim)""" + return self.net(mouse_data) + + +class VisualDecoder(nn.Module): + """视觉解码器：workspace 向量 → 图片""" + + def __init__(self, workspace_dim: int, img_size: int = 32): + super().__init__() + self.img_size = img_size + self.fc = nn.Linear(workspace_dim, 32 * 4 * 4) + self.deconv = nn.Sequential( + nn.ConvTranspose2d(32, 32, 3, stride=2, padding=1, output_padding=1), # 4→8 + nn.ReLU(), + nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1), # 8→16 + nn.ReLU(), + nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1), # 16→32 + nn.Tanh(), + ) + + def forward(self, w: torch.Tensor) -> torch.Tensor: + """w: (batch, workspace_dim) → (batch, 3, 32, 32)""" + x = self.fc(w).reshape(-1, 32, 4, 4) + return self.deconv(x) + + +class AudioDecoder(nn.Module): + """音频解码器：workspace 向量 → 音频波形""" + + def __init__(self, workspace_dim: int, audio_frame_size: int = 1024): + super().__init__() + self.audio_frame_size = audio_frame_size + self.fc = nn.Linear(workspace_dim, 32 * 64) + self.deconv = nn.Sequential( + nn.ConvTranspose1d(32, 32, 32, stride=4, padding=14, output_padding=2), + nn.ReLU(), + nn.ConvTranspose1d(32, 16, 16, stride=4, padding=6, output_padding=2), + nn.ReLU(), + nn.ConvTranspose1d(16, 1, 64, stride=4, padding=30, output_padding=2), + nn.Tanh(), + ) + + def forward(self, w: torch.Tensor) -> torch.Tensor: + """w: (batch, workspace_dim) → (batch, audio_frame_size)""" + x = self.fc(w).reshape(-1, 32, 64) + audio = self.deconv(x).squeeze(1) # (batch, L) + # 裁剪或填充到目标长度 + if audio.shape[-1] > self.audio_frame_size: + audio = audio[:, :self.audio_frame_size] + else: + audio = F.pad(audio, (0, self.audio_frame_size - audio.shape[-1])) + return audio + + +class TextDecoder(nn.Module): + """文本解码器：workspace 向量 → logits""" + + def __init__(self, workspace_dim: int, vocab_size: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(workspace_dim, 64), + nn.Tanh(), + nn.Linear(64, vocab_size), + ) + + def forward(self, w: torch.Tensor) -> torch.Tensor: + return self.net(w) + + +class MultimodalJSpaceModel(nn.Module): + """ + 多模态 JSpace 模型。 + + 12 个专家分工： + 专家 0,1: 视觉（处理摄像头图片） + 专家 2,3: 屏幕（处理屏幕截图） + 专家 4,5: 听觉（处理音频帧） + 专家 6,7: 语言（处理文本 token） + 专家 8,9: 鼠标（处理鼠标位置/点击） + 专家 10,11: 跨模态（对齐不同模态的概念） + + 所有专家共享 workspace w，通过 J-space 广播。 + + 支持的输入模态： + 'image' - 摄像头图片 (batch, 3, H, W) + 'screen' - 屏幕截图 (batch, 3, H, W)（用视觉编码器） + 'audio' - 音频帧 (batch, frame_size) + 'text' - 文本 token (batch,) + 'keyboard' - 键盘按键 (batch,) 或 (batch, T) + 'mouse' - 鼠标数据 (batch, 4) = (x, y, click_l, click_r) + """ + + def __init__(self, config: MultimodalConfig): + super().__init__() + self.config = config + + # 编码器（各模态 → input_dim） + self.visual_encoder = VisualEncoder(config.input_dim) # 摄像头 + 屏幕 + self.audio_encoder = AudioEncoder(config.input_dim, config.audio_frame_size) + self.text_encoder = TextEncoder(config.vocab_size, config.embed_dim, config.input_dim) + self.keyboard_encoder = KeyboardEncoder(config.keyboard_vocab, config.embed_dim, config.input_dim) + self.mouse_encoder = MouseEncoder(config.input_dim) + + # 解码器（workspace → 各模态） + self.visual_decoder = VisualDecoder(config.workspace_dim, config.img_size) + self.audio_decoder = AudioDecoder(config.workspace_dim, config.audio_frame_size) + self.text_decoder = TextDecoder(config.workspace_dim, config.vocab_size) + + # 专家池 + self.experts = nn.ModuleList([ + Expert( + expert_dim=config.expert_dim, + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_wells=config.num_wells, + sparsity=config.jacobian_sparsity, + ) + for _ in range(config.num_experts) + ]) + + # 工作空间 + self.workspace = JSpaceWorkspace( + workspace_dim=config.workspace_dim, + input_dim=config.input_dim, + num_experts=config.num_experts, + ) + + # 模态类型标记（12 个专家的分工） + self.expert_modality = ( + ['visual'] * 2 + # 0,1: 摄像头 + ['screen'] * 2 + # 2,3: 屏幕 + ['audio'] * 2 + # 4,5: 听觉 + ['text'] * 2 + # 6,7: 语言 + ['mouse'] * 2 + # 8,9: 鼠标 + ['cross'] * 2 # 10,11: 跨模态 + )[:config.num_experts] + + def init_state(self, batch_size: int, device: torch.device) -> dict: + return { + 'w': torch.zeros(batch_size, self.config.workspace_dim, device=device), + 'm': [torch.zeros(batch_size, self.config.expert_dim, device=device) + for _ in range(self.config.num_experts)], + } + + def step(self, state: dict, x: torch.Tensor, + record_trajectory: bool = False) -> tuple[dict, list]: + """ + 单步前向 + + Args: + state: {'w': ..., 'm': [...]} + x: (batch, input_dim) 已编码的输入（任意模态） + record_trajectory: 是否记录 w 轨迹 + + Returns: + new_state, w_trajectory (list) + """ + w = state['w'] + ms = state['m'] + cfg = self.config + w_trajectory = [] + + for substep in range(cfg.ode_steps): + contributions = [] + new_ms = [] + for i, expert in enumerate(self.experts): + m_next, contrib = expert( + ms[i], w, x, + dt=cfg.dt, noise_std=cfg.noise_std, + ) + new_ms.append(m_next) + contributions.append(contrib) + contributions = torch.stack(contributions, dim=1) + + w, alpha = self.workspace( + w, x, contributions, + dt=cfg.dt, tau_w=cfg.tau_w, + ) + ms = new_ms + + if record_trajectory: + w_trajectory.append(w.detach()) + + new_state = {'w': w, 'm': ms} + return new_state, w_trajectory + + def encode_modality(self, modality: str, data: torch.Tensor) -> torch.Tensor: + """编码任意模态到 input_dim + + 支持的模态：image, screen, audio, text, keyboard, mouse + - image/screen: 共用 visual_encoder（都是 RGB 图像） + - keyboard: 键盘按键 id + - mouse: (batch, 4) = (x_norm, y_norm, click_left, click_right) + """ + if modality in ('image', 'screen'): + return self.visual_encoder(data) + elif modality == 'audio': + return self.audio_encoder(data) + elif modality == 'text': + return self.text_encoder(data) + elif modality == 'keyboard': + return self.keyboard_encoder(data) + elif modality == 'mouse': + return self.mouse_encoder(data) + else: + raise ValueError(f"Unknown modality: {modality}") + + def decode_modality(self, modality: str, w: torch.Tensor) -> torch.Tensor: + """从 workspace 解码到任意模态 + + screen 用 visual_decoder（和 image 共享） + keyboard 用 text_decoder（都是 token） + """ + if modality in ('image', 'screen'): + return self.visual_decoder(w) + elif modality == 'audio': + return self.audio_decoder(w) + elif modality in ('text', 'keyboard'): + return self.text_decoder(w) + else: + raise ValueError(f"Unknown modality: {modality}") + + def forward_multimodal(self, modality: str, data: torch.Tensor, + state: dict | None = None, + record_trajectory: bool = False) -> tuple[dict, dict]: + """ + 多模态前向 + + Args: + modality: 'image' / 'audio' / 'text' + data: 模态原始数据 + state: 初始状态 + record_trajectory: 是否记录 w 轨迹 + + Returns: + outputs: {'w': workspace, 'logits/img/audio': 各模态解码} + info: {'alpha', 'w_norm', 'w_trajectory'} + """ + if state is None: + # 先编码确定 batch size + x = self.encode_modality(modality, data) + batch_size = x.shape[0] + device = x.device + state = self.init_state(batch_size, device) + else: + x = self.encode_modality(modality, data) + + # 处理序列或单步 + if x.dim() == 2: # 单步 (batch, input_dim) + state, w_traj = self.step(state, x, record_trajectory) + w = state['w'] + else: # 序列 (batch, T, input_dim) + w_traj_all = [] + for t in range(x.shape[1]): + state, w_traj = self.step(state, x[:, t], record_trajectory) + if record_trajectory: + w_traj_all.append(w_traj) + w = state['w'] + + # 解码到所有模态（workspace 是模态无关的，可解码到任意模态） + outputs = { + 'w': w, + 'image': self.visual_decoder(w), + 'audio': self.audio_decoder(w), + 'text_logits': self.text_decoder(w), + } + + info = { + 'w_norm': w.norm(dim=-1), + } + if record_trajectory and w_traj_all: + info['w_trajectory'] = w_traj_all + + return outputs, info diff --git a/jspaceai/platform.py b/jspaceai/platform.py new file mode 100644 index 0000000..df3308c --- /dev/null +++ b/jspaceai/platform.py @@ -0,0 +1,128 @@ +""" +平台抽象层——跨平台支持（macOS / Windows / Linux） + +统一封装各平台差异：屏幕尺寸、权限检查、路径等。 +所有平台相关代码集中在这里。 +""" +from __future__ import annotations + +import sys +import platform +from dataclasses import dataclass + + +@dataclass +class PlatformInfo: + """平台信息""" + system: str + machine: str + python_version: str + is_macos: bool + is_windows: bool + is_linux: bool + + @classmethod + def detect(cls) -> "PlatformInfo": + s = platform.system() + return cls( + system=s, + machine=platform.machine(), + python_version=sys.version, + is_macos=(s == 'Darwin'), + is_windows=(s == 'Windows'), + is_linux=(s == 'Linux'), + ) + + def __str__(self) -> str: + return f"{self.system}/{self.machine} (Python {self.python_version.split()[0]})" + + +PLATFORM = PlatformInfo.detect() + + +def get_screen_size() -> tuple[int, int]: + """获取主屏幕尺寸（跨平台）""" + try: + import mss + with mss.MSS() as sct: + mon = sct.monitors[1] + return (mon['width'], mon['height']) + except Exception: + pass + + if PLATFORM.is_windows: + try: + import ctypes + user32 = ctypes.windll.user32 + return (user32.GetSystemMetrics(0), user32.GetSystemMetrics(1)) + except Exception: + pass + + return (1920, 1080) + + +def check_camera_permission() -> bool: + """检查摄像头权限""" + if not PLATFORM.is_macos: + return True + try: + import cv2 + cap = cv2.VideoCapture(0) + ok = cap.isOpened() + cap.release() + return ok + except Exception: + return False + + +def check_microphone_permission() -> bool: + """检查麦克风权限""" + try: + import sounddevice as sd + sd.query_devices(kind='input') + return True + except Exception: + return False + + +def check_input_monitoring_permission() -> bool: + """检查键盘/鼠标监听权限""" + if PLATFORM.is_windows: + return True + if PLATFORM.is_linux: + return True + if PLATFORM.is_macos: + try: + import subprocess + result = subprocess.run( + ['osascript', '-e', + 'tell application "System Events" to keystroke ""'], + capture_output=True, timeout=2, + ) + return result.returncode == 0 + except Exception: + return False + return True + + +def print_permission_guide(): + """打印权限配置指南""" + print("\n权限配置指南:") + print("-" * 50) + if PLATFORM.is_macos: + print("macOS 权限设置:") + print(" 1. 摄像头: 系统设置 → 隐私与安全 → 摄像头") + print(" 2. 麦克风: 系统设置 → 隐私与安全 → 麦克风") + print(" 3. 键盘/鼠标监听: 系统设置 → 隐私与安全 → 辅助功能") + print(" 4. 屏幕录制: 系统设置 → 隐私与安全 → 屏幕录制") + elif PLATFORM.is_linux: + print("Linux 权限设置:") + print(" 1. 摄像头: 确保用户在 video 组 (sudo usermod -aG video $USER)") + print(" 2. 音频: 确保用户在 audio 组") + print(" 3. 键盘/鼠标: 需要 X11 或 Wayland 输入权限") + print(" 4. 屏幕捕获: 需要 X11 或安装 grim/slurp (Wayland)") + elif PLATFORM.is_windows: + print("Windows 权限设置:") + print(" 1. 摄像头/麦克风: 设置 → 隐私 → 摄像头/麦克风") + print(" 2. 键盘/鼠标: 通常不需要额外权限") + print("-" * 50) diff --git a/jspaceai/realtime.py b/jspaceai/realtime.py new file mode 100644 index 0000000..6d2f7e4 --- /dev/null +++ b/jspaceai/realtime.py @@ -0,0 +1,374 @@ +""" +实时多模态 I/O 层 + +使用 OpenCV（视频/图像）+ sounddevice（音频 I/O，基于 PortAudio 开源库） ++ macOS afplay（音频播放）实现实时感知-行动循环。 + +提供： + - CameraStream: 实时摄像头采集 + - MicrophoneStream: 实时麦克风采集 + - AudioPlayer: 音频播放（扬声器） + - MultimodalStream: 统一的实时多模态流管理 + +所有流都是非阻塞的，用队列缓冲。可以同时采集视频和音频。 +""" +from __future__ import annotations + +import cv2 +import numpy as np +import sounddevice as sd +import torch +import threading +import queue +import time +from typing import Optional, Callable +from dataclasses import dataclass + + +@dataclass +class Frame: + """统一的帧数据结构""" + timestamp: float + modality: str # 'image' / 'audio' / 'text' + data: np.ndarray # 图像 (H,W,3) / 音频 (frame_size,) / 文本 str + + +class CameraStream: + """ + OpenCV 摄像头流。 + + 非阻塞采集，帧放入队列。 + """ + + def __init__(self, camera_index: int = 0, frame_size: tuple = (32, 32), + fps: int = 10): + self.camera_index = camera_index + self.target_size = frame_size # resize 到这个尺寸（给模型用） + self.fps = fps + self.frame_queue: queue.Queue = queue.Queue(maxsize=30) + self.running = False + self.thread: Optional[threading.Thread] = None + self.cap: Optional[cv2.VideoCapture] = None + + def start(self): + """启动采集""" + self.cap = cv2.VideoCapture(self.camera_index) + if not self.cap.isOpened(): + raise RuntimeError(f"无法打开摄像头 {self.camera_index}（需要 macOS 权限）") + + self.running = True + self.thread = threading.Thread(target=self._capture_loop, daemon=True) + self.thread.start() + + def _capture_loop(self): + frame_interval = 1.0 / self.fps + while self.running: + ret, frame = self.cap.read() + if not ret: + continue + + # BGR → RGB + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + # resize 到目标尺寸 + frame_resized = cv2.resize(frame_rgb, self.target_size) + + try: + self.frame_queue.put_nowait( + Frame(time.time(), 'image', frame_resized) + ) + except queue.Full: + pass # 丢帧 + + time.sleep(frame_interval) + + def get_frame(self) -> Optional[Frame]: + """获取最新帧（非阻塞）""" + try: + return self.frame_queue.get_nowait() + except queue.Empty: + return None + + def stop(self): + self.running = False + if self.thread: + self.thread.join(timeout=1.0) + if self.cap: + self.cap.release() + + +class MicrophoneStream: + """ + sounddevice 麦克风流。 + + 持续采集音频，分帧放入队列。 + """ + + def __init__(self, sample_rate: int = 16000, frame_size: int = 1024, + channels: int = 1): + self.sample_rate = sample_rate + self.frame_size = frame_size + self.channels = channels + self.frame_queue: queue.Queue = queue.Queue(maxsize=100) + self.running = False + self.stream: Optional[sd.InputStream] = None + + def start(self): + """启动采集""" + self.running = True + self.stream = sd.InputStream( + samplerate=self.sample_rate, + blocksize=self.frame_size, + channels=self.channels, + dtype='float32', + callback=self._audio_callback, + ) + self.stream.start() + + def _audio_callback(self, indata, frames, time_info, status): + if not self.running: + return + # indata: (frame_size, channels) + audio = indata[:, 0] if self.channels > 1 else indata.flatten() + try: + self.frame_queue.put_nowait( + Frame(time.time(), 'audio', audio.copy()) + ) + except queue.Full: + pass + + def get_frame(self) -> Optional[Frame]: + try: + return self.frame_queue.get_nowait() + except queue.Empty: + return None + + def stop(self): + self.running = False + if self.stream: + self.stream.stop() + self.stream.close() + + +class AudioPlayer: + """ + 音频播放器。用 sounddevice 输出到扬声器。 + """ + + def __init__(self, sample_rate: int = 16000): + self.sample_rate = sample_rate + self.playing = False + + def play(self, audio: np.ndarray, blocking: bool = False): + """播放音频波形 + + Args: + audio: (n_samples,) float32, 范围 [-1, 1] + blocking: 是否阻塞等待播放完成 + """ + audio = np.clip(audio, -1, 1).astype(np.float32) + sd.play(audio, self.sample_rate) + if blocking: + sd.wait() + + def stop(self): + sd.stop() + + +class MultimodalStream: + """ + 统一管理摄像头 + 麦克风 + 扬声器的实时流。 + + 提供统一的接口： + - start(): 启动所有采集 + - get_frames(): 获取当前所有可用帧（按时间戳对齐） + - respond_audio(): 播放音频回应 + - stop(): 停止所有 + + 这是感知-行动循环的物理层。 + """ + + def __init__(self, use_camera: bool = True, use_mic: bool = True, + img_size: tuple = (32, 32), sample_rate: int = 16000, + audio_frame_size: int = 1024): + self.use_camera = use_camera + self.use_mic = use_mic + + self.camera = CameraStream(frame_size=img_size) if use_camera else None + self.mic = MicrophoneStream( + sample_rate=sample_rate, frame_size=audio_frame_size + ) if use_mic else None + self.player = AudioPlayer(sample_rate=sample_rate) + + def start(self): + """启动所有采集""" + if self.camera: + try: + self.camera.start() + print(" 摄像头已启动") + except RuntimeError as e: + print(f" 摄像头启动失败: {e}") + self.camera = None + + if self.mic: + try: + self.mic.start() + print(" 麦克风已启动") + except Exception as e: + print(f" 麦克风启动失败: {e}") + self.mic = None + + def get_latest_frames(self) -> dict[str, Optional[Frame]]: + """获取最新的各模态帧""" + result = { + 'image': None, + 'audio': None, + } + if self.camera: + # 清空队列，只保留最新 + latest = None + while True: + f = self.camera.get_frame() + if f is None: + break + latest = f + result['image'] = latest + + if self.mic: + latest = None + while True: + f = self.mic.get_frame() + if f is None: + break + latest = f + result['audio'] = latest + + return result + + def play_audio(self, audio: np.ndarray, blocking: bool = False): + """播放音频""" + self.player.play(audio, blocking) + + def stop(self): + if self.camera: + self.camera.stop() + if self.mic: + self.mic.stop() + self.player.stop() + + +class SensoryMotorLoop: + """ + 感知-行动循环：传感器 → 模型 → 执行器 → 传感器... + + 这是自进化语言网络的核心运行时。 + + 循环： + 1. 从传感器读帧（摄像头/麦克风） + 2. 编码到 workspace 输入 + 3. 模型 forward 一步（workspace 更新） + 4. 从 workspace 解码输出（音频/图像/文本） + 5. 执行器输出（扬声器播放/显示图像） + 6. 输出反馈到输入（自监督：预测自己的输出） + 7. 在线学习（更新参数） + + 这个循环永不停止——模型持续感知、思考、行动、学习。 + """ + + def __init__(self, model, stream: MultimodalStream, device: str = 'cpu'): + self.model = model + self.stream = stream + self.device = device + self.state = model.init_state(1, torch.device(device)) + self.running = False + self.step_count = 0 + + # 自监督学习记录 + self.last_output = None # 上一步的输出（用于反馈） + self.online_loss_history = [] + + def step_once(self) -> dict: + """执行一步感知-行动循环 + + Returns: + info: 包含输入模态、输出、loss 等 + """ + # 1. 读取传感器 + frames = self.stream.get_latest_frames() + + # 选择活跃模态（优先音频，其次图像） + modality = None + input_data = None + if frames['audio'] is not None: + modality = 'audio' + input_data = torch.tensor( + frames['audio'].data, dtype=torch.float32 + ).unsqueeze(0).to(self.device) + elif frames['image'] is not None: + modality = 'image' + img = frames['image'].data # (H, W, 3) uint8 + img_tensor = torch.tensor( + img, dtype=torch.float32 + ).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0 # 归一化到 [-1, 1] + input_data = img_tensor + + if modality is None: + return {'status': 'no_input'} + + # 2. 编码 + forward + with torch.no_grad(): + outputs, info = self.model.forward_multimodal( + modality, input_data, state=self.state + ) + self.state = self.model.init_state(1, torch.device(self.device)) + # 保留 state（forward_multimodal 修改了 state） + # 实际上 forward_multimodal 内部更新了 state，但接口设计问题 + # 这里简化：每次用上一步的 state + + # 3. 输出 + info['modality'] = modality + info['step'] = self.step_count + self.step_count += 1 + + # 4. 自监督：记录输出用于下一步反馈 + self.last_output = outputs + + return info + + def run(self, n_steps: int = 100, interval: float = 0.1, + on_step: Optional[Callable] = None): + """运行感知-行动循环 + + Args: + n_steps: 总步数（None = 无限） + interval: 每步间隔（秒） + on_step: 每步回调 + """ + self.running = True + self.stream.start() + print(f"感知-行动循环启动，{n_steps} 步，间隔 {interval}s") + print("（首次运行 macOS 会请求摄像头/麦克风权限）") + + try: + step = 0 + while self.running and (n_steps is None or step < n_steps): + info = self.step_once() + + if info.get('status') != 'no_input': + if on_step: + on_step(info) + + # 如果有音频输出，播放 + if self.last_output and 'audio' in self.last_output: + audio_out = self.last_output['audio'][0].cpu().numpy() + # 每 10 步播放一次（避免太频繁） + if step % 10 == 0: + self.stream.play_audio(audio_out) + + step += 1 + time.sleep(interval) + except KeyboardInterrupt: + print("\n用户中断") + finally: + self.running = False + self.stream.stop() + print(f"循环结束，共 {step} 步") diff --git a/jspaceai/task.py b/jspaceai/task.py new file mode 100644 index 0000000..af69fd6 --- /dev/null +++ b/jspaceai/task.py @@ -0,0 +1,101 @@ +""" +玩具任务：连续时间序列预测 + +生成多模态时间序列——不同时段由不同的潜在规则驱动： + 时段 A（t mod 4 ∈ [0,1)）：正弦波 + 慢漂移 + 时段 B（t mod 4 ∈ [1,2)）：快速振荡 + 时段 C（t mod 4 ∈ [2,3)）：脉冲信号 + 时段 D（t mod 4 ∈ [3,4)）：组合信号 + +任务：给定 x_t，预测 x_{t+1}。 + +设计意图： + - 不同时段需要不同的"专家"处理——验证工作空间能否学到分工 + - 模式切换时需要工作空间广播给正确的专家——验证 J-space 路由 + - 扁平 MLP 在模式切换时会糊掉，工作空间架构应该更鲁棒 +""" +from __future__ import annotations + +import numpy as np +import torch +from torch.utils.data import Dataset + + +class ContinuousSequenceTask: + """生成连续多模态时间序列""" + + def __init__(self, input_dim: int = 8, seq_len: int = 64, seed: int = 42): + self.input_dim = input_dim + self.seq_len = seq_len + self.rng = np.random.default_rng(seed) + + # 每个时段有自己的参数 + self.phase_len = 16 # 每个模式持续 16 步 + self.phase_params = self._gen_phase_params() + + def _gen_phase_params(self) -> list[dict]: + """4 个时段，每段不同参数""" + params = [] + for i in range(4): + params.append({ + 'freq': self.rng.uniform(0.1, 0.5), # 振荡频率 + 'phase': self.rng.uniform(0, 2 * np.pi), # 相位 + 'drift': self.rng.uniform(-0.02, 0.02), # 慢漂移 + 'amp': self.rng.uniform(0.3, 1.0), # 振幅 + 'noise': self.rng.uniform(0.01, 0.05), # 噪声 + }) + return params + + def generate_sequence(self, n_steps: int) -> np.ndarray: + """生成长度为 n_steps 的序列""" + x = np.zeros((n_steps, self.input_dim), dtype=np.float32) + state = np.zeros(self.input_dim, dtype=np.float32) + + for t in range(n_steps): + phase_idx = (t // self.phase_len) % 4 + p = self.phase_params[phase_idx] + + # 每个维度独立演化，但共享相位 + for d in range(self.input_dim): + d_offset = d * 0.3 + if phase_idx == 0: + # 正弦 + 漂移 + state[d] = p['amp'] * np.sin(p['freq'] * t + p['phase'] + d_offset) + p['drift'] * t + elif phase_idx == 1: + # 快速振荡 + state[d] = p['amp'] * np.sin(p['freq'] * 3 * t + p['phase'] + d_offset) + elif phase_idx == 2: + # 脉冲 + pulse = 1.0 if (t % 5 == 0) else 0.0 + state[d] = p['amp'] * pulse + 0.1 * np.sin(p['phase'] + d_offset) + else: + # 组合 + state[d] = (p['amp'] * 0.5 * np.sin(p['freq'] * t + p['phase'] + d_offset) + + p['amp'] * 0.3 * np.cos(p['freq'] * 2 * t + d_offset)) + + state[d] += self.rng.normal(0, p['noise']) + + x[t] = state + + # 归一化到 [-1, 1] + x = np.tanh(x) + return x + + def generate_batch(self, batch_size: int, n_steps: int | None = None) -> torch.Tensor: + """生成一个 batch 的序列""" + n_steps = n_steps or self.seq_len + batch = np.stack([self.generate_sequence(n_steps) for _ in range(batch_size)]) + return torch.from_numpy(batch) + + +class SequenceDataset(Dataset): + """用于 DataLoader 的数据集""" + + def __init__(self, task: ContinuousSequenceTask, n_samples: int, seq_len: int): + self.data = [task.generate_sequence(seq_len) for _ in range(n_samples)] + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return torch.from_numpy(self.data[idx]) diff --git a/jspaceai/trainer.py b/jspaceai/trainer.py new file mode 100644 index 0000000..f3c1d3a --- /dev/null +++ b/jspaceai/trainer.py @@ -0,0 +1,92 @@ +""" +训练器：预测学习目标 + 对比训练 + +学习目标（第一性原理）： + 智慧系统的核心是"内部建模世界"——即预测下一时刻世界状态。 + 所以 loss = ||x_{t+1} - pred_t||² + + pred_t = model(x_{0:t}) + +这是自监督的——不需要标签，只需要时间序列本身。 +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import numpy as np +from typing import Optional +from .task import ContinuousSequenceTask + + +class Trainer: + """通用训练器，适用于 JSpaceModel 和 FlatBaseline""" + + def __init__(self, model: nn.Module, task: ContinuousSequenceTask, + lr: float = 1e-3, device: str = 'cpu'): + self.model = model.to(device) + self.task = task + self.device = device + self.optimizer = torch.optim.Adam(model.parameters(), lr=lr) + self.loss_fn = nn.MSELoss() + self.history: list[float] = [] + + def train_step(self, xs: torch.Tensor) -> float: + """ + xs: (batch, T, input_dim) + returns: loss value + """ + xs = xs.to(self.device) + # 预测目标：x_{t+1} = xs[:, 1:], 输入 xs[:, :-1] + # 模型输出 preds[:, t] 应该预测 xs[:, t+1] + preds, _ = self.model(xs) + # preds: (batch, T, input_dim) —— preds[:, t] 是从 xs[:, :t+1] 预测的下一时刻 + # 对齐：preds[:, :-1] 预测 xs[:, 1:] + pred = preds[:, :-1] + target = xs[:, 1:] + + loss = self.loss_fn(pred, target) + + self.optimizer.zero_grad() + loss.backward() + # 梯度裁剪（ODE 训练容易爆炸） + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + self.optimizer.step() + + return loss.item() + + def evaluate(self, xs: torch.Tensor) -> float: + """评估，不更新参数""" + self.model.eval() + with torch.no_grad(): + xs = xs.to(self.device) + preds, _ = self.model(xs) + pred = preds[:, :-1] + target = xs[:, 1:] + loss = self.loss_fn(pred, target).item() + self.model.train() + return loss + + def train(self, n_steps: int = 500, batch_size: int = 32, + eval_interval: int = 50, verbose: bool = True) -> list[float]: + """完整训练循环""" + self.model.train() + for step in range(n_steps): + xs = self.task.generate_batch(batch_size) + loss = self.train_step(xs) + self.history.append(loss) + + if verbose and (step + 1) % eval_interval == 0: + eval_xs = self.task.generate_batch(64) + eval_loss = self.evaluate(eval_xs) + print(f" step {step+1:4d} | train_loss {loss:.4f} | eval_loss {eval_loss:.4f}") + + return self.history + + def get_attention(self, xs: torch.Tensor) -> Optional[torch.Tensor]: + """获取注意力权重（仅 JSpaceModel 有，用于可解释性）""" + self.model.eval() + with torch.no_grad(): + xs = xs.to(self.device) + _, info = self.model(xs) + self.model.train() + return info.get('alpha', None) diff --git a/main.py b/main.py new file mode 100644 index 0000000..03e62ba --- /dev/null +++ b/main.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +""" +JspaceAI —— 主程序入口 + +对比实验： + 1. JSpaceModel（工作空间 + J-space 广播 + ODE 动力学） + 2. FlatBaseline（同样参数量的扁平 MLP） + +任务：多模态连续时间序列预测（4 种模式切换） + +输出： + - 训练/评估 loss 曲线对比 + - 注意力热力图（哪个专家在哪个时段被激活） + - 工作空间 ||w|| 演化曲线（输出门控信号） + - 预测 vs 真实序列对比 + +运行： + python main.py + python main.py --steps 1000 --device cuda +""" +from __future__ import annotations + +import argparse +import torch +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path + +from jspaceai import ( + JSpaceConfig, JSpaceModel, FlatBaseline, + ContinuousSequenceTask, Trainer, +) + + +def count_params(model: torch.nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def run_experiment(n_steps: int, device: str, outdir: Path): + outdir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(42) + np.random.seed(42) + + # 任务 + task = ContinuousSequenceTask(input_dim=8, seq_len=64, seed=42) + + # 模型 1: JSpaceModel + config = JSpaceConfig( + input_dim=8, + workspace_dim=32, + expert_dim=16, + num_experts=5, + num_wells=4, + ode_steps=4, + dt=0.1, + tau_w=0.3, + output_threshold=0.5, + jacobian_sparsity=8, + noise_std=0.01, + ) + jspace_model = JSpaceModel(config) + + # 模型 2: FlatBaseline（参数量对齐） + # JSpaceModel 大约的参数量： + # 专家: 5 × (4×16 + 4 + 16×8 + 16×32 + 16×32) ≈ 5 × 1108 = 5540 + # 工作空间: (8+32)×32 + 32 + 32×32 + 32 ≈ 2144 + # 预测头: 32×32 + 32 + 32×8 + 8 ≈ 1384 + # 总计 ≈ 9000 + flat_model = FlatBaseline(input_dim=8, hidden_dim=90, num_layers=2) + + print("=" * 60) + print("JspaceAI 对比实验") + print("=" * 60) + print(f"任务: 多模态连续序列预测 (input_dim=8, seq_len=64)") + print(f"训练步数: {n_steps}") + print(f"设备: {device}") + print() + print(f"JSpaceModel 参数量: {count_params(jspace_model):,}") + print(f"FlatBaseline 参数量: {count_params(flat_model):,}") + print() + + # 训练 + print("-" * 60) + print("训练 JSpaceModel...") + trainer_js = Trainer(jspace_model, task, lr=1e-3, device=device) + history_js = trainer_js.train(n_steps=n_steps, batch_size=32, eval_interval=n_steps // 10) + + print() + print("-" * 60) + print("训练 FlatBaseline...") + trainer_flat = Trainer(flat_model, task, lr=1e-3, device=device) + history_flat = trainer_flat.train(n_steps=n_steps, batch_size=32, eval_interval=n_steps // 10) + + # 评估 + print() + print("=" * 60) + print("最终评估") + print("=" * 60) + eval_xs = task.generate_batch(128) + eval_loss_js = trainer_js.evaluate(eval_xs) + eval_loss_flat = trainer_flat.evaluate(eval_xs) + print(f"JSpaceModel eval MSE: {eval_loss_js:.6f}") + print(f"FlatBaseline eval MSE: {eval_loss_flat:.6f}") + winner = "JSpaceModel" if eval_loss_js < eval_loss_flat else "FlatBaseline" + improvement = abs(eval_loss_js - eval_loss_flat) / max(eval_loss_js, eval_loss_flat) * 100 + print(f"胜者: {winner} (相对优势 {improvement:.1f}%)") + + # 可视化 + print() + print("生成可视化...") + + # 1. Loss 曲线对比 + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + + ax = axes[0, 0] + ax.plot(history_js, label='JSpaceModel', alpha=0.7, linewidth=0.8) + ax.plot(history_flat, label='FlatBaseline', alpha=0.7, linewidth=0.8) + # 平滑曲线 + if len(history_js) > 20: + smooth_js = np.convolve(history_js, np.ones(20)/20, mode='valid') + smooth_flat = np.convolve(history_flat, np.ones(20)/20, mode='valid') + ax.plot(smooth_js, label='JSpace (smoothed)', linewidth=2) + ax.plot(smooth_flat, label='Flat (smoothed)', linewidth=2) + ax.set_xlabel('Step') + ax.set_ylabel('MSE Loss') + ax.set_title('Training Loss') + ax.legend() + ax.set_yscale('log') + ax.grid(True, alpha=0.3) + + # 2. 注意力热力图 + test_xs = task.generate_batch(1) # 单条序列 + alpha = trainer_js.get_attention(test_xs) # (1, T, num_experts) + if alpha is not None: + ax = axes[0, 1] + alpha_np = alpha[0].cpu().numpy() # (T, num_experts) + im = ax.imshow(alpha_np.T, aspect='auto', cmap='hot', interpolation='nearest') + ax.set_xlabel('Time step') + ax.set_ylabel('Expert index') + ax.set_title('J-space Attention (which expert is active)') + plt.colorbar(im, ax=ax) + # 标注时段边界 + for i in range(1, 64 // 16): + ax.axvline(x=i * 16, color='cyan', linestyle='--', alpha=0.7) + ax.set_xticks(range(0, 64, 16)) + + # 3. 工作空间 ||w|| 演化 + jspace_model.eval() + with torch.no_grad(): + xs = test_xs.to(device) + _, info = jspace_model(xs) + w_norm = info['w_norm'][0].cpu().numpy() + + ax = axes[1, 0] + ax.plot(w_norm, label='||w||', linewidth=2) + ax.axhline(y=config.output_threshold, color='r', linestyle='--', + label=f'threshold={config.output_threshold}', alpha=0.7) + ax.set_xlabel('Time step') + ax.set_ylabel('||w||') + ax.set_title('Workspace norm (output gating signal)') + ax.legend() + ax.grid(True, alpha=0.3) + for i in range(1, 64 // 16): + ax.axvline(x=i * 16, color='gray', linestyle='--', alpha=0.5) + + # 4. 预测 vs 真实（取第 0 维） + ax = axes[1, 1] + with torch.no_grad(): + xs = test_xs.to(device) + preds_js, _ = jspace_model(xs) + preds_flat, _ = flat_model(xs) + true_seq = test_xs[0, 1:, 0].numpy() + pred_js = preds_js[0, :-1, 0].cpu().numpy() + pred_flat = preds_flat[0, :-1, 0].cpu().numpy() + + ax.plot(true_seq, label='True', linewidth=2, color='black') + ax.plot(pred_js, label='JSpace', alpha=0.8) + ax.plot(pred_flat, label='Flat', alpha=0.8) + ax.set_xlabel('Time step') + ax.set_ylabel('Value (dim 0)') + ax.set_title('Prediction vs True (dim 0)') + ax.legend() + ax.grid(True, alpha=0.3) + for i in range(1, 64 // 16): + ax.axvline(x=i * 16, color='gray', linestyle='--', alpha=0.3) + + plt.tight_layout() + fig_path = outdir / 'experiment.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"可视化已保存: {fig_path}") + plt.close() + + # 保存模型 + torch.save({ + 'jspace_model': jspace_model.state_dict(), + 'flat_model': flat_model.state_dict(), + 'config': config, + 'eval_loss_js': eval_loss_js, + 'eval_loss_flat': eval_loss_flat, + }, outdir / 'models.pt') + + print() + print("=" * 60) + print("实验完成") + print("=" * 60) + print(f"JSpaceModel MSE: {eval_loss_js:.6f}") + print(f"FlatBaseline MSE: {eval_loss_flat:.6f}") + print(f"胜者: {winner}") + print() + print("关键观察点（看 experiment.png）:") + print(" 1. Loss 曲线: JSpace 是否收敛更快/更低？") + print(" 2. 注意力热力图: 不同时段是否激活不同专家？") + print(" 3. ||w|| 曲线: 模式切换时是否有明显尖峰？") + print(" 4. 预测对比: 模式切换处哪个模型更鲁棒？") + + +def main(): + parser = argparse.ArgumentParser(description='JspaceAI 对比实验') + parser.add_argument('--steps', type=int, default=500, + help='训练步数 (default: 500)') + parser.add_argument('--device', type=str, default='cpu', + help='设备 (cpu / cuda / mps)') + parser.add_argument('--outdir', type=str, default='outputs', + help='输出目录') + args = parser.parse_args() + + # 自动检测设备 + if args.device == 'auto': + if torch.cuda.is_available(): + device = 'cuda' + elif torch.backends.mps.is_available(): + device = 'mps' + else: + device = 'cpu' + else: + device = args.device + + run_experiment( + n_steps=args.steps, + device=device, + outdir=Path(args.outdir), + ) + + +if __name__ == '__main__': + main() diff --git a/main_autonomous.py b/main_autonomous.py new file mode 100644 index 0000000..91a7fe8 --- /dev/null +++ b/main_autonomous.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""JspaceAI 自主心智 demo —— 永不停止的自主进化 + +运行： + python main_autonomous.py --steps 100 + # 再次运行会从上次状态继续 + python main_autonomous.py --steps 100 +""" +from __future__ import annotations +import argparse, torch, numpy as np, matplotlib.pyplot as plt +from jspaceai import ( + MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent, + AutonomousMind, PLATFORM, +) + + +def get_config(): + return MultimodalConfig( + vocab_size=50, embed_dim=16, input_dim=8, workspace_dim=64, + expert_dim=24, num_experts=12, num_wells=4, ode_steps=3, + dt=0.1, tau_w=0.3, jacobian_sparsity=16, noise_std=0.01, + img_size=32, audio_frame_size=1024, keyboard_vocab=128, + ) + + +def main(): + p = argparse.ArgumentParser(description='JspaceAI 自主心智') + p.add_argument('--steps', type=int, default=100) + p.add_argument('--device', default='cpu') + p.add_argument('--unsafe', action='store_true') + args = p.parse_args() + + dev = args.device + if dev == 'auto': + dev = 'cuda' if torch.cuda.is_available() else ( + 'mps' if torch.backends.mps.is_available() else 'cpu') + + print(f"平台: {PLATFORM}") + config = get_config() + model = MultimodalJSpaceModel(config).to(dev) + model.eval() + + agent = EmbodiedAgent( + model, device=dev, + enable_mouse_output=args.unsafe, + enable_keyboard_output=args.unsafe, + enable_audio_output=True, enable_screen_output=False, + ) + mind = AutonomousMind(agent, save_dir='outputs/mind', device=dev) + + print("\n" + "=" * 60) + print("自主心智 - 永不停止的进化") + print("=" * 60) + print("1. 好奇心驱动 2. 状态持久化 3. 自我模型 4. 元学习") + print(f"运行 {args.steps} 步（Ctrl+C 中断，状态自动保存）\n") + + log = [] + + def on_step(info): + log.append(info) + 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']}") + + mind.run(n_steps=args.steps, interval=0.2, save_every=30, on_step=on_step) + print("\n" + mind.introspect()) + + if log: + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + steps = [s['step'] for s in log] + + axes[0, 0].plot(steps, [s['w_norm'] for s in log], 'b-') + axes[0, 0].set_title('Workspace ||w||'); axes[0, 0].grid(True, alpha=0.3) + + axes[0, 1].plot(steps, [s['curiosity'] for s in log], 'r-', label='curiosity') + axes[0, 1].plot(steps, [s['success'] for s in log], 'g-', label='success') + axes[0, 1].legend(); axes[0, 1].set_title('Curiosity & Success') + axes[0, 1].grid(True, alpha=0.3) + + fc = log[-1]['self_confidence'] + axes[1, 0].barh(list(fc.keys()), list(fc.values()), + color=plt.cm.RdYlGn(list(fc.values()))) + axes[1, 0].set_xlim(0, 1); axes[1, 0].set_title('Self Model') + + axes[1, 1].plot(steps, [s['world_loss'] for s in log], 'orange') + axes[1, 1].set_title('World Model Loss'); axes[1, 1].grid(True, alpha=0.3) + + plt.tight_layout() + Path('outputs').mkdir(exist_ok=True) + plt.savefig('outputs/autonomous_mind.png', dpi=150, bbox_inches='tight') + print(f"\n可视化: outputs/autonomous_mind.png") + + +if __name__ == '__main__': + from pathlib import Path + main() diff --git a/main_embodied.py b/main_embodied.py new file mode 100644 index 0000000..4d39d25 --- /dev/null +++ b/main_embodied.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +"""JspaceAI 具身 Agent - 完整神经系统 + 跨平台""" +from __future__ import annotations +import argparse, torch, numpy as np, matplotlib.pyplot as plt +from pathlib import Path +import time +from jspaceai import ( + MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent, + PLATFORM, get_screen_size, print_permission_guide, + check_camera_permission, check_microphone_permission, + check_input_monitoring_permission, +) + + +def get_config(): + return MultimodalConfig( + vocab_size=50, embed_dim=16, input_dim=8, workspace_dim=64, + expert_dim=24, num_experts=12, num_wells=4, ode_steps=3, + dt=0.1, tau_w=0.3, jacobian_sparsity=16, noise_std=0.01, + img_size=32, audio_frame_size=1024, keyboard_vocab=128, + ) + + +def test_subsystems(): + print("=" * 60) + print(f"具身 Agent 子系统测试 | 平台: {PLATFORM}") + print("=" * 60) + + print("\n1. 权限检查:") + checks = { + '摄像头': check_camera_permission(), + '麦克风': check_microphone_permission(), + '键盘/鼠标监听': check_input_monitoring_permission(), + } + for name, ok in checks.items(): + print(f" {name}: {'OK' if ok else '需要权限'}") + if not all(checks.values()): + print_permission_guide() + + sw, sh = get_screen_size() + print(f" 屏幕尺寸: {sw}x{sh}") + + print("\n2. 模型 + Agent 初始化:") + config = get_config() + model = MultimodalJSpaceModel(config) + print(f" 模型参数: {sum(p.numel() for p in model.parameters()):,}") + print(f" 专家分工: {model.expert_modality}") + + agent = EmbodiedAgent( + model, device='cpu', + enable_mouse_output=False, enable_keyboard_output=False, + enable_audio_output=True, enable_screen_output=False, + ) + print(f" 小脑参数: {sum(p.numel() for p in agent.cerebellum.parameters()):,}") + print(f" 海马体容量: {agent.hippocampus.capacity}") + print(f" 反射弧数: {len(agent.cns.reflexes)}") + + print("\n3. 单步循环测试（2秒采集）:") + agent.senses.start() + time.sleep(2) + for i 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"executed {info['action']['executed']} | mem {info['memories_count']}") + time.sleep(0.5) + agent.senses.stop() + agent.audio_actuator.stop() + + print("\n4. 海马体回忆测试:") + if agent.hippocampus and agent.hippocampus.size() > 0: + w = agent.state['w'][0].cpu().numpy() + for i, m in enumerate(agent.hippocampus.recall(w, top_k=3)): + print(f" 记忆 {i}: sim={m['similarity']:.3f} ctx={m['context']}") + + print("\n所有子系统测试完成") + + +def live_embodied(n_steps, device, safe_mode=False): + print(f"平台: {PLATFORM}") + sw, sh = get_screen_size() + print(f"屏幕: {sw}x{sh}") + if safe_mode: + print("\n安全模式：不执行鼠标/键盘动作") + print_permission_guide() + + config = get_config() + model = MultimodalJSpaceModel(config).to(device) + mp = Path('outputs/multimodal_model.pt') + if mp.exists(): + try: + ckpt = torch.load(mp, map_location=device, weights_only=False) + model.load_state_dict(ckpt['model'], strict=False) + print(f"已加载模型: {mp}") + except Exception: + print("模型加载失败，随机初始化") + else: + print("未找到训练模型，随机初始化") + model.eval() + + agent = EmbodiedAgent( + model, device=device, + enable_mouse_output=not safe_mode, + enable_keyboard_output=not safe_mode, + enable_audio_output=True, + enable_screen_output=not safe_mode, + risk_threshold=0.5 if safe_mode else 0.3, + ) + + log = [] + + def on_step(info): + log.append(info) + if info['step'] % 5 == 0: + a = info['action'] + print(f" step {info['step']:3d} | mod {info['modality']:8s} | " + f"||w|| {info['w_norm']:.3f} | action {a['action_idx']} " + f"str {a['action_strength']:.2f} " + f"exec {'Y' if a['executed'] else 'N'} | " + f"mem {info['memories_count']}") + + print(f"\n运行 {n_steps} 步具身循环...") + agent.run(n_steps=n_steps, interval=0.2, on_step=on_step) + + print("\n" + "=" * 60) + print("总结") + print("=" * 60) + print(f"总步数: {len(log)}") + if log: + wn = [s['w_norm'] for s in log] + print(f"||w||: [{min(wn):.3f}, {max(wn):.3f}] mean={np.mean(wn):.3f}") + mc = {} + for s in log: + mc[s['modality']] = mc.get(s['modality'], 0) + 1 + print("\n模态分布:") + for m, c in sorted(mc.items(), key=lambda x: -x[1]): + print(f" {m:10s}: {c:3d} ({c/len(log)*100:.0f}%)") + ex = sum(1 for s in log if s['action']['executed']) + print(f"\n动作执行: {ex}/{len(log)} ({ex/len(log)*100:.0f}%)") + print("\n基底神经节:") + for i, c in enumerate(agent.basal_ganglia.habit_counts): + h = " (习惯化)" if agent.basal_ganglia.is_habitual(i) else "" + print(f" 动作{i}: {int(c):4d}{h}") + if agent.hippocampus: + print(f"\n海马体: {agent.hippocampus.size()} 条记忆") + + if log: + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + steps = [s['step'] for s in log] + wn = [s['w_norm'] for s in log] + mods = [s['modality'] for s in log] + cm = {'image': 'green', 'screen': 'purple', 'audio': 'blue', + 'keyboard': 'orange', 'mouse': 'red', 'idle': 'gray'} + colors = [cm.get(m, 'gray') for m in mods] + + ax = axes[0, 0] + ax.scatter(steps, wn, c=colors, alpha=0.7, s=30) + ax.set_xlabel('Step'); ax.set_ylabel('||w||') + ax.set_title('Workspace Norm'); ax.grid(True, alpha=0.3) + from matplotlib.patches import Patch + ax.legend(handles=[Patch(facecolor=c, label=m) for m, c in cm.items()], + fontsize=7) + + ax = axes[0, 1] + mc2 = {} + for m in mods: mc2[m] = mc2.get(m, 0) + 1 + ax.bar(mc2.keys(), mc2.values(), + color=[cm.get(m, 'gray') for m in mc2.keys()]) + ax.set_title('Modality Distribution') + + ax = axes[1, 0] + ast = [s['action']['action_strength'] for s in log] + exf = [1 if s['action']['executed'] else 0 for s in log] + ax.plot(steps, ast, label='strength', alpha=0.7) + ax.scatter(steps, exf, c=['green' if e else 'red' for e in exf], + label='executed', alpha=0.5, s=20) + ax.set_title('Action Strength & Execution'); ax.legend(); ax.grid(True, alpha=0.3) + + ax = axes[1, 1] + ai = [s['action']['action_idx'] for s in log] + ax.hist(ai, bins=range(6), align='left', rwidth=0.8, color='steelblue', alpha=0.7) + ax.set_title('Basal Ganglia Action Selection') + ax.set_xticks(range(5)) + + plt.tight_layout() + Path('outputs').mkdir(exist_ok=True) + plt.savefig('outputs/embodied_live.png', dpi=150, bbox_inches='tight') + print(f"\n可视化: outputs/embodied_live.png") + + +def main(): + p = argparse.ArgumentParser(description='JspaceAI 具身 Agent') + p.add_argument('--mode', default='test', choices=['test', 'live', 'safe']) + p.add_argument('--steps', type=int, default=50) + p.add_argument('--device', default='cpu') + args = p.parse_args() + + dev = args.device + if dev == 'auto': + dev = 'cuda' if torch.cuda.is_available() else ( + 'mps' if torch.backends.mps.is_available() else 'cpu') + + if args.mode == 'test': + test_subsystems() + elif args.mode == 'live': + live_embodied(args.steps, dev, safe_mode=False) + elif args.mode == 'safe': + live_embodied(args.steps, dev, safe_mode=True) + + +if __name__ == '__main__': + main() diff --git a/main_full_sensory.py b/main_full_sensory.py new file mode 100644 index 0000000..eedb899 --- /dev/null +++ b/main_full_sensory.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +""" +JspaceAI 全感官交互 demo + +接入全部 5 个输入通道：摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标 + +模式： + --mode test: 测试各 I/O 通道是否工作 + --mode live: 实时全感官感知循环 + +运行： + python main_full_sensory.py --mode test + python main_full_sensory.py --mode live --steps 50 +""" +from __future__ import annotations + +import argparse +import torch +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +import time + +from jspaceai import ( + MultimodalConfig, MultimodalJSpaceModel, + FullSensoryStream, +) +from jspaceai.platform import get_screen_size, PLATFORM, print_permission_guide + + +def get_config() -> MultimodalConfig: + return MultimodalConfig( + vocab_size=50, embed_dim=16, input_dim=8, + workspace_dim=64, expert_dim=24, num_experts=12, + num_wells=4, ode_steps=3, dt=0.1, tau_w=0.3, + jacobian_sparsity=16, noise_std=0.01, + img_size=32, audio_frame_size=1024, keyboard_vocab=128, + ) + + +def test_io(): + """测试所有 I/O 通道""" + print("=" * 60) + print("测试全感官 I/O 通道") + print("=" * 60) + + stream = FullSensoryStream( + use_camera=True, use_mic=True, use_desktop=True, img_size=(32, 32), + ) + stream.start() + + print("\n采集 5 秒数据...（请移动鼠标、按键、对摄像头说话）") + log = {'camera': 0, 'audio': 0, 'screen': 0, 'keyboard': 0, 'mouse': 0} + mouse_positions = [] + + for i in range(25): + data = stream.get_latest() + if data['camera']: + log['camera'] += 1 + if data['audio']: + log['audio'] += 1 + if data['screen']: + log['screen'] += 1 + if data['keyboard']: + log['keyboard'] += len(data['keyboard']) + if data['mouse']: + log['mouse'] += len(data['mouse']) + mouse_positions.append(data['mouse_pos']) + time.sleep(0.2) + + stream.stop() + + print("\n" + "=" * 60) + print("采集结果") + print("=" * 60) + for ch, count in log.items(): + status = "OK" if count > 0 else "无数据" + print(f" {ch:10s}: {count:4d} 帧/事件 [{status}]") + + if mouse_positions: + xs = [p[0] for p in mouse_positions] + ys = [p[1] for p in mouse_positions] + print(f" 鼠标位置范围: x=[{min(xs)}, {max(xs)}], y=[{min(ys)}, {max(ys)}]") + + # 可视化 + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + + ax = axes[0] + if mouse_positions: + xs = [p[0] for p in mouse_positions] + ys = [p[1] for p in mouse_positions] + ax.plot(xs, ys, 'b.-', alpha=0.5) + ax.set_xlabel('X') + ax.set_ylabel('Y') + ax.set_title(f'Mouse Trajectory ({len(mouse_positions)} points)') + ax.invert_yaxis() + else: + ax.text(0.5, 0.5, "No mouse movement", transform=ax.transAxes, ha='center') + ax.set_title('Mouse') + + ax = axes[1] + channels = list(log.keys()) + counts = list(log.values()) + ax.bar(channels, counts, color=['green', 'blue', 'purple', 'orange', 'red']) + ax.set_ylabel('Event count') + ax.set_title('Channel Activity') + + plt.tight_layout() + outdir = Path('outputs') + outdir.mkdir(exist_ok=True) + fig_path = outdir / 'io_test.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"\n可视化: {fig_path}") + + +def live_full_sensory(n_steps: int, device: str): + """实时全感官感知循环""" + screen_w, screen_h = get_screen_size() + print(f"平台: {PLATFORM}") + print(f"屏幕: {screen_w}x{screen_h}") + + config = get_config() + model = MultimodalJSpaceModel(config).to(device) + + model_path = Path('outputs/multimodal_model.pt') + if model_path.exists(): + ckpt = torch.load(model_path, map_location=device, weights_only=False) + try: + model.load_state_dict(ckpt['model'], strict=False) + print(f"已加载模型: {model_path}") + except Exception as e: + print(f"模型加载失败，用随机初始化: {e}") + else: + print("未找到训练模型，用随机初始化") + model.eval() + + print(f"\n专家分工: {model.expert_modality}") + + print("\n启动全感官流（摄像头+麦克风+屏幕+键盘+鼠标）...") + stream = FullSensoryStream( + use_camera=True, use_mic=True, use_desktop=True, img_size=(32, 32), + ) + stream.start() + + state = model.init_state(1, torch.device(device)) + step_log = [] + + print(f"\n运行 {n_steps} 步感知循环...") + print("（移动鼠标、按键、对摄像头说话/做动作，模型会持续感知）") + print("=" * 60) + + try: + for step in range(n_steps): + data = stream.get_latest() + + modality = None + input_tensor = None + + if data['keyboard']: + key_events = data['keyboard'] + key_ids = [] + for ev in key_events[:8]: + k = ev.data if isinstance(ev.data, str) else ' ' + key_ids.append(ord(k[0]) if k and len(k) == 1 else 0) + if key_ids: + modality = 'keyboard' + input_tensor = torch.tensor([key_ids], dtype=torch.long).to(device) + + elif data['mouse']: + ev = data['mouse'][0] + x, y = ev.data + x_norm = x / screen_w + y_norm = y / screen_h + click_l = 1.0 if ev.modifiers.get('button') == 'left' else 0.0 + click_r = 1.0 if ev.modifiers.get('button') == 'right' else 0.0 + modality = 'mouse' + input_tensor = torch.tensor([[x_norm, y_norm, click_l, click_r]], + dtype=torch.float32).to(device) + + elif data['audio']: + audio = data['audio'].data + modality = 'audio' + input_tensor = torch.tensor([audio], dtype=torch.float32).to(device) + + elif data['screen']: + img = data['screen'].data + img_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device) / 127.5 - 1.0 + modality = 'screen' + input_tensor = img_tensor + + elif data['camera']: + img = data['camera'].data + img_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device) / 127.5 - 1.0 + modality = 'image' + input_tensor = img_tensor + + if modality is None: + time.sleep(0.2) + continue + + with torch.no_grad(): + x = model.encode_modality(modality, input_tensor) + if x.dim() == 1: + x = x.unsqueeze(0) + if x.dim() == 3: + x = x[:, -1, :] + if x.shape[0] != 1: + x = x[-1:] + state, _ = model.step(state, x) + w = state['w'] + w_norm = w.norm(dim=-1).mean().item() + + step_log.append({ + 'step': step, 'modality': modality, 'w_norm': w_norm, + }) + + if step % 5 == 0: + if step % 10 == 0 and step > 0: + with torch.no_grad(): + audio_out = model.audio_decoder(w)[0].cpu().numpy() + stream.play_audio(audio_out) + + mouse_pos = data.get('mouse_pos', (0, 0)) + key_buf = data.get('keyboard_buffer', '')[:20] + print(f" step {step:3d} | mod {modality:8s} | ||w|| {w_norm:.3f} | " + f"mouse={mouse_pos} | keys={repr(key_buf)}") + + time.sleep(0.15) + + except KeyboardInterrupt: + print("\n用户中断") + finally: + stream.stop() + + # 可视化 + if step_log: + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + + ax = axes[0] + steps = [s['step'] for s in step_log] + w_norms = [s['w_norm'] for s in step_log] + modalities = [s['modality'] for s in step_log] + color_map = {'image': 'green', 'screen': 'purple', 'audio': 'blue', + 'keyboard': 'orange', 'mouse': 'red'} + colors = [color_map.get(m, 'gray') for m in modalities] + ax.scatter(steps, w_norms, c=colors, alpha=0.7, s=30) + ax.set_xlabel('Step') + ax.set_ylabel('||w||') + ax.set_title('Workspace Norm by Modality') + ax.grid(True, alpha=0.3) + from matplotlib.patches import Patch + legend_elements = [Patch(facecolor=c, label=m) for m, c in color_map.items()] + ax.legend(handles=legend_elements) + + ax = axes[1] + mod_counts = {} + for m in modalities: + mod_counts[m] = mod_counts.get(m, 0) + 1 + ax.bar(mod_counts.keys(), mod_counts.values(), + color=[color_map.get(m, 'gray') for m in mod_counts.keys()]) + ax.set_ylabel('Count') + ax.set_title('Modality Distribution') + + plt.tight_layout() + outdir = Path('outputs') + outdir.mkdir(exist_ok=True) + fig_path = outdir / 'full_sensory_live.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"\n可视化: {fig_path}") + + print(f"\n完成，共 {len(step_log)} 步") + + +def main(): + parser = argparse.ArgumentParser(description='JspaceAI 全感官交互') + parser.add_argument('--mode', type=str, default='test', + choices=['test', 'live']) + parser.add_argument('--steps', type=int, default=50) + parser.add_argument('--device', type=str, default='cpu') + args = parser.parse_args() + + device = args.device + if device == 'auto': + if torch.cuda.is_available(): device = 'cuda' + elif torch.backends.mps.is_available(): device = 'mps' + else: device = 'cpu' + + if args.mode == 'test': + test_io() + elif args.mode == 'live': + live_full_sensory(args.steps, device) + + +if __name__ == '__main__': + main() diff --git a/main_language.py b/main_language.py new file mode 100644 index 0000000..3e9e8cd --- /dev/null +++ b/main_language.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +""" +JspaceAI 语言版 —— 自主进化实验 + +模型在 Shakespeare 文本上持续学习，边推理边进化。 +观察： + 1. loss 持续下降（在学习） + 2. 生成文本从乱码逐渐变成类 Shakespeare 风格 + 3. 专家分工涌现（不同专家处理不同字符模式） + 4. EWC + 经验回放防止灾难性遗忘 + +运行： + python main_language.py + python main_language.py --steps 300 --device mps +""" +from __future__ import annotations + +import argparse +import torch +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path + +from jspaceai import ( + LanguageConfig, JSpaceLanguageModel, EvolutionTrainer, + CharTokenizer, load_shakespeare, +) + + +def run_language_experiment(n_steps: int, device: str, outdir: Path): + outdir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(42) + np.random.seed(42) + + # 1. 数据 + text = load_shakespeare() + tokenizer = CharTokenizer.from_text(text) + print(f"文本长度: {len(text)} 字符") + print(f"词汇表大小: {tokenizer.vocab_size}") + print(f"词汇表: {''.join(tokenizer.chars[:50])}...") + + # 切成多段，模拟持续到来的文本流 + chunk_size = 200 + text_stream = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)] + print(f"文本流: {len(text_stream)} 段, 每段 {chunk_size} 字符") + + # 2. 模型 + config = LanguageConfig( + vocab_size=tokenizer.vocab_size, + embed_dim=16, + input_dim=8, # J-space 输入维度 + workspace_dim=32, # 工作空间维度 + expert_dim=16, # 每个专家内部状态 + num_experts=5, # 5 个并行专家 + num_wells=4, # 每个专家 4 个吸引子 + ode_steps=3, # ODE 积分子步 + dt=0.1, + tau_w=0.3, + jacobian_sparsity=8, + noise_std=0.005, + ) + model = JSpaceLanguageModel(config) + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"\n模型参数量: {n_params:,}") + + # 3. 自主进化训练器 + trainer = EvolutionTrainer( + model, config, + lr=5e-3, + ewc_lambda=0.05, # 较小的 EWC 权重，让模型能学新东西 + device=device, + ) + + # 4. 进化 + print("\n" + "=" * 70) + print("自主进化开始") + print("=" * 70) + + # 限制步数 + history = trainer.evolve( + text_stream, + tokenizer, + seq_len=48, + batch_size=4, + consolidate_every=30, + generate_every=30, + max_steps=n_steps, + prompt_text="To be", + ) + + # 5. 总结 + summary = trainer.get_evolution_summary() + print("\n" + "=" * 70) + print("进化总结") + print("=" * 70) + print(f"总步数: {summary['steps']}") + print(f"初始 loss: {summary['initial_loss']:.4f}") + print(f"最终 loss: {summary['final_loss']:.4f}") + print(f"最低 loss: {summary['min_loss']:.4f}") + print(f"loss 下降: {(1 - summary['final_loss']/summary['initial_loss'])*100:.1f}%") + print(f"\n专家最终使用率: {[f'{u:.3f}' for u in summary['expert_usage']]}") + + # 专家专业化（每个专家最常处理的 top-5 字符） + print("\n专家专业化（top-5 字符）:") + for i, spec in enumerate(summary['expert_specialization']): + chars = tokenizer.decode([s[0] for s in spec]) + weights = [f"{s[1]:.1f}" for s in spec] + display = ' '.join(f"{repr(c)}({w})" for c, w in zip(chars, weights)) + print(f" 专家 {i}: {display}") + + # 6. 最终生成对比 + print("\n" + "=" * 70) + print("最终生成对比") + print("=" * 70) + for prompt in ["To be", "Romeo", "The "]: + prompt_ids = tokenizer.encode(prompt) + generated = model.generate(prompt_ids, n_new=100, temperature=0.7, top_k=5) + sample = prompt + tokenizer.decode(generated) + print(f"\n提示 '{prompt}':") + print(f" {sample}") + + # 7. 可视化 + print("\n生成可视化...") + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + + # Loss 曲线 + ax = axes[0, 0] + losses = [h['loss'] for h in history] + ax.plot(losses, alpha=0.3, linewidth=0.5, color='blue', label='raw') + if len(losses) > 10: + smoothed = np.convolve(losses, np.ones(10)/10, mode='valid') + ax.plot(smoothed, linewidth=2, color='blue', label='smoothed') + ax.set_xlabel('Evolution step') + ax.set_ylabel('Cross-entropy loss') + ax.set_title('Loss during self-evolution') + ax.legend() + ax.grid(True, alpha=0.3) + + # 专家使用率演化 + ax = axes[0, 1] + alpha_history = np.array([h['alpha_mean'] for h in history]) + for i in range(config.num_experts): + ax.plot(alpha_history[:, i], label=f'Expert {i}', alpha=0.8) + ax.set_xlabel('Evolution step') + ax.set_ylabel('Attention weight (mean)') + ax.set_title('Expert usage during evolution') + ax.legend() + ax.grid(True, alpha=0.3) + + # ||w|| 演化 + ax = axes[1, 0] + w_norms = [h['w_norm_mean'] for h in history] + ax.plot(w_norms, linewidth=1.5, color='green') + ax.set_xlabel('Evolution step') + ax.set_ylabel('||w|| (mean)') + ax.set_title('Workspace norm during evolution') + ax.grid(True, alpha=0.3) + + # 生成样本随时间演化 + ax = axes[1, 1] + ax.axis('off') + samples = [(h['step'], h.get('sample', '')) for h in history if 'sample' in h] + text_lines = [] + for step, sample in samples[:8]: + s = sample[:60].replace('\n', ' ') + text_lines.append(f"step {step:3d}: {s}") + ax.text(0.05, 0.95, '\n'.join(text_lines), + transform=ax.transAxes, fontsize=9, verticalalignment='top', + fontfamily='monospace', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5)) + ax.set_title('Generated samples during evolution') + + plt.tight_layout() + fig_path = outdir / 'language_evolution.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"可视化已保存: {fig_path}") + plt.close() + + # 保存模型 + torch.save({ + 'model': model.state_dict(), + 'config': config, + 'tokenizer_chars': tokenizer.chars, + 'history': history, + 'summary': summary, + }, outdir / 'language_model.pt') + print(f"模型已保存: {outdir / 'language_model.pt'}") + + print("\n" + "=" * 70) + print("实验完成") + print("=" * 70) + + +def main(): + parser = argparse.ArgumentParser(description='JspaceAI 语言版自主进化实验') + parser.add_argument('--steps', type=int, default=100, + help='文本流段数 (default: 100)') + parser.add_argument('--device', type=str, default='cpu', + help='设备 (cpu / cuda / mps / auto)') + parser.add_argument('--outdir', type=str, default='outputs', + help='输出目录') + args = parser.parse_args() + + if args.device == 'auto': + if torch.cuda.is_available(): + device = 'cuda' + elif torch.backends.mps.is_available(): + device = 'mps' + else: + device = 'cpu' + else: + device = args.device + + run_language_experiment( + n_steps=args.steps, + device=device, + outdir=Path(args.outdir), + ) + + +if __name__ == '__main__': + main() diff --git a/main_multimodal.py b/main_multimodal.py new file mode 100644 index 0000000..bc03975 --- /dev/null +++ b/main_multimodal.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python3 +""" +JspaceAI 多模态实时交互 demo + +接入摄像头 + 麦克风 + 扬声器，原生支持图像/音频/文本。 + +模式： + 1. --mode train: 在合成多模态数据上训练（无需摄像头权限） + 2. --mode live: 实时感知-行动循环（需要摄像头/麦克风权限） + 3. --mode eval: 离线评估多模态对齐能力 + +运行： + python main_multimodal.py --mode train + python main_multimodal.py --mode live --steps 100 + python main_multimodal.py --mode eval +""" +from __future__ import annotations + +import argparse +import torch +import torch.nn.functional as F +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +import time + +from jspaceai import ( + MultimodalConfig, MultimodalJSpaceModel, + MultimodalStream, SensoryMotorLoop, +) + + +def get_config(vocab_size: int = 50) -> MultimodalConfig: + return MultimodalConfig( + vocab_size=vocab_size, + embed_dim=16, + input_dim=8, + workspace_dim=64, + expert_dim=24, + num_experts=8, + num_wells=4, + ode_steps=3, + dt=0.1, + tau_w=0.3, + jacobian_sparsity=16, + noise_std=0.01, + img_size=32, + audio_frame_size=1024, + ) + + +def generate_synthetic_multimodal(batch_size: int = 4, device: str = 'cpu'): + """生成合成多模态数据：3 个概念对应 3 种模态特征""" + import cv2 + concepts = ['A', 'B', 'C'] + colors = [(1, 0, 0), (0, 0, 1), (0, 1, 0)] + freqs = [200, 800, 400] + + images, audios, tokens = [], [], [] + for _ in range(batch_size): + idx = np.random.randint(3) + img = np.zeros((32, 32, 3), dtype=np.float32) + color = colors[idx] + if idx == 0: + cv2.circle(img, (16, 16), 8, color, -1) + elif idx == 1: + img[8:24, 8:24] = color + else: + for y in range(32): + w = min(y, 31 - y) + if 8 < y < 24: + img[31-y, 16-w:16+w] = color + images.append(img.transpose(2, 0, 1)) + + t = np.linspace(0, 1024/16000, 1024) + audio = np.sin(2 * np.pi * freqs[idx] * t).astype(np.float32) * 0.5 + audios.append(audio) + tokens.append(idx) + + images = torch.tensor(np.stack(images)).to(device) / 255.0 * 2 - 1 + audios = torch.tensor(np.stack(audios)).to(device) + tokens = torch.tensor(tokens, dtype=torch.long).to(device) + return images, audios, tokens + + +def train_multimodal(n_steps: int, device: str, outdir: Path): + outdir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(42); np.random.seed(42) + + config = get_config() + model = MultimodalJSpaceModel(config).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + + print(f"模型参数: {sum(p.numel() for p in model.parameters()):,}") + print(f"专家模态: {model.expert_modality}") + print(f"\n训练 {n_steps} 步...") + history = [] + + for step in range(n_steps): + images, audios, tokens = generate_synthetic_multimodal(8, device) + modality_choice = np.random.randint(3) + total_loss = 0 + + if modality_choice == 0: + outputs, _ = model.forward_multimodal('image', images) + text_loss = F.cross_entropy(outputs['text_logits'], tokens) + audio_loss = F.mse_loss(outputs['audio'], audios) + total_loss = text_loss + 0.1 * audio_loss + elif modality_choice == 1: + outputs, _ = model.forward_multimodal('audio', audios) + text_loss = F.cross_entropy(outputs['text_logits'], tokens) + img_loss = F.mse_loss(outputs['image'], images) + total_loss = text_loss + 0.1 * img_loss + else: + outputs, _ = model.forward_multimodal('text', tokens) + img_loss = F.mse_loss(outputs['image'], images) + audio_loss = F.mse_loss(outputs['audio'], audios) + total_loss = 0.1 * img_loss + 0.1 * audio_loss + + optimizer.zero_grad() + total_loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + history.append(total_loss.item()) + + if (step + 1) % 20 == 0: + print(f" step {step+1:4d} | loss {total_loss.item():.4f} | " + f"modality {['img','aud','txt'][modality_choice]}") + + print("\n" + "=" * 60) + print("跨模态对齐评估") + print("=" * 60) + + model.eval() + with torch.no_grad(): + for mod_name, mod_data in [('image', images), ('audio', audios), ('text', tokens)]: + outputs, _ = model.forward_multimodal(mod_name, mod_data[:3]) + pred_tokens = outputs['text_logits'].argmax(dim=-1) + print(f" 输入 {mod_name:6s} → 预测 token: {pred_tokens.cpu().tolist()} " + f"(真实: {tokens[:3].cpu().tolist()})") + + fig, axes = plt.subplots(2, 3, figsize=(12, 8)) + + ax = axes[0, 0] + ax.plot(history, alpha=0.5, linewidth=0.5) + if len(history) > 10: + smoothed = np.convolve(history, np.ones(10)/10, mode='valid') + ax.plot(smoothed, linewidth=2) + ax.set_title('Training Loss') + ax.set_xlabel('Step') + ax.grid(True, alpha=0.3) + + with torch.no_grad(): + test_img = images[:1] + outputs, _ = model.forward_multimodal('image', test_img) + + ax = axes[0, 1] + ax.imshow(((test_img[0].cpu().permute(1,2,0) + 1) / 2).numpy()) + ax.set_title('Input Image') + ax.axis('off') + + ax = axes[0, 2] + ax.imshow(((outputs['image'][0].cpu().permute(1,2,0) + 1) / 2).clamp(0,1).numpy()) + ax.set_title('Reconstructed Image') + ax.axis('off') + + ax = axes[1, 0] + input_audio = audios[0].cpu().numpy() + ax.plot(input_audio[:200], alpha=0.7, label='input') + out_audio = outputs['audio'][0].cpu().numpy() + ax.plot(out_audio[:200], alpha=0.7, label='reconstructed') + ax.set_title('Audio Waveform') + ax.legend() + ax.grid(True, alpha=0.3) + + ax = axes[1, 1] + logits = outputs['text_logits'][0].cpu().numpy() + probs = np.exp(logits) / np.exp(logits).sum() + ax.bar(['A', 'B', 'C'], probs[:3]) + ax.set_title('Text Logits (softmax)') + + ax = axes[1, 2] + w = outputs['w'][0].cpu().numpy() + ax.bar(range(len(w)), w) + ax.set_title('Workspace w (64-dim)') + + plt.tight_layout() + fig_path = outdir / 'multimodal_train.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"\n可视化: {fig_path}") + + torch.save({'model': model.state_dict(), 'config': config}, + outdir / 'multimodal_model.pt') + print(f"模型: {outdir / 'multimodal_model.pt'}") + + +def live_demo(n_steps: int, device: str, outdir: Path): + outdir.mkdir(parents=True, exist_ok=True) + config = get_config() + model = MultimodalJSpaceModel(config).to(device) + + model_path = outdir / 'multimodal_model.pt' + if model_path.exists(): + ckpt = torch.load(model_path, map_location=device) + model.load_state_dict(ckpt['model']) + print(f"已加载模型: {model_path}") + else: + print("未找到训练好的模型，用随机初始化运行") + model.eval() + + print("\n启动实时多模态流...") + print("（macOS 会请求摄像头和麦克风权限，请允许）") + stream = MultimodalStream( + use_camera=True, use_mic=True, + img_size=(32, 32), sample_rate=16000, audio_frame_size=1024, + ) + loop = SensoryMotorLoop(model, stream, device=device) + step_log = [] + + def on_step(info): + mod = info.get('modality', '?') + step = info.get('step', 0) + w_norm = info.get('w_norm', torch.tensor([0])).mean().item() + step_log.append({'step': step, 'modality': mod, 'w_norm': w_norm}) + if step % 5 == 0: + print(f" step {step:3d} | modality {mod:5s} | ||w|| {w_norm:.3f}") + + loop.run(n_steps=n_steps, interval=0.2, on_step=on_step) + + if step_log: + fig, ax = plt.subplots(1, 1, figsize=(10, 4)) + steps = [s['step'] for s in step_log] + w_norms = [s['w_norm'] for s in step_log] + modalities = [s['modality'] for s in step_log] + colors = ['blue' if m == 'audio' else 'green' for m in modalities] + ax.scatter(steps, w_norms, c=colors, alpha=0.6, s=20) + ax.set_xlabel('Step') + ax.set_ylabel('||w||') + ax.set_title('Workspace Norm During Live (blue=audio, green=image)') + ax.grid(True, alpha=0.3) + plt.tight_layout() + fig_path = outdir / 'multimodal_live.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"可视化: {fig_path}") + + +def eval_multimodal(device: str, outdir: Path): + model_path = outdir / 'multimodal_model.pt' + if not model_path.exists(): + print("请先运行 --mode train") + return + + config = get_config() + model = MultimodalJSpaceModel(config).to(device) + ckpt = torch.load(model_path, map_location=device) + model.load_state_dict(ckpt['model']) + model.eval() + + print("多模态对齐评估") + print("=" * 60) + + images, audios, tokens = generate_synthetic_multimodal(6, device) + + with torch.no_grad(): + for mod_name, mod_data in [('image', images), ('audio', audios), ('text', tokens)]: + outputs, _ = model.forward_multimodal(mod_name, mod_data[:3]) + pred_tokens = outputs['text_logits'].argmax(dim=-1) + correct = (pred_tokens == tokens[:3]).float().mean().item() + print(f" 输入 {mod_name:6s} → token 准确率: {correct:.3f}") + + +def main(): + parser = argparse.ArgumentParser(description='JspaceAI 多模态') + parser.add_argument('--mode', type=str, default='train', + choices=['train', 'live', 'eval']) + parser.add_argument('--steps', type=int, default=200) + parser.add_argument('--device', type=str, default='cpu') + parser.add_argument('--outdir', type=str, default='outputs') + args = parser.parse_args() + + device = args.device + if device == 'auto': + if torch.cuda.is_available(): device = 'cuda' + elif torch.backends.mps.is_available(): device = 'mps' + else: device = 'cpu' + + if args.mode == 'train': + train_multimodal(args.steps, device, Path(args.outdir)) + elif args.mode == 'live': + live_demo(args.steps, device, Path(args.outdir)) + elif args.mode == 'eval': + eval_multimodal(device, Path(args.outdir)) + + +if __name__ == '__main__': + main() diff --git a/main_v2.py b/main_v2.py new file mode 100644 index 0000000..a5ead0d --- /dev/null +++ b/main_v2.py @@ -0,0 +1,438 @@ +#!/usr/bin/env python3 +""" +JspaceAI v2 —— 基于 Anthropic J-space 论文优化版 + +新增功能： + 1. J-lens：观测模型内部"想法"（每个 ODE 子步的 workspace 读出） + 2. Directed Modulation：指令模型"想某概念"，验证 workspace 可被 top-down 调制 + 3. Selectivity 验证：ablate workspace，看是否只影响灵活推理 + 4. W 轨迹记录：forward 时记录每个子步的 w，用于 J-lens 训练和可视化 + +运行： + python main_v2.py + python main_v2.py --steps 200 --device mps +""" +from __future__ import annotations + +import argparse +import torch +import torch.nn.functional as F +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path + +from jspaceai import ( + LanguageConfig, JSpaceLanguageModel, EvolutionTrainer, + CharTokenizer, load_shakespeare, + JLensConfig, JLensSuite, + WorkspaceAblator, DirectedModulation, +) + + +def count_params(model): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def run_experiment(n_steps: int, device: str, outdir: Path): + outdir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(42) + np.random.seed(42) + + # 1. 数据 + text = load_shakespeare() + tokenizer = CharTokenizer.from_text(text) + print(f"文本: {len(text)} 字符, 词汇表: {tokenizer.vocab_size}") + + # 2. 模型 + config = LanguageConfig( + vocab_size=tokenizer.vocab_size, + embed_dim=16, input_dim=8, + workspace_dim=32, expert_dim=16, + num_experts=5, num_wells=4, + ode_steps=4, dt=0.1, tau_w=0.3, + jacobian_sparsity=8, noise_std=0.005, + ) + model = JSpaceLanguageModel(config) + print(f"模型参数: {count_params(model):,}") + + # 3. J-lens 套件 + jlens_config = JLensConfig( + n_substeps=config.ode_steps, + workspace_dim=config.workspace_dim, + vocab_size=config.vocab_size, + ) + jlens_suite = JLensSuite(jlens_config).to(device) + + # 4. 基础进化训练 + print("\n" + "=" * 70) + print("阶段 1：基础进化训练") + print("=" * 70) + + trainer = EvolutionTrainer(model, config, lr=5e-3, ewc_lambda=0.05, device=device) + chunks = [text[i:i+200] for i in range(0, len(text), 200)] + trainer.evolve( + chunks, tokenizer, + seq_len=48, batch_size=4, + consolidate_every=30, generate_every=50, + max_steps=n_steps, prompt_text="To be", + ) + + # 5. 训练 J-lens + print("\n" + "=" * 70) + print("阶段 2：训练 J-lens 探针") + print("=" * 70) + + jlens_optimizer = torch.optim.Adam(jlens_suite.parameters(), lr=3e-3) + model.eval() + all_tokens = tokenizer.encode(text) + + for jlens_epoch in range(60): + batch_tokens = [] + for _ in range(8): + start = np.random.randint(0, len(all_tokens) - 64) + seq = all_tokens[start:start+48] + batch_tokens.append(seq) + token_seq = torch.tensor(batch_tokens, dtype=torch.long).to(device) + + with torch.no_grad(): + logits, info = model(token_seq, record_trajectory=True) + + if 'w_trajectory' in info: + w_traj = info['w_trajectory'] # (batch, T, n_substeps, workspace_dim) + # target: 每个位置的下一个 token + targets = token_seq[:, 1:] # (batch, T-1) + + jlens_optimizer.zero_grad() + total_loss = 0 + for substep in range(jlens_config.n_substeps): + # (batch, T-1, workspace_dim) → 预测 targets + w_sub = w_traj[:, :-1, substep, :] # (batch, T-1, workspace_dim) + pred = jlens_suite.probes[substep](w_sub) # (batch, T-1, vocab) + loss = F.cross_entropy( + pred.reshape(-1, config.vocab_size), + targets.reshape(-1), + ) + total_loss += loss + + total_loss.backward() + jlens_optimizer.step() + + if (jlens_epoch + 1) % 15 == 0: + print(f" J-lens epoch {jlens_epoch+1}/60, loss={total_loss.item():.4f}") + + model.train() + + # 6. J-lens 观测 + print("\n" + "=" * 70) + print("阶段 3：J-lens 观测——模型在想什么") + print("=" * 70) + + model.eval() + with torch.no_grad(): + prompt = "To be" + prompt_ids = tokenizer.encode(prompt) + token_tensor = torch.tensor([prompt_ids], dtype=torch.long).to(device) + _, info = model(token_tensor, record_trajectory=True) + + if 'w_trajectory' in info: + w_traj = info['w_trajectory'][0] + T, n_sub, _ = w_traj.shape + sub_labels = ['sensory', 'workspace', 'workspace', 'motor'] + + print(f"\n提示: '{prompt}'") + for t in range(min(T, 6)): + char = prompt[t] if t < len(prompt) else '?' + print(f" pos {t} ('{char}'):") + for s in range(n_sub): + w = w_traj[t, s].unsqueeze(0) + probe = jlens_suite.probes[s] + logits_s = probe(w) + probs = F.softmax(logits_s, dim=-1) + topk_probs, topk_idx = probs[0].topk(5) + tokens_list = [tokenizer.idx_to_char.get(i.item(), '?') for i in topk_idx] + probs_str = [f"{p:.2f}" for p in topk_probs.tolist()] + pairs = ' '.join(f"{repr(tc)}({pr})" for tc, pr in zip(tokens_list, probs_str)) + label = sub_labels[s] if s < len(sub_labels) else f's{s}' + print(f" 子步{s} ({label}): {pairs}") + + # 7. Directed Modulation + print("\n" + "=" * 70) + print("阶段 4：Directed Modulation") + print("=" * 70) + + modulation = DirectedModulation(model, jlens_suite) + test_concepts = ['R', 'd', 'o', ' '] + + mod_results = [] + model.eval() + for concept_char in test_concepts: + if concept_char not in tokenizer.char_to_idx: + continue + concept_id = tokenizer.char_to_idx[concept_char] + + with torch.no_grad(): + # 正常生成 + prompt_ids = list(tokenizer.encode("To be")) + state = model.init_state(1, device) + for tok in prompt_ids: + state, _, _, _ = model.step(state, torch.tensor([tok], device=device)) + + normal_gen = [] + s = state + last_tok = prompt_ids[-1] + for _ in range(25): + s, logits, _, _ = model.step(s, torch.tensor([last_tok], device=device)) + next_tok = logits[0].argmax().item() + normal_gen.append(next_tok) + last_tok = next_tok + + # 调制生成 + mod_gen = [] + prompt_ids2 = list(tokenizer.encode("To be")) + s = model.init_state(1, device) + for tok in prompt_ids2: + s, _, _, _ = model.step(s, torch.tensor([tok], device=device)) + + last_tok = prompt_ids2[-1] + for _ in range(25): + s = modulation.modulate_state(s, concept_id, strength=3.0) + s, logits, _, _ = model.step(s, torch.tensor([last_tok], device=device)) + next_tok = logits[0].argmax().item() + mod_gen.append(next_tok) + last_tok = next_tok + + normal_str = tokenizer.decode(normal_gen) + mod_str = tokenizer.decode(mod_gen) + mod_results.append((concept_char, normal_str, mod_str)) + print(f"\n 注入 '{concept_char}':") + print(f" 正常: {repr(normal_str[:30])}") + print(f" 调制: {repr(mod_str[:30])}") + + # 8. Selectivity 验证——对比"自动任务"vs"需要 workspace 的任务" + print("\n" + "=" * 70) + print("阶段 5：Selectivity 验证") + print("=" * 70) + + # 任务 A: 简单续写（自动任务，应该 ablate 不影响） + test_seqs = [] + for _ in range(8): + start = np.random.randint(0, len(all_tokens) - 64) + test_seqs.append(all_tokens[start:start+48]) + test_tensor = torch.tensor(test_seqs, dtype=torch.long).to(device) + + model.eval() + with torch.no_grad(): + # 正常 forward + normal_logits, _ = model(test_tensor) + normal_pred = normal_logits[:, :-1].argmax(dim=-1) + targets = test_tensor[:, 1:] + normal_acc = (normal_pred == targets).float().mean().item() + + # Ablate workspace forward + ablate_state = model.init_state(test_tensor.shape[0], device) + ablate_preds = [] + for t in range(test_tensor.shape[1] - 1): + ablate_state, logits, _, _ = model.step(ablate_state, test_tensor[:, t]) + w = ablate_state['w'] + k = 5 # ablate top-5 J-lens 方向 + for b in range(w.shape[0]): + probe = jlens_suite.probes[2] + w_logits = probe(w[b:b+1]) + topk_vals, topk_idx = w_logits[0].topk(k) + for idx in topk_idx: + d = probe.lens.weight[idx] + d_norm = d / (d.norm() + 1e-8) + w[b] = w[b] - (w[b] @ d_norm) * d_norm + ablate_state['w'] = w + ablate_preds.append(logits.argmax(dim=-1)) + + ablate_preds = torch.stack(ablate_preds, dim=1) + ablate_acc = (ablate_preds == targets).float().mean().item() + + # 任务 B: 长程记忆（需要 workspace 持续装载信息） + # 构造序列：前半段是"key"，后半段需要回忆 key 的特征 + # 简化版：序列 [A, B, C, ..., A, ?] —— 第二次出现 A 后预测下一个 + # 对字符级，我们看"重复字符"任务：序列里某个字符重复出现，模型要"记住"它 + memory_seqs = [] + for _ in range(8): + # 构造：随机字符 X 出现在位置 0，然后在位置 40 重复 + x = np.random.choice(all_tokens) + seq = [np.random.choice(all_tokens) for _ in range(48)] + seq[0] = x + seq[40] = x # 40 步后重复 + memory_seqs.append(seq) + memory_tensor = torch.tensor(memory_seqs, dtype=torch.long).to(device) + + with torch.no_grad(): + # 正常 + mem_logits, _ = model(memory_tensor) + mem_pred = mem_logits[:, :-1].argmax(dim=-1) + mem_targets = memory_tensor[:, 1:] + # 只看位置 40 之后（需要记忆的位置） + mem_mask = torch.zeros_like(mem_targets, dtype=torch.bool) + mem_mask[:, 40:] = True # 位置 40+ 需要"回忆" + normal_mem_acc = (mem_pred[mem_mask] == mem_targets[mem_mask]).float().mean().item() + + # Ablate + abl_state = model.init_state(memory_tensor.shape[0], device) + abl_preds = [] + for t in range(memory_tensor.shape[1] - 1): + abl_state, logits, _, _ = model.step(abl_state, memory_tensor[:, t]) + w = abl_state['w'] + for b in range(w.shape[0]): + probe = jlens_suite.probes[2] + w_logits = probe(w[b:b+1]) + topk_vals, topk_idx = w_logits[0].topk(5) + for idx in topk_idx: + d = probe.lens.weight[idx] + d_norm = d / (d.norm() + 1e-8) + w[b] = w[b] - (w[b] @ d_norm) * d_norm + abl_state['w'] = w + abl_preds.append(logits.argmax(dim=-1)) + + abl_preds = torch.stack(abl_preds, dim=1) + ablate_mem_acc = (abl_preds[mem_mask] == mem_targets[mem_mask]).float().mean().item() + + print(f"\n 任务 A (简单续写 - 自动认知):") + print(f" 正常: {normal_acc:.3f} Ablate: {ablate_acc:.3f} 下降: {normal_acc - ablate_acc:.3f}") + print(f"\n 任务 B (长程记忆 - 需要 workspace):") + print(f" 正常: {normal_mem_acc:.3f} Ablate: {ablate_mem_acc:.3f} 下降: {normal_mem_acc - ablate_mem_acc:.3f}") + print(f"\n 解读: 任务 B 下降应大于任务 A —— workspace 对长程记忆更关键") + + # 9. 可视化 + print("\n" + "=" * 70) + print("阶段 6：可视化") + print("=" * 70) + + fig, axes = plt.subplots(2, 3, figsize=(18, 10)) + + ax = axes[0, 0] + losses = [h['loss'] for h in trainer.history] + ax.plot(losses, alpha=0.3, linewidth=0.5, color='blue') + if len(losses) > 10: + smoothed = np.convolve(losses, np.ones(10)/10, mode='valid') + ax.plot(smoothed, linewidth=2, color='blue') + ax.set_xlabel('Step') + ax.set_ylabel('Loss') + ax.set_title('Evolution Loss') + ax.grid(True, alpha=0.3) + + ax = axes[0, 1] + alpha_history = np.array([h['alpha_mean'] for h in trainer.history]) + for i in range(config.num_experts): + ax.plot(alpha_history[:, i], label=f'Expert {i}', alpha=0.8) + ax.set_xlabel('Step') + ax.set_ylabel('Usage') + ax.set_title('Expert Usage') + ax.legend() + ax.grid(True, alpha=0.3) + + ax = axes[0, 2] + w_norms = [h['w_norm_mean'] for h in trainer.history] + ax.plot(w_norms, linewidth=1.5, color='green') + ax.set_xlabel('Step') + ax.set_ylabel('||w||') + ax.set_title('Workspace Norm') + ax.grid(True, alpha=0.3) + + # J-lens 热力图 + ax = axes[1, 0] + if 'w_trajectory' in info: + w_traj = info['w_trajectory'][0] + T, n_sub, _ = w_traj.shape + sub_labels = ['sensory', 'workspace', 'workspace', 'motor'] + for s in range(min(n_sub, 4)): + chars_row = [] + for t in range(T): + w = w_traj[t, s].unsqueeze(0) + probe = jlens_suite.probes[s] + logits_s = probe(w) + top1 = logits_s.argmax().item() + chars_row.append(tokenizer.idx_to_char.get(top1, '?')) + chars = ''.join(chars_row) + label = sub_labels[s] if s < len(sub_labels) else f's{s}' + ax.text(0.05, 0.95 - s*0.2, f"{label}: {chars}", + transform=ax.transAxes, fontsize=9, fontfamily='monospace', + verticalalignment='top') + ax.set_title('J-lens Top-1 (position x substep)') + ax.axis('off') + + # Modulation 对比 + ax = axes[1, 1] + ax.axis('off') + mod_text = "Directed Modulation:\n\n" + for concept_char, normal_str, mod_str in mod_results[:3]: + mod_text += f"Inject '{concept_char}':\n" + mod_text += f" N: {repr(normal_str[:20])}\n" + mod_text += f" M: {repr(mod_str[:20])}\n\n" + ax.text(0.05, 0.95, mod_text, transform=ax.transAxes, fontsize=8, + verticalalignment='top', fontfamily='monospace', + bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8)) + ax.set_title('Directed Modulation') + + # Selectivity + ax = axes[1, 2] + tasks = ['Auto\n(continuation)', 'Memory\n(long-range)'] + normal_vals = [normal_acc, normal_mem_acc] + ablate_vals = [ablate_acc, ablate_mem_acc] + x = np.arange(len(tasks)) + width = 0.35 + ax.bar(x - width/2, normal_vals, width, label='Normal', color='green', alpha=0.7) + ax.bar(x + width/2, ablate_vals, width, label='Ablated', color='red', alpha=0.7) + ax.set_ylabel('Accuracy') + ax.set_title('Selectivity: Workspace Ablation') + ax.set_xticks(x) + ax.set_xticklabels(tasks) + ax.legend() + for i, (n, a) in enumerate(zip(normal_vals, ablate_vals)): + ax.text(i - width/2, n + 0.005, f'{n:.3f}', ha='center', fontsize=8) + ax.text(i + width/2, a + 0.005, f'{a:.3f}', ha='center', fontsize=8) + + plt.tight_layout() + fig_path = outdir / 'experiment_v2.png' + plt.savefig(fig_path, dpi=150, bbox_inches='tight') + print(f"可视化: {fig_path}") + plt.close() + + torch.save({ + 'model': model.state_dict(), + 'jlens': jlens_suite.state_dict(), + 'config': config, + 'tokenizer_chars': tokenizer.chars, + 'selectivity': { + 'auto': {'normal': normal_acc, 'ablate': ablate_acc}, + 'memory': {'normal': normal_mem_acc, 'ablate': ablate_mem_acc}, + }, + }, outdir / 'model_v2.pt') + + print("\n" + "=" * 70) + print("完成") + print("=" * 70) + print(f"基础训练: {len(trainer.history)} 步") + print(f"J-lens: {jlens_config.n_substeps} 个探针") + print(f"Selectivity:") + print(f" 自动任务下降: {normal_acc - ablate_acc:.3f}") + print(f" 记忆任务下降: {normal_mem_acc - ablate_mem_acc:.3f}") + + +def main(): + parser = argparse.ArgumentParser(description='JspaceAI v2') + parser.add_argument('--steps', type=int, default=200) + parser.add_argument('--device', type=str, default='cpu') + parser.add_argument('--outdir', type=str, default='outputs') + args = parser.parse_args() + + device = args.device + if device == 'auto': + if torch.cuda.is_available(): + device = 'cuda' + elif torch.backends.mps.is_available(): + device = 'mps' + else: + device = 'cpu' + + run_experiment(n_steps=args.steps, device=device, outdir=Path(args.outdir)) + + +if __name__ == '__main__': + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..258a24c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +torch>=2.0 +numpy>=1.24,<2 # torch 2.2 不兼容 numpy 2.x +matplotlib>=3.7 +opencv-python>=4.5 +sounddevice>=0.4 +soundfile>=0.12 +mss>=9.0 +pynput>=1.7 +av>=10.0