pulsareon/JspaceAI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor JspaceAI codebase: remove main_multimodal.py and main_v2.py, add main_chat.py for interactive dialogue mode with online learning capabilities.

Browse Source
This commit is contained in:
XiuchengWu
2026-07-07 11:22:53 +08:00
parent e782ef4db1
commit d5affe4458
15 changed files with 651 additions and 2795 deletions
Download Patch File Download Diff File Expand all files Collapse all files

258
README.md
View File

@@ -2,9 +2,9 @@
**全局工作空间 + J-space 广播的智慧系统**
两个版本：
1. **连续序列版** (`main.py`)：验证架构在多模态序列预测上的优势
2. **语言版 + 自主进化** (`main_language.py`)：字符级语言建模 + 边推理边学习
两个入口：
1. **全部接入** (`main.py`)：具身 Agent + 完整神经系统 + 5 感官 + 4 输出 + 自主心智
2. **只有对话** (`main_chat.py`)：字符级语言模型 + 交互对话 + 在线学习
基于第一性原理推导的智慧架构：不是 Transformer，不是 RNN，而是 **ODE 动力系统 + 并行专家 + J-space 工作空间广播 + 在线进化**。
@@ -59,67 +59,34 @@ pip install -r requirements.txt
## 运行
### 连续序列版（验证架构）
项目只有两个入口：
| 入口 | 用途 |
|---|---|
| `main.py` | 全部接入：具身 Agent + 完整神经系统 + 5 感官 + 4 输出 + 自主心智 |
| `main_chat.py` | 只有对话：字符级语言模型 + 交互对话 + 在线学习 |
### 全部接入版（`main.py`）
```bash
python main.py
python main.py --steps 2000 --device cuda
# 子系统自检（权限 + I/O + 模型 + 海马体回忆）
python main.py --mode test
# 实时具身循环 + 自主心智（默认安全：不控制鼠标键盘）
python main.py --mode live --steps 100
# 安全模式（更高动作门控阈值，不显示屏幕输出）
python main.py --mode safe --steps 50
# 允许鼠标键盘输出（需谨慎，加 --unsafe）
python main.py --mode live --steps 200 --unsafe
```
### 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
```
接入全部 5 个输入通道 + 4 个输出 + 神经系统：
- **输入**：摄像头(OpenCV) + 麦克风(PortAudio) + 屏幕(mss) + 键盘(pynput) + 鼠标(pynput)
- **输出**：鼠标 + 键盘 + 音频(PortAudio) + 屏幕(OpenCV 窗口)
- **神经系统**：小脑(前向+逆模型) / 中枢神经(反射弧+门控) / 海马体(向量检索) / 基底神经节(Q-learning)
- **自主心智**：好奇心驱动 + 状态持久化 + 自我模型 + 元学习
完整的人类神经系统对应：
@@ -140,44 +107,38 @@ python main_embodied.py --mode safe --steps 30
→ 4. 决策（基底神经节）→ 5. 精化（小脑）→ 6. 门控（中枢神经）
→ 7. 行动（手足口）→ 8. 学习（更新所有系统）
### 守护进程（后台持续运行）
状态保存在 `outputs/mind/`：workspace + 海马体 + 基底神经节 + 自我模型，再次运行从上次继续。
**macOS 权限**：键盘/鼠标监听需要在 系统设置 → 隐私与安全 → 辅助功能 中授权终端。
### 对话版（`main_chat.py`）
```bash
# 启动后台守护进程（静默运行，不干扰使用）
python daemon.py start
# 交互对话（默认）
python main_chat.py
# 低功耗模式（更长间隔，适合长期后台）
python daemon.py start --low-power --interval 2.0
# 先预训练再对话
python main_chat.py --mode train --steps 100
python main_chat.py --mode chat
# 前台运行（调试用，可看实时输出）
python daemon.py start --fg
# 查看状态
python daemon.py status
# 查看日志
python daemon.py log -n 30
# 心智自我认知
python daemon.py introspect
# 停止
python daemon.py stop
# 一次性生成
python main_chat.py --mode generate --prompt "To be"
```
守护进程特性：
- **用户主动控制**：手动 start/stop，不开机自启
- **静默后台运行**：不弹窗、不发声、不控制鼠标键盘
- **持续感知学习**：感知屏幕/键盘/鼠标/摄像头/麦克风 + 好奇心驱动学习
- **状态持久化**：停止后状态保存，再次启动从上次继续
- **跨会话持续**：海洋不蒸发，每次启动还是同一个"自己"
- **崩溃恢复**：异常不崩溃，等待后继续
对话特性：
- **在线学习**：每次用户输入都会被模型学习（next-token loss + 梯度更新）
- **持续进化**：EWC + 经验回放 + 专家可塑性防灾难性遗忘
- **状态持久化**：对话状态保存到 `outputs/chat_model.pt`，再次启动从上次继续
- **命令**：`/quit` 退出 `/save` 保存 `/reset` 重置 `/train N` 预训练N步
状态保存在 `~/.jspaceai/`：
- `state/`：workspace + 海马体 + 基底神经节 + 自我模型
- `logs/daemon.log`：运行日志
- `daemon.pid`：进程 PID
- `status.json`：实时状态
对话模式示例：
```
你: Hello
AI: ...
你: /save
你: /train 50
你: /quit
```
## 外挂模块系统（可热插拔）
@@ -212,38 +173,29 @@ dock.status() # 各模块连接状态
核心原则：**心智不依赖外挂**。断开任何外挂，心智继续工作，只是"知道得少"。
## 改进版 v2（异构专家 + 大 workspace + 蒸馏）
## 核心实现特性
四大改进：
`JSpaceConfig` 支持以下可配置项（默认值已开启推荐配置）：
1. **异构专家**：视觉/听觉/语言/跨模态专家有不同架构（不再全是相同 ODE）
2. **workspace 扩容**：64 → 256 维，加 LayerNorm 防衰减
3. **RK4 积分**：比 Euler 稳定（||w|| 从 0.05 提升到 16.0）
4. **小模型蒸馏**：接 GPT-2/Qwen 等，迁移理解能力
1. **RK4 积分**（`use_rk4=True`）：比 Euler 稳定，workspace 范数从 ~0.3 提升到 ~5.6
2. **LayerNorm**（`use_layer_norm=True`）：workspace 和专家状态都加 LayerNorm，防止长期衰减
3. **势能系数归一化**（`well_coeff=None` 自动 `1/sqrt(num_wells)`）：避免井数多时 softplus 项主导梯度导致专家状态发散
4. **异构专家支持**：`Expert` 构造可传入 `encoder` 参数，让不同模态专家有独立编码器
5. **轨迹记录**：`forward(xs, record_trajectory=True)` 返回 `info['w_trajectory']`，供 J-lens 训练
```python
from jspaceai import JSpaceConfigV2, JSpaceModelV2, SmallModelEncoder, DistillationTrainer
from jspaceai import JSpaceConfig, JSpaceModel
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)
# 蒸馏后小模型可断开，编码器已有理解能力
config = JSpaceConfig(
workspace_dim=32,
num_experts=5,
use_rk4=True, # RK4 积分
use_layer_norm=True, # LayerNorm 防衰减
)
model = JSpaceModel(config)
preds, info = model(xs, record_trajectory=True) # info['w_trajectory'] 可用
```
对比 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：
@@ -261,7 +213,7 @@ trainer.train(n_steps=100)
## 自主进化机制
语言版不只是"训练一次"，而是**天生支持持续进化**：
模型不只是"训练一次"，而是**天生支持持续进化**：
1. **在线学习**：每个 forward 都累积梯度并更新参数，模型永不"训练完成"
2. **EWC（Elastic Weight Consolidation）**：保护重要参数防灾难性遗忘
@@ -274,63 +226,60 @@ trainer.train(n_steps=100)
## 输出
### 连续序列版
- `outputs/experiment.png` — Loss/注意力/||w||/预测对比
- `outputs/models.pt` — 模型权重
### 全部接入版
- `outputs/embodied_live.png` — ||w||/模态分布/好奇心/自我模型/世界模型loss/海马体
- `outputs/mind/` — 自主心智状态（跨会话持久化）
### v2 版
- `outputs/experiment_v2.png` — Loss/专家使用率/||w||/J-lens 读出/Modulation/Selectivity
- `outputs/model_v2.pt` — 模型 + J-lens 权重
### 对话版
- `outputs/chat_model.pt` — 对话模型状态（跨会话持久化）
## 基于 Anthropic 2026 论文的优化
## 基于 Anthropic 2026 论文的设计
v2 基于 [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) 的发现：
基于 [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) 的发现：
1. **J-lens 探针**：学习线性映射 L: workspace → vocab，观测每个子步 workspace "准备说什么"
1. **J-lens 探针**：学习线性映射 L: workspace → vocab，观测每个子步 workspace "准备说什么"（`jlens.py` 提供）
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. **预测对比**：模式切换处哪个模型更鲁棒？
### 具身版
- 实时 5 通道感知（摄像头+麦克风+屏幕+键盘+鼠标）全部采集成功
- 感知-思考-行动闭环完整运行
- 海马体回忆正常（sim 0.99+）
- 小脑计算动作参数，基底神经节选择动作，中枢神经门控
## 文件结构
```
JspaceAI/
├── main.py # 主程序：对比实验
├── main.py # 全部接入（具身 + 自主心智）
├── main_chat.py # 只有对话
├── requirements.txt
├── README.md
└── jspaceai/
├── __init__.py
├── core.py # 核心：Expert + JSpaceWorkspace + JSpaceModel
├── core.py # 核心：Expert + JSpaceWorkspace + JSpaceModel（RK4+LayerNorm）
├── baselines.py # 对比基线：FlatBaseline
├── task.py # 玩具任务：多模态连续序列
└── trainer.py # 训练器：预测学习
├── trainer.py # 训练器：预测学习
├── language_data.py # 字符级 tokenizer + Shakespeare 语料
├── language_model.py # 语言版 + EWC + 经验回放 + 专家可塑性
├── jlens.py # J-lens 探针 + WorkspaceAblator + DirectedModulation
├── multimodal.py # 多模态核心（6 模态编解码 + 12 专家）
├── realtime.py # 摄像头 + 麦克风 + SensoryMotorLoop
├── desktop.py # 屏幕 + 键盘鼠标 + FullSensoryStream
├── platform.py # 跨平台抽象 + 权限检查
├── embodied.py # 神经系统（小脑/海马体/基底神经节/中枢神经）+ EmbodiedAgent
├── autonomous.py # 自主心智（好奇心/持久化/自我模型/元学习）
├── modules.py # 可热插拔外挂模块
└── evolution.py # 自主进化训练器
```
## 数学细节
@@ -339,7 +288,7 @@ JspaceAI/
$$\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)$
势能 $U_i(m_i) = \frac{1}{2}\|m_i\|^2 - \frac{c}{2}\sum_k \text{softplus}(a_k \cdot m_i + b_k)$，其中 $c = 1/\sqrt{\text{num\_wells}}$ 防梯度被 softplus 主导。
多井结构让专家在吸引子之间漫游 = 内部"思考"。
@@ -364,12 +313,11 @@ $$\mathcal{L} = \mathbb{E}\left[\|x_{t+\Delta t} - Q(w_t)\|^2\right]$$
这是**最小验证版本**，不是生产系统：
- 玩具任务（8 维序列），不是语言/图像
- ODE 用 Euler 积分（精度低，但简单可 backprop）
- 专家数 5，不是大规模
- 无持续跨会话状态（每次 forward 从零初始化）
- 对话版字符级，不是 subword
**下一步扩展方向**：
1. 用 adjoint method 替换 Euler（Neural ODE 路线）
2. 加跨会话状态持久化（真正的"海洋"）
3. Scale 到更大任务（MNIST 连续预测 → 语言建模）
4. 加输出门控的真正"自发输出"（非每步都预测）
1. 用 adjoint method 替换 RK4（Neural ODE 路线，更低显存）
2. Scale 到更大任务（MNIST 连续预测 → 语言建模）
3. 加输出门控的真正"自发输出"（非每步都预测）
4. 接 Whisper/VITS 真实语音输入输出

461
daemon.py
View File

@@ -1,461 +0,0 @@
"""
守护进程 —— 自主心智后台持续运行
用户主动控制：
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()

8
jspaceai/__init__.py
View File

@@ -2,7 +2,7 @@
JspaceAI —— 全局工作空间 + J-space 广播的智慧系统
模块：
1. core.py: 连续序列版（Expert + JSpaceWorkspace + JSpaceModel）
1. core.py: 核心架构（Expert + JSpaceWorkspace + JSpaceModel，含 RK4/LayerNorm/异构专家）
2. language_model.py: 语言版 + 自主进化
3. jlens.py: J-lens 可解释性工具
4. multimodal.py: 多模态（图像/音频/视频/文本）
@@ -77,9 +77,6 @@ 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__ = [
@@ -121,9 +118,6 @@ __all__ = [
# 外挂模块系统（可热插拔）
"ExternalModule", "SmallModelModule", "KnowledgeBaseModule",
"ToolModule", "ModuleDock",
# 改进版核心 v2
"JSpaceConfigV2", "JSpaceModelV2",
"SmallModelEncoder", "DistillationTrainer",
# 自主进化
"EvolutionTrainer",
]

178
jspaceai/core.py
View File

@@ -1,13 +1,13 @@
"""
核心架构：专家模块 + J-space 工作空间 + ODE 动力学
数学形式（每个 forward 时间步内做 Euler 积分若干子步）：
数学形式（每个 forward 时间步内做若干子步积分）：
专家 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 形成的局部吸引子）
U_i(m_i) = ½ ||m_i||² - ½ Σ_k softplus(a_ik · m_ik + b_ik) / sqrt(num_wells)
（多井势能：阻尼项 + softplus 形成的局部吸引子，系数按井数归一化防梯度被 softplus 项主导）
工作空间 w:
τ_w · dw/dt = -w + Σ_i α_i · P_i_out(m_i)
@@ -16,6 +16,10 @@
Jacobian 路由 J_i: 稀疏线性映射，每个专家只对 w 的少数维度敏感
输出门控：当 ||w|| > θ 时触发输出 R(w)
支持两种 ODE 积分：Euler（简单可 backprop）和 RK4（更稳定，||w|| 量级显著提升）。
可选 LayerNorm 防止 workspace 长期衰减。
专家可携带模态特定编码器（异构专家），用于多模态场景。
所有参数都可 backprop 训练。学习目标是预测下一时刻的输入。
"""
from __future__ import annotations
@@ -23,7 +27,7 @@ from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from dataclasses import dataclass, field
@dataclass
@@ -40,6 +44,10 @@ class JSpaceConfig:
output_threshold: float = 0.5 # 输出门控阈值（船舶涌出阈值）
jacobian_sparsity: int = 8 # 每个 J_i 只保留前 k 大的连接
noise_std: float = 0.01 # 内部噪声 ξ(t) 的标准差
# —— v2 合并进来的改进字段（带默认值，向后兼容）——
use_rk4: bool = True # 用 RK4 积分（比 Euler 稳定）
use_layer_norm: bool = True # workspace 加 LayerNorm 防衰减
well_coeff: float | None = None # softplus 项系数，None 则自动 1/sqrt(num_wells)
class Expert(nn.Module):
@@ -47,25 +55,36 @@ 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
势能：U_i(m_i) = ½||m_i||² - ½ · c · Σ_k softplus(a_k · m_i + b_k)
- ½||m_i||² 是阻尼项（拉回原点）
- softplus 项创造多个局部吸引子（多井势能 → 内部"思考"）
- c = well_coeff，默认 1/sqrt(num_wells)，避免井数多时梯度被 softplus 主导
动力学：dm_i/dt = -∇U_i(m_i) + J_i · w + P_in · x + ξ
J_i 是稀疏 Jacobian：从工作空间 w 路由信息进来。
支持 RK4 积分（use_rk4=True）和 LayerNorm（use_layer_norm=True）。
可选 encoder：把模态特定输入投影到 expert_dim（异构专家用）。
"""
def __init__(self, expert_dim: int, workspace_dim: int, input_dim: int,
num_wells: int, sparsity: int):
num_wells: int, sparsity: int,
use_rk4: bool = True, use_layer_norm: bool = True,
well_coeff: float | None = None,
encoder: nn.Module | None = None):
super().__init__()
self.expert_dim = expert_dim
self.workspace_dim = workspace_dim
self.num_wells = num_wells
self.use_rk4 = use_rk4
self.use_layer_norm = use_layer_norm
# softplus 项系数：默认按井数归一化，防止梯度被 softplus 项主导
if well_coeff is None:
well_coeff = 1.0 / (num_wells ** 0.5)
self.well_coeff = well_coeff
# 势能景观参数：每个井是一个 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_a = nn.Parameter(torch.randn(num_wells, expert_dim) * 0.2)
self.well_b = nn.Parameter(torch.zeros(num_wells))
# P_in: 输入投影 x -> m_i 的扰动
@@ -75,42 +94,49 @@ class Expert(nn.Module):
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 保留
# 可选模态特定编码器（异构专家）
self.encoder = encoder
# 可选 LayerNorm
if use_layer_norm:
self.m_norm = nn.LayerNorm(expert_dim)
self.x_norm = nn.LayerNorm(expert_dim)
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)
_, 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
U(m) = 0.5||m||^2 - 0.5 * c * sum_k softplus(a_k @ m + b_k)
∇U(m) = m - 0.5 * c * 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)
sig = torch.sigmoid(am) # (batch, num_wells)
grad_wells = torch.matmul(sig, self.well_a) # (batch, expert_dim)
return m - 0.5 * grad_wells
return m - 0.5 * self.well_coeff * grad_wells
def deriv(self, m: torch.Tensor, w: torch.Tensor, x_feat: torch.Tensor) -> torch.Tensor:
"""ODE 右端项 dm/dt = -∇U + J·w + P_in·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: torch.Tensor, w: torch.Tensor, x: torch.Tensor,
dt: float, noise_std: float) -> tuple[torch.Tensor, torch.Tensor]:
"""一步 ODE 积分（Euler 法）
"""一步 ODE 积分（Euler 或 RK4）
Args:
m: (batch, expert_dim) 当前状态
@@ -123,19 +149,30 @@ class Expert(nn.Module):
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)
# 可选编码器：把输入投影到专家内部空间
if self.encoder is not None:
x_feat = self.encoder(x)
if self.use_layer_norm:
x_feat = self.x_norm(x_feat)
else:
x_feat = x # P_in 会处理
noise = torch.randn_like(m) * noise_std if noise_std > 0 else 0.0
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
dm = -grad_U + w_proj + x_proj + noise
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) # (batch, workspace_dim)
contribution = self.P_out(m_next)
return m_next, contribution
@@ -147,9 +184,11 @@ class JSpaceWorkspace(nn.Module):
α_i = softmax(<q, P_i_out(m_i)>) q = MLP(x, w)
这个 α_i 是"注意力"——决定哪个专家的内容进入工作空间。
可选 LayerNorm 防止 workspace 长期衰减。
"""
def __init__(self, workspace_dim: int, input_dim: int, num_experts: int):
def __init__(self, workspace_dim: int, input_dim: int, num_experts: int,
use_layer_norm: bool = True):
super().__init__()
self.workspace_dim = workspace_dim
@@ -159,6 +198,9 @@ class JSpaceWorkspace(nn.Module):
nn.Tanh(),
nn.Linear(32, workspace_dim),
)
self.use_layer_norm = use_layer_norm
if use_layer_norm:
self.ln = nn.LayerNorm(workspace_dim)
def forward(self, w: torch.Tensor, x: torch.Tensor,
contributions: torch.Tensor, dt: float,
@@ -176,22 +218,15 @@ class JSpaceWorkspace(nn.Module):
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 = F.softmax(scores, dim=-1)
aggregated = (alpha.unsqueeze(-1) * contributions).sum(dim=1)
# 加权聚合
# 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
if self.use_layer_norm:
w_next = self.ln(w_next)
return w_next, alpha
@@ -219,6 +254,9 @@ class JSpaceModel(nn.Module):
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,
well_coeff=config.well_coeff,
)
for _ in range(config.num_experts)
])
@@ -227,12 +265,13 @@ class JSpaceModel(nn.Module):
workspace_dim=config.workspace_dim,
input_dim=config.input_dim,
num_experts=config.num_experts,
use_layer_norm=config.use_layer_norm,
)
# 输出门控：R(w) → action（这里 action = 预测的下一时刻输入）
self.predictor = nn.Sequential(
nn.Linear(config.workspace_dim, 32),
nn.Tanh(),
nn.GELU(),
nn.Linear(32, config.input_dim),
)
@@ -244,22 +283,25 @@ class JSpaceModel(nn.Module):
for _ in range(self.config.num_experts)],
}
def step(self, state: dict, x: torch.Tensor) -> tuple[dict, torch.Tensor, torch.Tensor, torch.Tensor]:
def step(self, state: dict, x: torch.Tensor,
record_trajectory: bool = False) -> tuple[dict, torch.Tensor, torch.Tensor, torch.Tensor, list]:
"""单时间步前向
Args:
state: {'w': ..., 'm': [...]}
x: (batch, input_dim) 输入
record_trajectory: 是否记录每个子步的 w（J-lens 用）
Returns:
new_state: 更新后的状态
pred: (batch, input_dim) 预测的下一时刻输入
alpha: (batch, num_experts) 注意力权重（可解释性）
w_norm: (batch,) 工作空间范数（输出门控信号）
new_state, pred, alpha, w_norm, w_trajectory
"""
w = state['w']
ms = state['m']
cfg = self.config
w_trajectory = []
# ODE 子步积分
for _ in range(cfg.ode_steps):
# 1. 每个专家更新
contributions = []
new_ms = []
for i, expert in enumerate(self.experts):
@@ -269,31 +311,34 @@ class JSpaceModel(nn.Module):
)
new_ms.append(m_next)
contributions.append(contrib)
contributions = torch.stack(contributions, dim=1) # (batch, num_experts, workspace_dim)
contributions = torch.stack(contributions, dim=1)
# 2. 工作空间更新
w, alpha = self.workspace(
w, x, contributions,
dt=cfg.dt, tau_w=cfg.tau_w,
)
ms = new_ms
# 3. 输出：预测下一时刻输入（船舶涌出，但这里为了训练简化为每步都预测）
if record_trajectory:
w_trajectory.append(w.detach())
pred = self.predictor(w)
w_norm = w.norm(dim=-1)
new_state = {'w': w, 'm': ms}
return new_state, pred, alpha, w_norm
return new_state, pred, alpha, w_norm, w_trajectory
def forward(self, xs: torch.Tensor, state: dict | None = None) -> tuple[torch.Tensor, dict]:
def forward(self, xs: torch.Tensor, state: dict | None = None,
record_trajectory: bool = False) -> tuple[torch.Tensor, dict]:
"""
Args:
xs: (batch, T, input_dim) 输入序列
state: 初始状态，None 则初始化
record_trajectory: 是否记录每个时间步每个子步的 w
Returns:
preds: (batch, T, input_dim) 每步对下一时刻的预测
info: 包含注意力、w_norm 等可解释性信息
info: 包含注意力、w_norm、可选 w_trajectory 等可解释性信息
"""
batch_size, T, _ = xs.shape
device = xs.device
@@ -304,17 +349,24 @@ class JSpaceModel(nn.Module):
preds = []
alphas = []
w_norms = []
w_traj_per_step = []
for t in range(T):
state, pred, alpha, w_norm = self.step(state, xs[:, t])
state, pred, alpha, w_norm, w_traj = self.step(
state, xs[:, t], record_trajectory=record_trajectory
)
preds.append(pred)
alphas.append(alpha)
w_norms.append(w_norm)
if record_trajectory and w_traj:
w_traj_per_step.append(torch.stack(w_traj, dim=1)) # (batch, n_substeps, workspace_dim)
preds = torch.stack(preds, dim=1) # (batch, T, input_dim)
preds = torch.stack(preds, dim=1)
info = {
'alpha': torch.stack(alphas, dim=1), # (batch, T, num_experts)
'w_norm': torch.stack(w_norms, dim=1), # (batch, T)
'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'],
}
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 preds, info

208
jspaceai/core_v2.py
View File

@@ -1,208 +0,0 @@
"""
改进版核心架构 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),
}

59
jspaceai/distill_encoder.py
View File

@@ -1,59 +0,0 @@
"""小模型感知编码器——把小模型表征投影到 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)

82
jspaceai/distill_trainer.py
View File

@@ -1,82 +0,0 @@
"""蒸馏训练器——把小模型理解能力迁移到 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

397
main.py
View File

@@ -1,245 +1,242 @@
#!/usr/bin/env python3
"""
JspaceAI —— 主程序入口
JspaceAI —— 全部接入版（具身 Agent + 完整神经系统 + 自主心智）
对比实验：
1. JSpaceModel（工作空间 + J-space 广播 + ODE 动力学）
2. FlatBaseline（同样参数量的扁平 MLP）
接入全部 5 个感官输入 + 4 个输出执行器 + 神经系统（小脑/海马体/基底神经节/中枢神经）
+ 自主心智（好奇心驱动 + 状态持久化 + 自我模型 + 元学习）。
任务：多模态连续时间序列预测（4 种模式切换）
输出：
- 训练/评估 loss 曲线对比
- 注意力热力图（哪个专家在哪个时段被激活）
- 工作空间 ||w|| 演化曲线（输出门控信号）
- 预测 vs 真实序列对比
模式：
--mode test: 子系统自检（权限 + I/O + 模型 + 海马体回忆）
--mode live: 实时具身循环 + 自主心智（默认安全：不控制鼠标键盘）
--mode safe: 同 live 但更保守（更高动作门控阈值）
运行：
python main.py
python main.py --steps 1000 --device cuda
python main.py --mode test
python main.py --mode live --steps 100
python main.py --mode safe --steps 50
python main.py --mode live --steps 200 --unsafe # 允许鼠标键盘输出
再次运行会从上次状态继续（海洋不蒸发）。
"""
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 (
JSpaceConfig, JSpaceModel, FlatBaseline,
ContinuousSequenceTask, Trainer,
MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent,
AutonomousMind, PLATFORM,
get_screen_size, print_permission_guide,
check_camera_permission, check_microphone_permission,
check_input_monitoring_permission,
)
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,
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,
)
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)
def test_subsystems():
"""子系统自检"""
print("=" * 60)
print("JspaceAI 对比实验")
print(f"具身 Agent 子系统测试 | 平台: {PLATFORM}")
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("\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()
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)
sw, sh = get_screen_size()
print(f" 屏幕尺寸: {sw}x{sh}")
# 评估
print()
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 _ 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(n_steps: int, device: str, safe_mode: bool = False, unsafe: bool = False):
"""实时具身循环 + 自主心智"""
print(f"平台: {PLATFORM}")
sw, sh = get_screen_size()
print(f"屏幕: {sw}x{sh}")
allow_output = unsafe and not safe_mode
if safe_mode:
print("\n安全模式：不执行鼠标/键盘动作")
print_permission_guide()
elif not allow_output:
print("\n默认模式：仅感知 + 音频/屏幕输出（不控制鼠标键盘）")
print(" 如需鼠标键盘输出，加 --unsafe")
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=allow_output,
enable_keyboard_output=allow_output,
enable_audio_output=True,
enable_screen_output=not safe_mode,
risk_threshold=0.5 if safe_mode else 0.3,
)
mind = AutonomousMind(agent, save_dir='outputs/mind', device=device)
print("\n" + "=" * 60)
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(f"运行 {n_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']} | "
f"mem {info['memory_count']}")
mind.run(n_steps=n_steps, interval=0.2, save_every=30, on_step=on_step)
print("\n" + mind.introspect())
# 可视化
print()
print("生成可视化...")
if log:
fig, axes = plt.subplots(2, 3, figsize=(18, 10))
steps = [s['step'] for s in log]
# 1. Loss 曲线对比
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
# ||w|| 按模态着色
ax = axes[0, 0]
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.scatter(steps, wn, 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
ax.legend(handles=[Patch(facecolor=c, label=m) for m, c in cm.items()],
fontsize=7)
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))
ax.plot(steps, [s['curiosity'] for s in log], 'r-', label='curiosity')
ax.plot(steps, [s['success'] for s in log], 'g-', label='success')
ax.legend(); ax.set_title('Curiosity & Success')
ax.grid(True, alpha=0.3)
# 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[0, 2]
mc = {}
for m in mods:
mc[m] = mc.get(m, 0) + 1
ax.bar(mc.keys(), mc.values(),
color=[cm.get(m, 'gray') for m in mc.keys()])
ax.set_title('Modality Distribution')
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)
# 自我模型
ax = axes[1, 0]
fc = log[-1]['self_confidence']
ax.barh(list(fc.keys()), list(fc.values()),
color=plt.cm.RdYlGn(list(fc.values())))
ax.set_xlim(0, 1); ax.set_title('Self Model (confidence)')
# 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()
# 世界模型 loss
ax = axes[1, 1]
ax.plot(steps, [s['world_loss'] for s in log], 'orange')
ax.set_title('World Model Loss'); ax.grid(True, alpha=0.3)
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)
# 记忆数
ax = axes[1, 2]
ax.plot(steps, [s['memory_count'] for s in log], 'teal')
ax.set_title('Hippocampus Memory Count'); ax.grid(True, 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. 预测对比: 模式切换处哪个模型更鲁棒？")
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():
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()
p = argparse.ArgumentParser(description='JspaceAI 全部接入版（具身 + 自主心智）')
p.add_argument('--mode', default='test', choices=['test', 'live', 'safe'],
help='运行模式: test=自检, live=实时循环, safe=安全模式')
p.add_argument('--steps', type=int, default=100, help='live/safe 模式步数')
p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
p.add_argument('--unsafe', action='store_true',
help='允许鼠标键盘输出（默认禁用）')
args = p.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
dev = args.device
if dev == 'auto':
dev = 'cuda' if torch.cuda.is_available() else (
'mps' if torch.backends.mps.is_available() else 'cpu')
run_experiment(
n_steps=args.steps,
device=device,
outdir=Path(args.outdir),
)
if args.mode == 'test':
test_subsystems()
elif args.mode == 'live':
live(args.steps, dev, safe_mode=False, unsafe=args.unsafe)
elif args.mode == 'safe':
live(args.steps, dev, safe_mode=True, unsafe=False)
if __name__ == '__main__':

97
main_autonomous.py
View File

@@ -1,97 +0,0 @@
#!/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()

235
main_chat.py Normal file
View File

@@ -0,0 +1,235 @@
#!/usr/bin/env python3
"""
JspaceAI —— 对话版（只有语言，控制台交互）
基于字符级语言模型 + 自主进化（EWC + 经验回放 + 专家可塑性）。
交互式对话：用户输入文本，模型生成回复，同时持续学习。
模式：
--mode chat: 交互对话（默认）
--mode train: 先在 Shakespeare 语料上预训练若干步，再进入对话
--mode generate: 给定提示词一次性生成
运行：
python main_chat.py
python main_chat.py --mode train --steps 100
python main_chat.py --mode generate --prompt "To be"
"""
from __future__ import annotations
import argparse
import torch
import numpy as np
from pathlib import Path
from jspaceai import (
LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
CharTokenizer, load_shakespeare,
)
def get_config(vocab_size: int) -> LanguageConfig:
return LanguageConfig(
vocab_size=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,
use_rk4=True, use_layer_norm=True,
)
def load_or_init_model(device: str):
"""加载已保存的模型或初始化新模型"""
text = load_shakespeare()
tokenizer = CharTokenizer.from_text(text)
config = get_config(tokenizer.vocab_size)
model = JSpaceLanguageModel(config).to(device)
mp = Path('outputs/chat_model.pt')
if mp.exists():
try:
ckpt = torch.load(mp, map_location=device, weights_only=False)
model.load_state_dict(ckpt['model'])
print(f"已加载模型: {mp}（上次保存的对话状态）")
except Exception:
print("模型加载失败，全新初始化")
else:
print("全新初始化（首次对话）")
return model, config, tokenizer, text
def save_model(model, config, tokenizer):
Path('outputs').mkdir(exist_ok=True)
torch.save({
'model': model.state_dict(),
'config': config,
'tokenizer_chars': tokenizer.chars,
}, 'outputs/chat_model.pt')
def train(model, tokenizer, text, n_steps: int, device: str):
"""在 Shakespeare 语料上预训练"""
print("\n" + "=" * 60)
print(f"预训练 {n_steps} 步（Shakespeare 语料）")
print("=" * 60)
config = model.config
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",
)
save_model(model, config, tokenizer)
print(f"\n模型已保存: outputs/chat_model.pt")
def generate_response(model, tokenizer, prompt: str, n_new: int = 100,
temperature: float = 0.8, top_k: int = 5) -> str:
"""生成回复"""
# 把用户输入编码（未知字符用 0）
prompt_ids = tokenizer.encode(prompt)
if not prompt_ids:
prompt_ids = [0]
generated = model.generate(
prompt_ids, n_new=n_new, temperature=temperature, top_k=top_k,
)
return tokenizer.decode(generated)
def online_learn(model, config, tokenizer, text: str, user_input: str, device: str):
"""在线学习用户输入 + 回放一段 Shakespeare 防遗忘"""
import torch.nn.functional as F
# 把用户输入作为新语料学习
user_tokens = tokenizer.encode(user_input)
if len(user_tokens) < 4:
return # 太短不学
# 重复用户输入凑够 seq_len
seq_len = 48
if len(user_tokens) < seq_len:
user_tokens = user_tokens * (seq_len // len(user_tokens) + 1)
user_seq = user_tokens[:seq_len]
token_tensor = torch.tensor([user_seq], dtype=torch.long).to(device)
# forward + next-token loss
model.train()
model.zero_grad()
logits, _ = model(token_tensor)
pred = logits[:, :-1]
target = token_tensor[:, 1:]
loss = F.cross_entropy(
pred.reshape(-1, config.vocab_size),
target.reshape(-1),
)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# 手动 SGD step（EvolutionTrainer 内部有 EWC，这里简化用直接 step）
with torch.no_grad():
for p in model.parameters():
if p.grad is not None:
p -= 5e-3 * p.grad
model.eval()
return loss.item()
def chat(model, config, tokenizer, text, device: str):
"""交互对话循环"""
print("\n" + "=" * 60)
print("JspaceAI 对话模式")
print("=" * 60)
print("输入文本与模型对话，模型会持续学习你的输入。")
print("命令: /quit 退出 /save 保存 /reset 重置 /train N 预训练N步")
print("=" * 60 + "\n")
model.eval()
while True:
try:
user_input = input("你: ").strip()
except (EOFError, KeyboardInterrupt):
print("\n再见")
break
if not user_input:
continue
if user_input.startswith('/'):
cmd = user_input.lower()
if cmd in ('/quit', '/exit', '/q'):
save_model(model, config, tokenizer)
print(f"(状态已保存到 outputs/chat_model.pt)")
print("再见")
break
elif cmd == '/save':
save_model(model, config, tokenizer)
print(f"(已保存到 outputs/chat_model.pt)")
continue
elif cmd == '/reset':
model = JSpaceLanguageModel(config).to(device)
print("(模型已重置为随机初始化)")
continue
elif cmd.startswith('/train'):
parts = cmd.split()
n = int(parts[1]) if len(parts) > 1 else 50
train(model, tokenizer, text, n, device)
model.eval()
continue
else:
print("未知命令。可用: /quit /save /reset /train N")
continue
# 在线学习用户输入
loss = online_learn(model, config, tokenizer, text, user_input, device)
# 生成回复
response = generate_response(
model, tokenizer, user_input,
n_new=80, temperature=0.8, top_k=5,
)
print(f"AI: {response}")
if loss is not None:
print(f" (学习 loss={loss:.3f})")
def generate_once(model, tokenizer, prompt: str, n_new: int = 200):
"""一次性生成"""
model.eval()
response = generate_response(model, tokenizer, prompt, n_new=n_new,
temperature=0.7, top_k=5)
print(f"提示: {prompt}")
print(f"生成: {response}")
def main():
p = argparse.ArgumentParser(description='JspaceAI 对话版')
p.add_argument('--mode', default='chat',
choices=['chat', 'train', 'generate'],
help='运行模式: chat=交互, train=预训练, generate=一次性生成')
p.add_argument('--steps', type=int, default=100, help='train 模式步数')
p.add_argument('--prompt', default='To be', help='generate 模式提示词')
p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
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')
model, config, tokenizer, text = load_or_init_model(dev)
if args.mode == 'chat':
chat(model, config, tokenizer, text, dev)
elif args.mode == 'train':
train(model, tokenizer, text, args.steps, dev)
elif args.mode == 'generate':
generate_once(model, tokenizer, args.prompt)
if __name__ == '__main__':
main()

212
main_embodied.py
View File

@@ -1,212 +0,0 @@
#!/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()

295
main_full_sensory.py
View File

@@ -1,295 +0,0 @@
#!/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()

222
main_language.py
View File

@@ -1,222 +0,0 @@
#!/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()

296
main_multimodal.py
View File

@@ -1,296 +0,0 @@
#!/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()

438
main_v2.py
View File

@@ -1,438 +0,0 @@
#!/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()