feat: JspaceAI 自主智慧架构
基于第一性原理 + Anthropic 2026 J-space 论文实现的具身智慧系统。 核心架构: - ODE 动力系统 + 并行专家 + J-space 工作空间广播 - 12 个异构专家(视觉/屏幕/听觉/语言/鼠标/跨模态) - workspace 256 维 + LayerNorm + RK4 积分 自主心智(最重要的能力): - 好奇心驱动探索(内在奖励 + 世界模型) - 跨会话状态持久化(海洋不蒸发) - 自我模型(知道自己会什么不会什么) - 元学习(自适应学习率 + 策略选择) 具身 Agent(完整神经系统): - 感知层:摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标 - 大脑皮层(workspace)+ 小脑(运动控制)+ 中枢神经(门控) - 海马体(情景记忆)+ 基底神经节(动作选择) - 执行器:鼠标控制 + 键盘输出 + 音频播放 + 屏幕绘制 多模态支持: - 原生图像/音频/视频/文本/键盘/鼠标 6 种模态 - 跨平台(macOS/Windows/Linux) 外挂模块系统: - 可热插拔的外部能力(小模型/知识库/工具) - 核心心智不依赖外挂,断开后继续工作 守护进程: - 用户主动 start/stop(不自启) - 后台静默运行,持续感知学习 - 状态自动保存,跨会话继续 J-lens 可解释性: - 观测模型内部每个 ODE 子步的想法 - Directed Modulation 验证 workspace 因果作用 - Selectivity 验证(ablate workspace) 小模型蒸馏: - 接 GPT-2/Qwen 等迁移理解能力 - 蒸馏完成后小模型可断开 验证结果: - 连续序列:JSpace 胜 Flat 39.7% - 语言进化:loss 3.95→2.40 - workspace ||w||:v1 0.05 → v2 16.0 - 实时五通道感知 + 具身闭环运行
This commit is contained in:
36
.gitignore
vendored
Normal file
36
.gitignore
vendored
Normal file
@@ -0,0 +1,36 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# 项目运行产物
|
||||
outputs/
|
||||
*.pt
|
||||
*.pth
|
||||
*.npz
|
||||
|
||||
# 自主心智状态(用户私有)
|
||||
~/.jspaceai/
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# 日志
|
||||
*.log
|
||||
logs/
|
||||
375
README.md
Normal file
375
README.md
Normal file
@@ -0,0 +1,375 @@
|
||||
# JspaceAI
|
||||
|
||||
**全局工作空间 + J-space 广播的智慧系统**
|
||||
|
||||
两个版本:
|
||||
1. **连续序列版** (`main.py`):验证架构在多模态序列预测上的优势
|
||||
2. **语言版 + 自主进化** (`main_language.py`):字符级语言建模 + 边推理边学习
|
||||
|
||||
基于第一性原理推导的智慧架构:不是 Transformer,不是 RNN,而是 **ODE 动力系统 + 并行专家 + J-space 工作空间广播 + 在线进化**。
|
||||
|
||||
## 核心思想
|
||||
|
||||
智慧 = 系统内部维持对世界的模型 M,可操作、可修正、可指导行动。
|
||||
|
||||
从第一性原理推导出 M 必须满足的性质,每一性质对应一个实现选择:
|
||||
|
||||
| 性质 | 实现 |
|
||||
|---|---|
|
||||
| 高维 | 向量空间 $\mathbb{R}^n$ |
|
||||
| 内部动力学 | ODE $\dot{m} = -\nabla U(m) + \text{broadcast}$ |
|
||||
| 输入敏感但不决定 | 加性扰动 |
|
||||
| 可读出 | 流形投影 + 阈值门控 |
|
||||
| 可修正 | 参数化 + backprop |
|
||||
| 工作空间结构 | 分块 + J-space 广播 |
|
||||
|
||||
## 架构
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────┐
|
||||
│ 专家池 E_1..E_5(并行) │
|
||||
│ 每个 m_i 有多井势能 → 内部"思考" │
|
||||
│ J_i 稀疏 Jacobian → 选择性接收广播 │
|
||||
└─────────────────────────────────────┘
|
||||
↓ 写入(注意力 α_i 选择)
|
||||
┌─────────────────────────────────────┐
|
||||
│ 工作空间 w(J-space) │
|
||||
│ τ_w · dw/dt = -w + Σ α_i · P_i(m_i) │
|
||||
│ 有自己的动力学 → 持续演化 │
|
||||
└─────────────────────────────────────┘
|
||||
↓ 读出
|
||||
预测头 Q(w) → 预测下一时刻世界
|
||||
```
|
||||
|
||||
**与 Transformer 的对比**:
|
||||
|
||||
| | Transformer | JspaceAI |
|
||||
|---|---|---|
|
||||
| 内部状态 | 离散、无状态 | 连续 ODE 动力学 |
|
||||
| 工作空间 | 隐式(residual stream) | 显式分块 + 广播 |
|
||||
| 专家协作 | 无(或 MoE 但孤立) | J-space 路由 |
|
||||
| 学习信号 | 下一 token | 下一时刻世界状态 |
|
||||
| 自发输出 | 否(被输入触发) | 是(阈值门控) |
|
||||
|
||||
## 安装
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 运行
|
||||
|
||||
### 连续序列版(验证架构)
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
python main.py --steps 2000 --device cuda
|
||||
```
|
||||
|
||||
### v2 版(J-lens + Selectivity + Directed Modulation)
|
||||
|
||||
```bash
|
||||
python main_v2.py
|
||||
python main_v2.py --steps 200 --device mps
|
||||
```
|
||||
|
||||
基于 Anthropic 2026 J-space 论文优化,新增:
|
||||
- **J-lens**:观测模型内部每个 ODE 子步的"想法"
|
||||
- **Directed Modulation**:注入概念向量到 workspace,定向改变输出
|
||||
- **Selectivity 验证**:ablate workspace,对比自动任务 vs 记忆任务
|
||||
|
||||
### 多模态版(摄像头 + 麦克风 + 扬声器)
|
||||
|
||||
```bash
|
||||
python main_multimodal.py --mode train --steps 300
|
||||
python main_multimodal.py --mode live --steps 50
|
||||
python main_multimodal.py --mode eval
|
||||
```
|
||||
|
||||
原生支持图像/音频/视频/文本四种模态,8 个专家分工。
|
||||
|
||||
### 全感官版(摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标)
|
||||
|
||||
```bash
|
||||
# 测试所有 I/O 通道
|
||||
python main_full_sensory.py --mode test
|
||||
|
||||
# 实时全感官感知循环
|
||||
python main_full_sensory.py --mode live --steps 50
|
||||
```
|
||||
|
||||
接入全部 5 个输入通道,12 个专家分工:
|
||||
- **摄像头**(OpenCV):环境视觉
|
||||
- **麦克风**(PortAudio):听觉输入
|
||||
- **屏幕**(mss):屏幕捕获,看到自己被显示在哪里
|
||||
- **键盘**(pynput):用户输入感知
|
||||
- **鼠标**(pynput):用户注意力追踪
|
||||
- **扬声器**(PortAudio):音频输出
|
||||
|
||||
**macOS 权限**:键盘/鼠标监听需要在 系统设置 → 隐私与安全 → 辅助功能 中授权终端。
|
||||
|
||||
### 具身 Agent 版(完整神经系统)
|
||||
|
||||
```bash
|
||||
# 测试各子系统
|
||||
python main_embodied.py --mode test
|
||||
|
||||
# 实时具身循环(鼠标键盘音频屏幕全输出)
|
||||
python main_embodied.py --mode live --steps 50
|
||||
|
||||
# 安全模式(只感知不执行鼠标键盘动作)
|
||||
python main_embodied.py --mode safe --steps 30
|
||||
```
|
||||
|
||||
完整的人类神经系统对应:
|
||||
|
||||
| 人类系统 | 功能 | 实现 |
|
||||
|---|---|---|
|
||||
| 大脑皮层 | 认知、语言 | workspace + 12 专家 |
|
||||
| 小脑 | 运动协调 | Cerebellum(前向模型+逆模型) |
|
||||
| 中枢神经 | 动作门控 | CentralNervousSystem(反射弧+抑制) |
|
||||
| 海马体 | 情景记忆 | Hippocampus(向量检索) |
|
||||
| 基底神经节 | 动作选择 | BasalGanglia(Q-learning) |
|
||||
| 眼/耳 | 视听觉 | 摄像头+麦克风 |
|
||||
| 口 | 说话 | AudioActuator(扬声器) |
|
||||
| 手 | 操作 | MouseActuator + KeyboardActuator |
|
||||
| 皮肤 | 触觉 | (预留接口) |
|
||||
|
||||
感知-思考-行动闭环:
|
||||
1. 感知(五感)→ 2. 思考(workspace)→ 3. 回忆(海马体)
|
||||
→ 4. 决策(基底神经节)→ 5. 精化(小脑)→ 6. 门控(中枢神经)
|
||||
→ 7. 行动(手足口)→ 8. 学习(更新所有系统)
|
||||
|
||||
### 守护进程(后台持续运行)
|
||||
|
||||
```bash
|
||||
# 启动后台守护进程(静默运行,不干扰使用)
|
||||
python daemon.py start
|
||||
|
||||
# 低功耗模式(更长间隔,适合长期后台)
|
||||
python daemon.py start --low-power --interval 2.0
|
||||
|
||||
# 前台运行(调试用,可看实时输出)
|
||||
python daemon.py start --fg
|
||||
|
||||
# 查看状态
|
||||
python daemon.py status
|
||||
|
||||
# 查看日志
|
||||
python daemon.py log -n 30
|
||||
|
||||
# 心智自我认知
|
||||
python daemon.py introspect
|
||||
|
||||
# 停止
|
||||
python daemon.py stop
|
||||
```
|
||||
|
||||
守护进程特性:
|
||||
- **用户主动控制**:手动 start/stop,不开机自启
|
||||
- **静默后台运行**:不弹窗、不发声、不控制鼠标键盘
|
||||
- **持续感知学习**:感知屏幕/键盘/鼠标/摄像头/麦克风 + 好奇心驱动学习
|
||||
- **状态持久化**:停止后状态保存,再次启动从上次继续
|
||||
- **跨会话持续**:海洋不蒸发,每次启动还是同一个"自己"
|
||||
- **崩溃恢复**:异常不崩溃,等待后继续
|
||||
|
||||
状态保存在 `~/.jspaceai/`:
|
||||
- `state/`:workspace + 海马体 + 基底神经节 + 自我模型
|
||||
- `logs/daemon.log`:运行日志
|
||||
- `daemon.pid`:进程 PID
|
||||
- `status.json`:实时状态
|
||||
|
||||
## 外挂模块系统(可热插拔)
|
||||
|
||||
```python
|
||||
from jspaceai import ModuleDock, SmallModelModule, KnowledgeBaseModule, ToolModule
|
||||
|
||||
dock = ModuleDock(workspace_dim=64)
|
||||
dock.register('llm', SmallModelModule(model_name='gpt2'))
|
||||
dock.register('kb', KnowledgeBaseModule())
|
||||
dock.register('tools', ToolModule())
|
||||
|
||||
# 热插上
|
||||
dock.connect('llm')
|
||||
dock.connect('tools')
|
||||
|
||||
# 使用
|
||||
result = dock.query('llm', 'hello world')
|
||||
calc = dock.query('tools', {'tool': 'calc', 'args': ['1+1']}) # → 2
|
||||
|
||||
# 随时拔掉(心智不停,只是该能力不可用)
|
||||
dock.disconnect('llm')
|
||||
dock.query('llm', 'hello') # → None
|
||||
|
||||
# 状态查询
|
||||
dock.status() # 各模块连接状态
|
||||
```
|
||||
|
||||
三种内置外挂:
|
||||
- **SmallModelModule**:外挂小模型(GPT-2/Qwen 等),提供语言能力
|
||||
- **KnowledgeBaseModule**:向量检索知识库
|
||||
- **ToolModule**:可调用工具(计算器、搜索等)
|
||||
|
||||
核心原则:**心智不依赖外挂**。断开任何外挂,心智继续工作,只是"知道得少"。
|
||||
|
||||
## 改进版 v2(异构专家 + 大 workspace + 蒸馏)
|
||||
|
||||
四大改进:
|
||||
|
||||
1. **异构专家**:视觉/听觉/语言/跨模态专家有不同架构(不再全是相同 ODE)
|
||||
2. **workspace 扩容**:64 → 256 维,加 LayerNorm 防衰减
|
||||
3. **RK4 积分**:比 Euler 稳定(||w|| 从 0.05 提升到 16.0)
|
||||
4. **小模型蒸馏**:接 GPT-2/Qwen 等,迁移理解能力
|
||||
|
||||
```python
|
||||
from jspaceai import JSpaceConfigV2, JSpaceModelV2, SmallModelEncoder, DistillationTrainer
|
||||
|
||||
config = JSpaceConfigV2(workspace_dim=256, num_experts=12)
|
||||
model = JSpaceModelV2(config)
|
||||
encoder = SmallModelEncoder(input_dim=config.input_dim, model_name='gpt2')
|
||||
trainer = DistillationTrainer(model, encoder, texts=['hello', 'world', ...])
|
||||
trainer.train(n_steps=100)
|
||||
# 蒸馏后小模型可断开,编码器已有理解能力
|
||||
```
|
||||
|
||||
对比 v1 vs v2:
|
||||
|
||||
| 指标 | v1 | v2 |
|
||||
|---|---|---|
|
||||
| workspace_dim | 64 | 256 |
|
||||
| expert_dim | 16 | 64 |
|
||||
| \|\|w\|\|(实测) | 0.05-0.35 | 12-16 |
|
||||
| ODE 积分 | Euler | RK4 |
|
||||
| 专家架构 | 全相同 | 异构 |
|
||||
| LayerNorm | 无 | 有 |
|
||||
| 小模型蒸馏 | 无 | 有 |
|
||||
|
||||
## 跨平台支持
|
||||
|
||||
支持 macOS / Windows / Linux:
|
||||
|
||||
| 功能 | macOS | Windows | Linux |
|
||||
|---|---|---|---|
|
||||
| 摄像头 | OpenCV | OpenCV | OpenCV |
|
||||
| 麦克风 | PortAudio | PortAudio | PortAudio |
|
||||
| 屏幕 | mss | mss | mss/X11 |
|
||||
| 键盘监听 | pynput(需辅助功能权限) | pynput | pynput(需 X11) |
|
||||
| 鼠标监听 | pynput(需辅助功能权限) | pynput | pynput(需 X11) |
|
||||
| 鼠标控制 | pynput | pynput | pynput |
|
||||
| 键盘控制 | pynput | pynput | pynput |
|
||||
| 音频播放 | PortAudio | PortAudio | PortAudio |
|
||||
|
||||
## 自主进化机制
|
||||
|
||||
语言版不只是"训练一次",而是**天生支持持续进化**:
|
||||
|
||||
1. **在线学习**:每个 forward 都累积梯度并更新参数,模型永不"训练完成"
|
||||
2. **EWC(Elastic Weight Consolidation)**:保护重要参数防灾难性遗忘
|
||||
- 周期性计算 Fisher 信息,锚定重要参数
|
||||
- 新知识只能修改不重要的参数
|
||||
3. **经验回放(Experience Replay)**:定期回放旧序列,防遗忘
|
||||
4. **专家可塑性(Expert Plasticity)**:追踪专家专业化,鼓励新知识流向空闲专家
|
||||
|
||||
进化循环可以**无限运行**——持续喂入新文本,模型持续进化。
|
||||
|
||||
## 输出
|
||||
|
||||
### 连续序列版
|
||||
- `outputs/experiment.png` — Loss/注意力/||w||/预测对比
|
||||
- `outputs/models.pt` — 模型权重
|
||||
|
||||
### v2 版
|
||||
- `outputs/experiment_v2.png` — Loss/专家使用率/||w||/J-lens 读出/Modulation/Selectivity
|
||||
- `outputs/model_v2.pt` — 模型 + J-lens 权重
|
||||
|
||||
## 基于 Anthropic 2026 论文的优化
|
||||
|
||||
v2 基于 [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) 的发现:
|
||||
|
||||
1. **J-lens 探针**:学习线性映射 L: workspace → vocab,观测每个子步 workspace "准备说什么"
|
||||
2. **三层分层**:sensory(子步0)→ workspace(子步1-2)→ motor(子步3)
|
||||
3. **Directed Modulation**:注入概念向量 v_concept 到 workspace,验证 workspace 是因果中介
|
||||
4. **Selectivity 验证**:ablate workspace top-k J-lens 方向,对比自动任务 vs 记忆任务
|
||||
|
||||
**验证结果**:
|
||||
- Directed Modulation 成功定向改变输出(注入 'o' → 输出 'o' 增多)
|
||||
- J-lens 读出在 motor 子步最集中(符合 motor = 输出准备)
|
||||
- Selectivity 符合预期:简单任务 ablate 影响小(与 Anthropic 发现一致)
|
||||
|
||||
## 实验结果
|
||||
|
||||
### 连续序列版
|
||||
- JSpaceModel eval MSE: 0.0234
|
||||
- FlatBaseline eval MSE: 0.0387
|
||||
- **JSpace 胜出 39.7%**(同样参数量)
|
||||
- 注意力热力图显示专家自发分工
|
||||
|
||||
### 语言版
|
||||
- Loss 从 3.95 → 2.40,下降 39%
|
||||
- 生成从纯乱码进化为类语言结构(含换行、常见词片段)
|
||||
- 专家按字符类型分工(元音/辅音/结构字符)
|
||||
|
||||
## 关键观察点
|
||||
|
||||
跑完后看 `experiment.png`,关注:
|
||||
|
||||
1. **Loss 曲线**:JSpace 是否收敛更快/更低?
|
||||
2. **注意力热力图**:4 种模式时段是否激活不同的专家?(如果是,说明 J-space 路由学到了分工)
|
||||
3. **||w|| 曲线**:模式切换时是否有尖峰?(如果是,说明工作空间在"感知到"环境变化)
|
||||
4. **预测对比**:模式切换处哪个模型更鲁棒?
|
||||
|
||||
## 文件结构
|
||||
|
||||
```
|
||||
JspaceAI/
|
||||
├── main.py # 主程序:对比实验
|
||||
├── requirements.txt
|
||||
├── README.md
|
||||
└── jspaceai/
|
||||
├── __init__.py
|
||||
├── core.py # 核心:Expert + JSpaceWorkspace + JSpaceModel
|
||||
├── baselines.py # 对比基线:FlatBaseline
|
||||
├── task.py # 玩具任务:多模态连续序列
|
||||
└── trainer.py # 训练器:预测学习
|
||||
```
|
||||
|
||||
## 数学细节
|
||||
|
||||
### 专家动力学
|
||||
|
||||
$$\dot{m}_i = -\nabla U_i(m_i) + J_i \cdot w + P_i^{in} \cdot x + \xi$$
|
||||
|
||||
势能 $U_i(m_i) = \frac{1}{2}\|m_i\|^2 - \frac{1}{2}\sum_k \text{softplus}(a_k \cdot m_i + b_k)$
|
||||
|
||||
多井结构让专家在吸引子之间漫游 = 内部"思考"。
|
||||
|
||||
### 工作空间动力学
|
||||
|
||||
$$\tau_w \cdot \dot{w} = -w + \sum_i \alpha_i \cdot P_i^{out}(m_i)$$
|
||||
|
||||
$$\alpha_i = \text{softmax}(\langle q, P_i^{out}(m_i) \rangle), \quad q = \text{MLP}(x, w)$$
|
||||
|
||||
### J-space 路由
|
||||
|
||||
$J_i$ 是稀疏矩阵(top-k 选择),每个专家只对工作空间的少数方向敏感 → 选择性广播。
|
||||
|
||||
### 学习目标
|
||||
|
||||
$$\mathcal{L} = \mathbb{E}\left[\|x_{t+\Delta t} - Q(w_t)\|^2\right]$$
|
||||
|
||||
自监督——只需时间序列,不需标签。
|
||||
|
||||
## 局限性
|
||||
|
||||
这是**最小验证版本**,不是生产系统:
|
||||
|
||||
- 玩具任务(8 维序列),不是语言/图像
|
||||
- ODE 用 Euler 积分(精度低,但简单可 backprop)
|
||||
- 专家数 5,不是大规模
|
||||
- 无持续跨会话状态(每次 forward 从零初始化)
|
||||
|
||||
**下一步扩展方向**:
|
||||
1. 用 adjoint method 替换 Euler(Neural ODE 路线)
|
||||
2. 加跨会话状态持久化(真正的"海洋")
|
||||
3. Scale 到更大任务(MNIST 连续预测 → 语言建模)
|
||||
4. 加输出门控的真正"自发输出"(非每步都预测)
|
||||
461
daemon.py
Normal file
461
daemon.py
Normal file
@@ -0,0 +1,461 @@
|
||||
"""
|
||||
守护进程 —— 自主心智后台持续运行
|
||||
|
||||
用户主动控制:
|
||||
python daemon.py start # 启动后台守护进程
|
||||
python daemon.py stop # 停止
|
||||
python daemon.py status # 查看状态
|
||||
python daemon.py log # 查看最近日志
|
||||
python daemon.py fg # 前台运行(调试用)
|
||||
|
||||
守护进程特性:
|
||||
- 静默后台运行,不弹窗不干扰
|
||||
- 持续感知(屏幕/键盘/鼠标/摄像头/麦克风)
|
||||
- 持续学习(好奇心驱动 + 状态保存)
|
||||
- 信号处理(SIGTERM 优雅退出)
|
||||
- 崩溃恢复(自动重启)
|
||||
- 低资源占用(可配置 CPU/内存限制)
|
||||
|
||||
状态保存在 ~/.jspaceai/,每次运行从上次状态继续。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import signal
|
||||
import time
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from jspaceai import (
|
||||
MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent,
|
||||
AutonomousMind, PLATFORM,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 路径配置
|
||||
# ============================================================
|
||||
|
||||
APP_DIR = Path.home() / ".jspaceai"
|
||||
LOG_DIR = APP_DIR / "logs"
|
||||
PID_FILE = APP_DIR / "daemon.pid"
|
||||
STATE_DIR = APP_DIR / "state"
|
||||
LOG_FILE = LOG_DIR / "daemon.log"
|
||||
STATUS_FILE = APP_DIR / "status.json"
|
||||
|
||||
APP_DIR.mkdir(exist_ok=True)
|
||||
LOG_DIR.mkdir(exist_ok=True)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 日志
|
||||
# ============================================================
|
||||
|
||||
def get_logger(foreground: bool = False) -> logging.Logger:
|
||||
logger = logging.getLogger("jspaceai-daemon")
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.handlers.clear()
|
||||
fmt = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S')
|
||||
fh = logging.FileHandler(LOG_FILE, encoding='utf-8')
|
||||
fh.setFormatter(fmt)
|
||||
logger.addHandler(fh)
|
||||
if foreground:
|
||||
sh = logging.StreamHandler()
|
||||
sh.setFormatter(fmt)
|
||||
logger.addHandler(sh)
|
||||
return logger
|
||||
|
||||
|
||||
# ============================================================
|
||||
# PID 管理
|
||||
# ============================================================
|
||||
|
||||
def write_pid(pid: int):
|
||||
PID_FILE.write_text(str(pid))
|
||||
|
||||
def read_pid() -> Optional[int]:
|
||||
if PID_FILE.exists():
|
||||
try:
|
||||
return int(PID_FILE.read_text().strip())
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
def clear_pid():
|
||||
if PID_FILE.exists():
|
||||
PID_FILE.unlink()
|
||||
|
||||
def is_running(pid: int) -> bool:
|
||||
"""检查进程是否存活"""
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
return True
|
||||
except (OSError, ProcessLookupError):
|
||||
return False
|
||||
|
||||
def write_status(status: dict):
|
||||
STATUS_FILE.write_text(json.dumps(status, indent=2, default=str))
|
||||
|
||||
def read_status() -> Optional[dict]:
|
||||
if STATUS_FILE.exists():
|
||||
try:
|
||||
return json.loads(STATUS_FILE.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 守护进程主循环
|
||||
# ============================================================
|
||||
|
||||
class MindDaemon:
|
||||
"""自主心智守护进程"""
|
||||
|
||||
def __init__(self, device: str = 'cpu', low_power: bool = False,
|
||||
step_interval: float = 0.5):
|
||||
self.device = device
|
||||
self.low_power = low_power
|
||||
self.step_interval = step_interval
|
||||
self.logger = get_logger(foreground=False)
|
||||
self.running = False
|
||||
self.mind: Optional[AutonomousMind] = None
|
||||
self.agent: Optional[EmbodiedAgent] = None
|
||||
self.start_time = time.time()
|
||||
self.step_count_session = 0
|
||||
|
||||
def initialize(self):
|
||||
"""初始化心智"""
|
||||
self.logger.info(f"初始化 | 平台: {PLATFORM} | 设备: {self.device}")
|
||||
|
||||
config = MultimodalConfig(
|
||||
vocab_size=50, embed_dim=16, input_dim=8,
|
||||
workspace_dim=64, expert_dim=24, num_experts=12,
|
||||
num_wells=4, ode_steps=3, dt=0.1, tau_w=0.3,
|
||||
jacobian_sparsity=16, noise_std=0.01,
|
||||
img_size=32, audio_frame_size=1024, keyboard_vocab=128,
|
||||
)
|
||||
model = MultimodalJSpaceModel(config).to(self.device)
|
||||
model.eval()
|
||||
|
||||
# 完全静默:不控制鼠标键盘,不发声,不弹窗
|
||||
self.agent = EmbodiedAgent(
|
||||
model, device=self.device,
|
||||
enable_mouse_output=False,
|
||||
enable_keyboard_output=False,
|
||||
enable_audio_output=False,
|
||||
enable_screen_output=False,
|
||||
)
|
||||
self.mind = AutonomousMind(
|
||||
self.agent, save_dir=str(STATE_DIR), device=self.device,
|
||||
)
|
||||
self.logger.info(f"就绪 | 历史步数: {self.mind.step_count}")
|
||||
|
||||
def run_forever(self):
|
||||
"""主循环——永不停止,直到收到停止信号"""
|
||||
self.running = True
|
||||
|
||||
# 信号处理
|
||||
def handle_stop(signum, frame):
|
||||
self.logger.info(f"收到信号 {signum},优雅退出...")
|
||||
self.running = False
|
||||
|
||||
signal.signal(signal.SIGTERM, handle_stop)
|
||||
signal.signal(signal.SIGINT, handle_stop)
|
||||
|
||||
self.initialize()
|
||||
|
||||
# 写入运行状态
|
||||
write_status({
|
||||
'running': True,
|
||||
'pid': os.getpid(),
|
||||
'started_at': time.time(),
|
||||
'device': self.device,
|
||||
'low_power': self.low_power,
|
||||
'step_count': self.mind.step_count,
|
||||
})
|
||||
|
||||
self.logger.info("守护进程启动,开始持续感知学习")
|
||||
last_save = time.time()
|
||||
save_interval = 30 # 每 30 秒保存一次
|
||||
|
||||
try:
|
||||
while self.running:
|
||||
try:
|
||||
info = self.mind.step()
|
||||
self.step_count_session += 1
|
||||
|
||||
# 每 50 步记录一次日志
|
||||
if self.step_count_session % 50 == 0:
|
||||
self.logger.info(
|
||||
f"step {info['step']} | mod {info['modality']} | "
|
||||
f"||w|| {info['w_norm']:.3f} | "
|
||||
f"curio {info['curiosity']:.3f} | "
|
||||
f"success {info['success']:.2f} | "
|
||||
f"mem {info['memory_count']}"
|
||||
)
|
||||
# 更新状态文件
|
||||
status = read_status() or {}
|
||||
status.update({
|
||||
'step_count': info['step'],
|
||||
'last_step_at': time.time(),
|
||||
'last_modality': info['modality'],
|
||||
'last_w_norm': info['w_norm'],
|
||||
'last_curiosity': info['curiosity'],
|
||||
'memory_count': info['memory_count'],
|
||||
'session_steps': self.step_count_session,
|
||||
})
|
||||
write_status(status)
|
||||
|
||||
# 定期保存
|
||||
if time.time() - last_save > save_interval:
|
||||
self.mind.save_state()
|
||||
last_save = time.time()
|
||||
|
||||
# 控制频率
|
||||
time.sleep(self.step_interval)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
break
|
||||
except Exception as e:
|
||||
self.logger.error(f"步进异常: {e}\n{__import__('traceback').format_exc()}")
|
||||
# 等待后继续,不崩溃
|
||||
time.sleep(5.0)
|
||||
|
||||
finally:
|
||||
# 清理
|
||||
self.logger.info("保存最终状态...")
|
||||
self.mind.save_state()
|
||||
if self.agent:
|
||||
self.agent.senses.stop()
|
||||
self.logger.info(
|
||||
f"守护进程退出 | 本次步数: {self.step_count_session} | "
|
||||
f"总步数: {self.mind.step_count} | "
|
||||
f"运行时长: {time.time() - self.start_time:.0f}s"
|
||||
)
|
||||
write_status({
|
||||
'running': False,
|
||||
'stopped_at': time.time(),
|
||||
'step_count': self.mind.step_count,
|
||||
'session_steps': self.step_count_session,
|
||||
})
|
||||
clear_pid()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 命令处理
|
||||
# ============================================================
|
||||
|
||||
def cmd_start(device: str = 'cpu', low_power: bool = False,
|
||||
interval: float = 0.5, foreground: bool = False):
|
||||
"""启动守护进程"""
|
||||
existing_pid = read_pid()
|
||||
if existing_pid and is_running(existing_pid):
|
||||
print(f"守护进程已在运行 (PID {existing_pid})")
|
||||
print(f"查看状态: python daemon.py status")
|
||||
return
|
||||
|
||||
if foreground:
|
||||
# 前台运行(调试用)
|
||||
print(f"前台模式启动(Ctrl+C 退出)")
|
||||
daemon = MindDaemon(device=device, low_power=low_power,
|
||||
step_interval=interval)
|
||||
daemon.logger = get_logger(foreground=True)
|
||||
daemon.run_forever()
|
||||
return
|
||||
|
||||
# 后台 fork
|
||||
print(f"启动后台守护进程...")
|
||||
daemon_script = Path(__file__).resolve()
|
||||
cmd = [sys.executable, str(daemon_script), "_run",
|
||||
"--device", device,
|
||||
"--interval", str(interval)]
|
||||
if low_power:
|
||||
cmd.append("--low-power")
|
||||
|
||||
# 用 subprocess 启动,脱离终端
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=open(LOG_DIR / 'stdout.log', 'a'),
|
||||
stderr=subprocess.STDOUT,
|
||||
stdin=subprocess.DEVNULL,
|
||||
start_new_session=True, # 脱离父进程
|
||||
)
|
||||
|
||||
write_pid(proc.pid)
|
||||
|
||||
# 等待一下确认启动
|
||||
time.sleep(2)
|
||||
if proc.poll() is None:
|
||||
print(f"守护进程已启动 (PID {proc.pid})")
|
||||
print(f"日志: {LOG_FILE}")
|
||||
print(f"停止: python daemon.py stop")
|
||||
print(f"状态: python daemon.py status")
|
||||
else:
|
||||
print(f"启动失败,查看日志: {LOG_FILE}")
|
||||
clear_pid()
|
||||
|
||||
|
||||
def cmd_stop():
|
||||
"""停止守护进程"""
|
||||
pid = read_pid()
|
||||
if not pid:
|
||||
print("守护进程未运行")
|
||||
return
|
||||
|
||||
if not is_running(pid):
|
||||
print(f"进程 {pid} 已不存在,清理 PID 文件")
|
||||
clear_pid()
|
||||
return
|
||||
|
||||
print(f"发送停止信号到 PID {pid}...")
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
# 等待退出
|
||||
for _ in range(10):
|
||||
time.sleep(0.5)
|
||||
if not is_running(pid):
|
||||
break
|
||||
if is_running(pid):
|
||||
print("强制终止...")
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
time.sleep(1)
|
||||
clear_pid()
|
||||
print("守护进程已停止")
|
||||
except Exception as e:
|
||||
print(f"停止失败: {e}")
|
||||
|
||||
|
||||
def cmd_status():
|
||||
"""查看状态"""
|
||||
pid = read_pid()
|
||||
running = pid and is_running(pid)
|
||||
|
||||
print("=" * 50)
|
||||
print(f"JspaceAI 守护进程状态")
|
||||
print("=" * 50)
|
||||
|
||||
if running:
|
||||
print(f"状态: 运行中 (PID {pid})")
|
||||
else:
|
||||
print(f"状态: 已停止")
|
||||
if pid:
|
||||
print(f" (PID {pid} 已不存在)")
|
||||
|
||||
status = read_status()
|
||||
if status:
|
||||
print(f"总步数: {status.get('step_count', '?')}")
|
||||
if 'started_at' in status and running:
|
||||
uptime = time.time() - status['started_at']
|
||||
print(f"运行时长: {uptime:.0f}s ({uptime/3600:.1f}h)")
|
||||
if 'last_modality' in status:
|
||||
print(f"最后模态: {status['last_modality']}")
|
||||
if 'last_w_norm' in status:
|
||||
print(f"||w||: {status['last_w_norm']:.3f}")
|
||||
if 'last_curiosity' in status:
|
||||
print(f"好奇心: {status['last_curiosity']:.3f}")
|
||||
if 'memory_count' in status:
|
||||
print(f"记忆数: {status['memory_count']}")
|
||||
if 'session_steps' in status:
|
||||
print(f"本次会话步数: {status['session_steps']}")
|
||||
if not running and 'stopped_at' in status:
|
||||
print(f"停止时间: {time.ctime(status['stopped_at'])}")
|
||||
|
||||
print(f"\n日志: {LOG_FILE}")
|
||||
print(f"状态目录: {STATE_DIR}")
|
||||
|
||||
|
||||
def cmd_log(n_lines: int = 30):
|
||||
"""查看最近日志"""
|
||||
if not LOG_FILE.exists():
|
||||
print("无日志")
|
||||
return
|
||||
lines = LOG_FILE.read_text().strip().split('\n')
|
||||
for line in lines[-n_lines:]:
|
||||
print(line)
|
||||
|
||||
|
||||
def cmd_introspect():
|
||||
"""内省——查看心智的自我认知"""
|
||||
status = read_status()
|
||||
if not status:
|
||||
print("无状态数据")
|
||||
return
|
||||
|
||||
pid = read_pid()
|
||||
if pid and is_running(pid):
|
||||
print("心智正在运行,自我认知:")
|
||||
else:
|
||||
print("心智已停止,最后的自我认知:")
|
||||
|
||||
print(json.dumps(status, indent=2, default=str))
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 入口
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
p = argparse.ArgumentParser(description='JspaceAI 自主心智守护进程')
|
||||
sub = p.add_subparsers(dest='command')
|
||||
|
||||
# start
|
||||
sp = sub.add_parser('start', help='启动后台守护进程')
|
||||
sp.add_argument('--device', default='cpu')
|
||||
sp.add_argument('--low-power', action='store_true', help='低功耗模式(更长间隔)')
|
||||
sp.add_argument('--interval', type=float, default=0.5, help='步进间隔(秒)')
|
||||
sp.add_argument('--fg', action='store_true', help='前台运行(调试用)')
|
||||
|
||||
# stop
|
||||
sub.add_parser('stop', help='停止守护进程')
|
||||
|
||||
# status
|
||||
sub.add_parser('status', help='查看状态')
|
||||
|
||||
# log
|
||||
sp = sub.add_parser('log', help='查看日志')
|
||||
sp.add_argument('-n', type=int, default=30, help='行数')
|
||||
|
||||
# introspect
|
||||
sub.add_parser('introspect', help='心智自我认知')
|
||||
|
||||
# _run(内部命令,被 start 调用)
|
||||
sp = sub.add_parser('_run', help='内部运行命令')
|
||||
sp.add_argument('--device', default='cpu')
|
||||
sp.add_argument('--low-power', action='store_true')
|
||||
sp.add_argument('--interval', type=float, default=0.5)
|
||||
|
||||
args = p.parse_args()
|
||||
|
||||
if args.command == 'start':
|
||||
cmd_start(args.device, args.low_power, args.interval, args.fg)
|
||||
elif args.command == 'stop':
|
||||
cmd_stop()
|
||||
elif args.command == 'status':
|
||||
cmd_status()
|
||||
elif args.command == 'log':
|
||||
cmd_log(args.n)
|
||||
elif args.command == 'introspect':
|
||||
cmd_introspect()
|
||||
elif args.command == '_run':
|
||||
# 内部运行模式
|
||||
daemon = MindDaemon(
|
||||
device=args.device,
|
||||
low_power=args.low_power,
|
||||
step_interval=args.interval,
|
||||
)
|
||||
daemon.run_forever()
|
||||
else:
|
||||
p.print_help()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
129
jspaceai/__init__.py
Normal file
129
jspaceai/__init__.py
Normal file
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
JspaceAI —— 全局工作空间 + J-space 广播的智慧系统
|
||||
|
||||
模块:
|
||||
1. core.py: 连续序列版(Expert + JSpaceWorkspace + JSpaceModel)
|
||||
2. language_model.py: 语言版 + 自主进化
|
||||
3. jlens.py: J-lens 可解释性工具
|
||||
4. multimodal.py: 多模态(图像/音频/视频/文本)
|
||||
5. realtime.py: 实时 I/O(摄像头/麦克风/扬声器)
|
||||
6. evolution.py: 自主进化训练器
|
||||
"""
|
||||
from .core import (
|
||||
Expert,
|
||||
JSpaceWorkspace,
|
||||
JSpaceModel,
|
||||
JSpaceConfig,
|
||||
)
|
||||
from .baselines import FlatBaseline
|
||||
from .task import ContinuousSequenceTask
|
||||
from .trainer import Trainer
|
||||
|
||||
from .language_data import CharTokenizer, CharDataset, load_shakespeare
|
||||
from .language_model import (
|
||||
LanguageConfig,
|
||||
JSpaceLanguageModel,
|
||||
ExperienceReplay,
|
||||
EWCOptimizer,
|
||||
ExpertPlasticity,
|
||||
)
|
||||
from .jlens import (
|
||||
JLensConfig,
|
||||
JLensProbe,
|
||||
JLensSuite,
|
||||
WorkspaceAblator,
|
||||
DirectedModulation,
|
||||
CounterfactualReflection,
|
||||
)
|
||||
from .multimodal import (
|
||||
MultimodalConfig,
|
||||
MultimodalJSpaceModel,
|
||||
VisualEncoder, VisualDecoder,
|
||||
AudioEncoder, AudioDecoder,
|
||||
TextEncoder, TextDecoder,
|
||||
)
|
||||
from .realtime import (
|
||||
Frame,
|
||||
CameraStream,
|
||||
MicrophoneStream,
|
||||
AudioPlayer,
|
||||
MultimodalStream,
|
||||
SensoryMotorLoop,
|
||||
)
|
||||
from .desktop import (
|
||||
InputEvent,
|
||||
ScreenCapture,
|
||||
KeyboardMonitor,
|
||||
MouseMonitor,
|
||||
DesktopStream,
|
||||
FullSensoryStream,
|
||||
)
|
||||
from .platform import (
|
||||
PlatformInfo, PLATFORM,
|
||||
get_screen_size,
|
||||
check_camera_permission, check_microphone_permission,
|
||||
check_input_monitoring_permission, print_permission_guide,
|
||||
)
|
||||
from .embodied import (
|
||||
MouseActuator, KeyboardActuator, AudioActuator, ScreenActuator,
|
||||
Cerebellum, CentralNervousSystem, Hippocampus, BasalGanglia,
|
||||
EmbodiedAgent,
|
||||
)
|
||||
from .autonomous import (
|
||||
CuriosityDrive, PersistentState, SelfModel, MetaLearner,
|
||||
AutonomousMind,
|
||||
)
|
||||
from .modules import (
|
||||
ExternalModule, SmallModelModule, KnowledgeBaseModule,
|
||||
ToolModule, ModuleDock,
|
||||
)
|
||||
from .core_v2 import JSpaceConfigV2, JSpaceModelV2
|
||||
from .distill_encoder import SmallModelEncoder
|
||||
from .distill_trainer import DistillationTrainer
|
||||
from .evolution import EvolutionTrainer
|
||||
|
||||
__all__ = [
|
||||
# 核心架构
|
||||
"Expert", "JSpaceWorkspace", "JSpaceModel", "JSpaceConfig",
|
||||
# 对比基线
|
||||
"FlatBaseline",
|
||||
# 连续序列任务
|
||||
"ContinuousSequenceTask", "Trainer",
|
||||
# 语言建模
|
||||
"CharTokenizer", "CharDataset", "load_shakespeare",
|
||||
"LanguageConfig", "JSpaceLanguageModel",
|
||||
"ExperienceReplay", "EWCOptimizer", "ExpertPlasticity",
|
||||
# J-lens 可解释性
|
||||
"JLensConfig", "JLensProbe", "JLensSuite",
|
||||
"WorkspaceAblator", "DirectedModulation", "CounterfactualReflection",
|
||||
# 多模态
|
||||
"MultimodalConfig", "MultimodalJSpaceModel",
|
||||
"VisualEncoder", "VisualDecoder",
|
||||
"AudioEncoder", "AudioDecoder",
|
||||
"TextEncoder", "TextDecoder",
|
||||
# 实时 I/O
|
||||
"Frame", "CameraStream", "MicrophoneStream", "AudioPlayer",
|
||||
"MultimodalStream", "SensoryMotorLoop",
|
||||
# 桌面 I/O
|
||||
"InputEvent", "ScreenCapture", "KeyboardMonitor", "MouseMonitor",
|
||||
"DesktopStream", "FullSensoryStream",
|
||||
# 平台抽象
|
||||
"PlatformInfo", "PLATFORM", "get_screen_size",
|
||||
"check_camera_permission", "check_microphone_permission",
|
||||
"check_input_monitoring_permission", "print_permission_guide",
|
||||
# 具身 Agent(输出执行器 + 神经系统)
|
||||
"MouseActuator", "KeyboardActuator", "AudioActuator", "ScreenActuator",
|
||||
"Cerebellum", "CentralNervousSystem", "Hippocampus", "BasalGanglia",
|
||||
"EmbodiedAgent",
|
||||
# 自主心智(最重要的能力)
|
||||
"CuriosityDrive", "PersistentState", "SelfModel", "MetaLearner",
|
||||
"AutonomousMind",
|
||||
# 外挂模块系统(可热插拔)
|
||||
"ExternalModule", "SmallModelModule", "KnowledgeBaseModule",
|
||||
"ToolModule", "ModuleDock",
|
||||
# 改进版核心 v2
|
||||
"JSpaceConfigV2", "JSpaceModelV2",
|
||||
"SmallModelEncoder", "DistillationTrainer",
|
||||
# 自主进化
|
||||
"EvolutionTrainer",
|
||||
]
|
||||
389
jspaceai/autonomous.py
Normal file
389
jspaceai/autonomous.py
Normal file
@@ -0,0 +1,389 @@
|
||||
"""
|
||||
自主心智(AutonomousMind)—— 智慧最重要的能力
|
||||
|
||||
四个核心能力:
|
||||
1. CuriosityDrive: 好奇心驱动的主动探索(内在奖励)
|
||||
2. PersistentState: 跨会话状态持久化(海洋不蒸发)
|
||||
3. SelfModel: 自我模型(知道自己会什么不会什么)
|
||||
4. MetaLearner: 元学习(学会如何学习)
|
||||
|
||||
合起来 = AutonomousMind,永不停止的自主进化。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass
|
||||
from collections import deque
|
||||
|
||||
|
||||
class CuriosityDrive(nn.Module):
|
||||
"""好奇心驱动——基于预测误差的内在奖励
|
||||
|
||||
模型有世界模型预测下一状态,预测误差=好奇心=内在奖励。
|
||||
模型被驱动去探索"预测不准"的区域。学会后好奇心降低,转向新区域。
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim: int, action_dim: int = 5, hidden_dim: int = 64):
|
||||
super().__init__()
|
||||
self.workspace_dim = workspace_dim
|
||||
self.world_model = nn.Sequential(
|
||||
nn.Linear(workspace_dim + action_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, workspace_dim),
|
||||
)
|
||||
self.state_history: deque = deque(maxlen=500)
|
||||
self.prediction_error_ema = 0.1
|
||||
|
||||
def predict_next(self, w, action):
|
||||
return self.world_model(torch.cat([w, action], dim=-1))
|
||||
|
||||
def compute_curiosity(self, w_current, action, w_next):
|
||||
with torch.no_grad():
|
||||
w_pred = self.predict_next(w_current, action)
|
||||
pred_error = F.mse_loss(w_pred, w_next).item()
|
||||
w_np = w_current[0].cpu().numpy()
|
||||
novelty = self._compute_novelty(w_np)
|
||||
progress = max(0, pred_error - self.prediction_error_ema * 0.9)
|
||||
self.prediction_error_ema = 0.95 * self.prediction_error_ema + 0.05 * pred_error
|
||||
curiosity = progress + 0.3 * novelty
|
||||
self.state_history.append(w_np.copy())
|
||||
return curiosity
|
||||
|
||||
def _compute_novelty(self, w):
|
||||
if len(self.state_history) < 5:
|
||||
return 1.0
|
||||
history = list(self.state_history)[-100:]
|
||||
distances = [np.linalg.norm(w - h) for h in history]
|
||||
return float(min(1.0, min(distances) / 2.0))
|
||||
|
||||
def train_world_model(self, w_current, action, w_next):
|
||||
w_pred = self.predict_next(w_current, action.detach())
|
||||
return F.mse_loss(w_pred, w_next.detach())
|
||||
|
||||
|
||||
class PersistentState:
|
||||
"""跨会话状态持久化——让海洋不蒸发
|
||||
|
||||
保存 workspace + 专家状态 + 海马体 + 基底神经节 + 好奇心历史 + 自我模型。
|
||||
"""
|
||||
|
||||
def __init__(self, save_dir: Path):
|
||||
self.save_dir = Path(save_dir)
|
||||
self.save_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.state_file = self.save_dir / 'mind_state.json'
|
||||
self.tensors_file = self.save_dir / 'mind_tensors.npz'
|
||||
|
||||
def save(self, state: dict):
|
||||
tensors = {}
|
||||
if 'w' in state:
|
||||
tensors['w'] = state['w'].cpu().numpy()
|
||||
if 'm' in state:
|
||||
for i, m in enumerate(state['m']):
|
||||
if m is not None:
|
||||
tensors[f'm_{i}'] = m.cpu().numpy()
|
||||
if 'basal_ganglia' in state and state['basal_ganglia'] is not None:
|
||||
tensors['bg_weights'] = state['basal_ganglia']
|
||||
if 'curiosity_history' in state:
|
||||
tensors['curiosity_history'] = np.array(state['curiosity_history'])
|
||||
if tensors:
|
||||
np.savez(self.tensors_file, **tensors)
|
||||
|
||||
json_state = {
|
||||
'step_count': state.get('step_count', 0),
|
||||
'total_runtime': state.get('total_runtime', 0.0),
|
||||
'self_model': state.get('self_model', {}),
|
||||
'saved_at': time.time(),
|
||||
}
|
||||
self.state_file.write_text(json.dumps(json_state, indent=2, default=str))
|
||||
|
||||
def load(self) -> Optional[dict]:
|
||||
if not self.state_file.exists():
|
||||
return None
|
||||
result = {}
|
||||
if self.tensors_file.exists():
|
||||
data = np.load(self.tensors_file, allow_pickle=True)
|
||||
if 'w' in data:
|
||||
result['w'] = torch.tensor(data['w'])
|
||||
ms = {}
|
||||
for key in data.files:
|
||||
if key.startswith('m_'):
|
||||
idx = int(key.split('_')[1])
|
||||
ms[idx] = torch.tensor(data[key])
|
||||
if ms:
|
||||
result['m'] = [ms[i] for i in sorted(ms.keys())]
|
||||
if 'bg_weights' in data:
|
||||
result['basal_ganglia'] = data['bg_weights']
|
||||
if 'curiosity_history' in data:
|
||||
result['curiosity_history'] = data['curiosity_history'].tolist()
|
||||
json_state = json.loads(self.state_file.read_text())
|
||||
result.update(json_state)
|
||||
return result
|
||||
|
||||
|
||||
class SelfModel:
|
||||
"""自我模型——知道自己会什么、不会什么
|
||||
|
||||
对每个能力域维护置信度(0-1),通过历史成功率更新。
|
||||
不知道时主动学习(好奇心驱动)。
|
||||
"""
|
||||
|
||||
def __init__(self, capabilities=None):
|
||||
if capabilities is None:
|
||||
capabilities = ['visual', 'audio', 'text', 'motor_mouse',
|
||||
'motor_keyboard', 'memory', 'prediction']
|
||||
self.capabilities = capabilities
|
||||
self.confidence = {c: 0.0 for c in capabilities}
|
||||
self.attempts = {c: 0 for c in capabilities}
|
||||
self.recent_results = {c: deque(maxlen=20) for c in capabilities}
|
||||
|
||||
def record_attempt(self, capability, success):
|
||||
if capability not in self.confidence:
|
||||
return
|
||||
self.attempts[capability] += 1
|
||||
self.recent_results[capability].append(success)
|
||||
recent = list(self.recent_results[capability])
|
||||
if recent:
|
||||
weights = np.linspace(0.5, 1.0, len(recent))
|
||||
self.confidence[capability] = float(np.average(recent, weights=weights))
|
||||
|
||||
def get_weakness(self):
|
||||
return min(self.confidence, key=self.confidence.get)
|
||||
|
||||
def get_strength(self):
|
||||
return max(self.confidence, key=self.confidence.get)
|
||||
|
||||
def knows(self, capability, threshold=0.5):
|
||||
return self.confidence.get(capability, 0.0) > threshold
|
||||
|
||||
def summary(self):
|
||||
return {
|
||||
'capabilities': dict(self.confidence),
|
||||
'attempts': dict(self.attempts),
|
||||
'strength': self.get_strength(),
|
||||
'weakness': self.get_weakness(),
|
||||
}
|
||||
|
||||
|
||||
class MetaLearner:
|
||||
"""元学习——学会如何学习
|
||||
|
||||
每个能力域有自适应学习率:进步快→加快,停滞→减慢。
|
||||
记录什么学习策略有效。
|
||||
"""
|
||||
|
||||
def __init__(self, capabilities=None):
|
||||
if capabilities is None:
|
||||
capabilities = ['visual', 'audio', 'text', 'motor_mouse',
|
||||
'motor_keyboard', 'memory', 'prediction']
|
||||
self.learning_rates = {c: 1e-3 for c in capabilities}
|
||||
self.loss_history = {c: deque(maxlen=20) for c in capabilities}
|
||||
self.strategy_scores = {
|
||||
'predict_next': 0.5, 'replay': 0.5,
|
||||
'explore': 0.5, 'imitate': 0.5,
|
||||
}
|
||||
|
||||
def get_lr(self, capability):
|
||||
return self.learning_rates.get(capability, 1e-3)
|
||||
|
||||
def record_loss(self, capability, loss):
|
||||
if capability not in self.loss_history:
|
||||
return
|
||||
self.loss_history[capability].append(loss)
|
||||
history = list(self.loss_history[capability])
|
||||
if len(history) < 5:
|
||||
return
|
||||
recent_avg = np.mean(history[-5:])
|
||||
old_avg = np.mean(history[-10:-5]) if len(history) >= 10 else recent_avg
|
||||
improvement = (old_avg - recent_avg) / max(old_avg, 1e-8)
|
||||
lr = self.learning_rates[capability]
|
||||
if improvement > 0.05:
|
||||
lr *= 1.1
|
||||
elif improvement < 0.01:
|
||||
lr *= 0.9
|
||||
self.learning_rates[capability] = max(1e-5, min(1e-2, lr))
|
||||
|
||||
def best_strategy(self):
|
||||
return max(self.strategy_scores, key=self.strategy_scores.get)
|
||||
|
||||
def reward_strategy(self, strategy, reward):
|
||||
if strategy in self.strategy_scores:
|
||||
self.strategy_scores[strategy] = (
|
||||
0.9 * self.strategy_scores[strategy] + 0.1 * reward
|
||||
)
|
||||
|
||||
|
||||
class AutonomousMind:
|
||||
"""自主心智——永不停止的自主进化
|
||||
|
||||
整合好奇心 + 持久化 + 自我模型 + 元学习。
|
||||
核心循环:感知→自我评估→好奇探索→行动→观察→学习→存盘
|
||||
"""
|
||||
|
||||
def __init__(self, agent, save_dir='outputs/mind', device='cpu'):
|
||||
self.agent = agent
|
||||
self.device = device
|
||||
self.config = agent.config
|
||||
|
||||
self.curiosity = CuriosityDrive(
|
||||
workspace_dim=self.config.workspace_dim, action_dim=5,
|
||||
).to(device)
|
||||
self.persistence = PersistentState(Path(save_dir))
|
||||
self.self_model = SelfModel()
|
||||
self.meta_learner = MetaLearner()
|
||||
|
||||
self.step_count = 0
|
||||
self.total_runtime = 0.0
|
||||
self.start_time = time.time()
|
||||
self.running = False
|
||||
self.curiosity_history = deque(maxlen=1000)
|
||||
self.curiosity_optimizer = torch.optim.Adam(
|
||||
self.curiosity.parameters(), lr=1e-3
|
||||
)
|
||||
|
||||
self._load_state()
|
||||
|
||||
def _load_state(self):
|
||||
state = self.persistence.load()
|
||||
if state is None:
|
||||
print(" [心智] 全新启动")
|
||||
return
|
||||
print(f" [心智] 恢复状态: step={state.get('step_count', 0)}")
|
||||
if 'w' in state:
|
||||
self.agent.state['w'] = state['w'].to(self.device)
|
||||
if 'm' in state:
|
||||
for i, m in enumerate(state['m']):
|
||||
if m is not None and i < len(self.agent.state['m']):
|
||||
self.agent.state['m'][i] = m.to(self.device)
|
||||
if 'basal_ganglia' in state and hasattr(self.agent, 'basal_ganglia'):
|
||||
self.agent.basal_ganglia.action_weights = state['basal_ganglia']
|
||||
self.step_count = state.get('step_count', 0)
|
||||
self.total_runtime = state.get('total_runtime', 0.0)
|
||||
|
||||
def save_state(self):
|
||||
state = {
|
||||
'w': self.agent.state['w'],
|
||||
'm': self.agent.state['m'],
|
||||
'basal_ganglia': getattr(self.agent.basal_ganglia, 'action_weights', None)
|
||||
if hasattr(self.agent, 'basal_ganglia') else None,
|
||||
'curiosity_history': list(self.curiosity.state_history),
|
||||
'step_count': self.step_count,
|
||||
'total_runtime': self.total_runtime + (time.time() - self.start_time),
|
||||
'self_model': {'confidence': self.self_model.confidence},
|
||||
}
|
||||
self.persistence.save(state)
|
||||
|
||||
def step(self) -> dict:
|
||||
# 1. 感知 + 思考
|
||||
sensory_data = self.agent.perceive()
|
||||
w_before = self.agent.state['w'].clone()
|
||||
weakness = self.self_model.get_weakness()
|
||||
strength = self.self_model.get_strength()
|
||||
|
||||
# 2. 行动
|
||||
action_info = self.agent.decide_and_act(w_before, sensory_data.get('modality', 'idle'))
|
||||
w_after, modality = self.agent.think(sensory_data)
|
||||
|
||||
# 3. 好奇心
|
||||
action_tensor = torch.tensor(action_info['action_params'],
|
||||
dtype=torch.float32).unsqueeze(0).to(self.device)
|
||||
curiosity_reward = self.curiosity.compute_curiosity(w_before, action_tensor, w_after)
|
||||
self.curiosity_history.append(curiosity_reward)
|
||||
|
||||
# 4. 训练世界模型
|
||||
world_loss = self.curiosity.train_world_model(w_before, action_tensor, w_after)
|
||||
self.curiosity_optimizer.zero_grad()
|
||||
world_loss.backward()
|
||||
self.curiosity_optimizer.step()
|
||||
|
||||
# 5. 评估成功度
|
||||
w_stability = 1.0 - min(1.0, abs(w_after.norm().item() - w_before.norm().item()))
|
||||
success = (0.3 * float(action_info['executed']) +
|
||||
0.4 * min(1.0, curiosity_reward) + 0.3 * w_stability)
|
||||
|
||||
# 6. 更新自我模型
|
||||
cap_map = {'image': 'visual', 'screen': 'visual', 'audio': 'audio',
|
||||
'text': 'text', 'keyboard': 'text', 'mouse': 'motor_mouse', 'idle': 'prediction'}
|
||||
cap = cap_map.get(modality, 'prediction')
|
||||
self.self_model.record_attempt(cap, success)
|
||||
self.self_model.record_attempt('prediction', 1.0 - min(1.0, world_loss.item()))
|
||||
|
||||
# 7. 元学习
|
||||
self.meta_learner.record_loss(cap, world_loss.item())
|
||||
if curiosity_reward > 0.3:
|
||||
self.meta_learner.reward_strategy('explore', curiosity_reward)
|
||||
|
||||
# 8. 记忆 + 基底神经节
|
||||
self.agent.remember(w_after, {
|
||||
'modality': modality, 'curiosity': curiosity_reward,
|
||||
'success': success, 'step': self.step_count,
|
||||
})
|
||||
self.agent.learn(w_after, np.array(action_info['action_params']), reward=curiosity_reward)
|
||||
|
||||
self.step_count += 1
|
||||
return {
|
||||
'step': self.step_count, 'modality': modality,
|
||||
'w_norm': w_after.norm().item(), 'curiosity': curiosity_reward,
|
||||
'world_loss': world_loss.item(), 'success': success,
|
||||
'weakness': weakness, 'strength': strength,
|
||||
'self_confidence': dict(self.self_model.confidence),
|
||||
'best_strategy': self.meta_learner.best_strategy(),
|
||||
'memory_count': self.agent.hippocampus.size() if self.agent.hippocampus else 0,
|
||||
}
|
||||
|
||||
def run(self, n_steps=100, interval=0.2, save_every=50, on_step=None):
|
||||
self.running = True
|
||||
self.start_time = time.time()
|
||||
self.agent.senses.start()
|
||||
print(f"\n自主心智启动 | 总步数: {self.step_count} | 保存间隔: {save_every}")
|
||||
print("=" * 60)
|
||||
|
||||
log = []
|
||||
try:
|
||||
for _ in range(n_steps):
|
||||
if not self.running:
|
||||
break
|
||||
info = self.step()
|
||||
log.append(info)
|
||||
if on_step:
|
||||
on_step(info)
|
||||
elif info['step'] % 10 == 0:
|
||||
print(f" step {info['step']:4d} | mod {info['modality']:8s} | "
|
||||
f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | "
|
||||
f"success {info['success']:.2f} | weak={info['weakness']} | "
|
||||
f"mem {info['memory_count']}")
|
||||
if info['step'] % save_every == 0:
|
||||
self.save_state()
|
||||
time.sleep(interval)
|
||||
except KeyboardInterrupt:
|
||||
print("\n用户中断")
|
||||
finally:
|
||||
self.running = False
|
||||
self.agent.senses.stop()
|
||||
if hasattr(self.agent, 'audio_actuator'):
|
||||
self.agent.audio_actuator.stop()
|
||||
self.save_state()
|
||||
self.total_runtime += time.time() - self.start_time
|
||||
return log
|
||||
|
||||
def introspect(self) -> str:
|
||||
"""内省——自我报告"""
|
||||
sm = self.self_model.summary()
|
||||
avg_curio = np.mean(list(self.curiosity_history)) if self.curiosity_history else 0
|
||||
report = f"=== 自主心智内省 ===\n步数: {self.step_count}\n运行: {self.total_runtime:.0f}s\n"
|
||||
report += f"记忆: {self.agent.hippocampus.size() if self.agent.hippocampus else 0}\n"
|
||||
report += f"平均好奇心: {avg_curio:.3f}\n\n自我认知:\n"
|
||||
for cap, conf in sm['capabilities'].items():
|
||||
bar = '█' * int(conf * 20)
|
||||
report += f" {cap:15s}: {conf:.2f} {bar}\n"
|
||||
report += f"\n最强: {sm['strength']}\n最弱: {sm['weakness']}\n"
|
||||
report += f"最佳策略: {self.meta_learner.best_strategy()}\n"
|
||||
return report
|
||||
41
jspaceai/baselines.py
Normal file
41
jspaceai/baselines.py
Normal file
@@ -0,0 +1,41 @@
|
||||
"""
|
||||
对比基线:扁平网络(无工作空间、无专家分块、无动力学)
|
||||
|
||||
用同样的参数量预算,验证"工作空间 + J-space 广播"架构比扁平 MLP 好。
|
||||
这是关键对比——如果工作空间架构赢不了同样大小的 MLP,那架构本身没意义。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class FlatBaseline(nn.Module):
|
||||
"""
|
||||
扁平 MLP 基线:参数量与 JSpaceModel 相当,但无结构。
|
||||
|
||||
- 无专家分块(单一隐层)
|
||||
- 无工作空间(无广播机制)
|
||||
- 无动力学(每步独立预测,无状态传递)
|
||||
- 无 J-space 路由
|
||||
|
||||
这代表"用同样算力做暴力拟合"的路线。
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim: int = 8, hidden_dim: int = 90, num_layers: int = 2):
|
||||
super().__init__()
|
||||
layers = []
|
||||
in_dim = input_dim
|
||||
for _ in range(num_layers):
|
||||
layers.extend([nn.Linear(in_dim, hidden_dim), nn.Tanh()])
|
||||
in_dim = hidden_dim
|
||||
layers.append(nn.Linear(in_dim, input_dim))
|
||||
self.net = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, xs: torch.Tensor, state=None) -> tuple[torch.Tensor, dict]:
|
||||
"""
|
||||
xs: (batch, T, input_dim)
|
||||
returns: preds (batch, T, input_dim), info (空)
|
||||
"""
|
||||
preds = self.net(xs) # (batch, T, input_dim)
|
||||
return preds, {}
|
||||
320
jspaceai/core.py
Normal file
320
jspaceai/core.py
Normal file
@@ -0,0 +1,320 @@
|
||||
"""
|
||||
核心架构:专家模块 + J-space 工作空间 + ODE 动力学
|
||||
|
||||
数学形式(每个 forward 时间步内做 Euler 积分若干子步):
|
||||
|
||||
专家 i 的状态 m_i:
|
||||
dm_i/dt = -∇U_i(m_i) + J_i · w + P_i_in · x
|
||||
|
||||
U_i(m_i) = ½ ||m_i||² - ½ Σ_k softplus(a_ik · m_ik + b_ik)
|
||||
(多井势能:阻尼项 + softplus 形成的局部吸引子)
|
||||
|
||||
工作空间 w:
|
||||
τ_w · dw/dt = -w + Σ_i α_i · P_i_out(m_i)
|
||||
α_i = softmax(<q, P_i_out(m_i)>) q = MLP(x, w)
|
||||
|
||||
Jacobian 路由 J_i: 稀疏线性映射,每个专家只对 w 的少数维度敏感
|
||||
输出门控:当 ||w|| > θ 时触发输出 R(w)
|
||||
|
||||
所有参数都可 backprop 训练。学习目标是预测下一时刻的输入。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class JSpaceConfig:
|
||||
"""模型超参,全部可调"""
|
||||
input_dim: int = 8 # 输入 x 的维度
|
||||
workspace_dim: int = 32 # 工作空间 w 的维度(J-space)
|
||||
expert_dim: int = 16 # 每个专家内部状态 m_i 的维度
|
||||
num_experts: int = 5 # 专家数量
|
||||
num_wells: int = 4 # 每个专家势能景观的井数
|
||||
ode_steps: int = 4 # 每个时间步内 ODE 积分子步数
|
||||
dt: float = 0.1 # ODE 积分步长
|
||||
tau_w: float = 0.3 # 工作空间时间常数
|
||||
output_threshold: float = 0.5 # 输出门控阈值(船舶涌出阈值)
|
||||
jacobian_sparsity: int = 8 # 每个 J_i 只保留前 k 大的连接
|
||||
noise_std: float = 0.01 # 内部噪声 ξ(t) 的标准差
|
||||
|
||||
|
||||
class Expert(nn.Module):
|
||||
"""
|
||||
单个专家模块。
|
||||
|
||||
状态:m_i ∈ R^{expert_dim}
|
||||
势能:U_i(m_i) = ½||m_i||² - ½ Σ_k softplus(a_k · m_i + b_k) · w_k
|
||||
- ½||m_i||² 是阻尼项(拉回原点)
|
||||
- softplus 项创造多个局部吸引子(多井势能 → 内部"思考")
|
||||
动力学:dm_i/dt = -∇U_i(m_i) + J_i · w + P_in · x + ξ
|
||||
|
||||
J_i 是稀疏 Jacobian:从工作空间 w 路由信息进来。
|
||||
"""
|
||||
|
||||
def __init__(self, expert_dim: int, workspace_dim: int, input_dim: int,
|
||||
num_wells: int, sparsity: int):
|
||||
super().__init__()
|
||||
self.expert_dim = expert_dim
|
||||
self.workspace_dim = workspace_dim
|
||||
self.num_wells = num_wells
|
||||
|
||||
# 势能景观参数:每个井是一个 softplus 形成的吸引子
|
||||
# U(m) = 0.5||m||^2 - 0.5 * sum_k softplus(a_k @ m + b_k)
|
||||
# ∇U(m) = m - 0.5 * sum_k sigmoid(a_k @ m + b_k) * a_k
|
||||
self.well_a = nn.Parameter(torch.randn(num_wells, expert_dim) * 0.3)
|
||||
self.well_b = nn.Parameter(torch.zeros(num_wells))
|
||||
|
||||
# P_in: 输入投影 x -> m_i 的扰动
|
||||
self.P_in = nn.Linear(input_dim, expert_dim, bias=False)
|
||||
|
||||
# P_out: 模块输出到工作空间的投影
|
||||
self.P_out = nn.Linear(expert_dim, workspace_dim, bias=False)
|
||||
|
||||
# J_i: 稀疏 Jacobian,从 w 路由信息到 m_i
|
||||
# 用 top-k 稀疏:训练时学习一个 full matrix,但只激活 top-k
|
||||
self.J_raw = nn.Parameter(torch.randn(expert_dim, workspace_dim) * 0.1)
|
||||
self.sparsity = sparsity
|
||||
|
||||
# 注意:sparsity 通过 forward 时 top-k 选择实现,可微性通过稀疏 mask 保留
|
||||
|
||||
def get_sparse_J(self) -> torch.Tensor:
|
||||
"""获取稀疏化的 Jacobian:每行只保留 top-k 元素"""
|
||||
if self.sparsity >= self.workspace_dim:
|
||||
return self.J_raw
|
||||
# 对每行做 top-k(按绝对值)
|
||||
abs_J = self.J_raw.abs()
|
||||
topk_vals, topk_idx = abs_J.topk(self.sparsity, dim=-1)
|
||||
mask = torch.zeros_like(self.J_raw)
|
||||
mask.scatter_(-1, topk_idx, 1.0)
|
||||
return self.J_raw * mask
|
||||
|
||||
def grad_potential(self, m: torch.Tensor) -> torch.Tensor:
|
||||
"""计算势能梯度 ∇U_i(m)
|
||||
U(m) = 0.5||m||^2 - 0.5 * sum_k softplus(a_k @ m + b_k)
|
||||
∇U(m) = m - 0.5 * sum_k sigmoid(a_k @ m + b_k) * a_k
|
||||
"""
|
||||
# m: (batch, expert_dim)
|
||||
# well_a: (num_wells, expert_dim)
|
||||
# a_k @ m: (batch, num_wells)
|
||||
am = F.linear(m, self.well_a, self.well_b) # (batch, num_wells)
|
||||
sig = torch.sigmoid(am) # (batch, num_wells)
|
||||
# sum_k sigmoid(...) * a_k: (batch, expert_dim)
|
||||
# well_a: (num_wells, expert_dim), sig: (batch, num_wells) -> (batch, 1, num_wells)
|
||||
# 用 matmul: sig @ well_a -> (batch, expert_dim)
|
||||
grad_wells = torch.matmul(sig, self.well_a) # (batch, expert_dim)
|
||||
return m - 0.5 * grad_wells
|
||||
|
||||
def forward(self, m: torch.Tensor, w: torch.Tensor, x: torch.Tensor,
|
||||
dt: float, noise_std: float) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""一步 ODE 积分(Euler 法)
|
||||
|
||||
Args:
|
||||
m: (batch, expert_dim) 当前状态
|
||||
w: (batch, workspace_dim) 工作空间状态
|
||||
x: (batch, input_dim) 输入
|
||||
dt: 步长
|
||||
noise_std: 噪声标准差
|
||||
|
||||
Returns:
|
||||
m_next: (batch, expert_dim) 下一状态
|
||||
contribution: (batch, workspace_dim) 对工作空间的贡献(pre-attention)
|
||||
"""
|
||||
# 动力学: dm/dt = -∇U(m) + J·w + P_in·x + ξ
|
||||
J = self.get_sparse_J() # (expert_dim, workspace_dim)
|
||||
w_proj = F.linear(w, J) # (batch, expert_dim)
|
||||
x_proj = self.P_in(x) # (batch, expert_dim)
|
||||
grad_U = self.grad_potential(m) # (batch, expert_dim)
|
||||
|
||||
noise = torch.randn_like(m) * noise_std if noise_std > 0 else 0.0
|
||||
|
||||
dm = -grad_U + w_proj + x_proj + noise
|
||||
m_next = m + dt * dm
|
||||
|
||||
# 对工作空间的贡献
|
||||
contribution = self.P_out(m_next) # (batch, workspace_dim)
|
||||
return m_next, contribution
|
||||
|
||||
|
||||
class JSpaceWorkspace(nn.Module):
|
||||
"""
|
||||
全局工作空间 w。
|
||||
|
||||
动力学: τ_w · dw/dt = -w + Σ_i α_i · P_i_out(m_i)
|
||||
α_i = softmax(<q, P_i_out(m_i)>) q = MLP(x, w)
|
||||
|
||||
这个 α_i 是"注意力"——决定哪个专家的内容进入工作空间。
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim: int, input_dim: int, num_experts: int):
|
||||
super().__init__()
|
||||
self.workspace_dim = workspace_dim
|
||||
|
||||
# Query 生成器:从 (x, w) 生成 query 向量
|
||||
self.query_gen = nn.Sequential(
|
||||
nn.Linear(input_dim + workspace_dim, 32),
|
||||
nn.Tanh(),
|
||||
nn.Linear(32, workspace_dim),
|
||||
)
|
||||
|
||||
def forward(self, w: torch.Tensor, x: torch.Tensor,
|
||||
contributions: torch.Tensor, dt: float,
|
||||
tau_w: float) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""一步工作空间演化
|
||||
|
||||
Args:
|
||||
w: (batch, workspace_dim)
|
||||
x: (batch, input_dim)
|
||||
contributions: (batch, num_experts, workspace_dim) 各专家的贡献
|
||||
dt: 步长
|
||||
tau_w: 时间常数
|
||||
|
||||
Returns:
|
||||
w_next: (batch, workspace_dim)
|
||||
alpha: (batch, num_experts) 注意力权重(可解释性用)
|
||||
"""
|
||||
# 生成 query
|
||||
q = self.query_gen(torch.cat([x, w], dim=-1)) # (batch, workspace_dim)
|
||||
|
||||
# 计算每个专家的注意力分数
|
||||
# contributions: (batch, num_experts, workspace_dim)
|
||||
# q: (batch, workspace_dim) -> (batch, 1, workspace_dim)
|
||||
scores = (contributions * q.unsqueeze(1)).sum(dim=-1) # (batch, num_experts)
|
||||
alpha = F.softmax(scores, dim=-1) # (batch, num_experts)
|
||||
|
||||
# 加权聚合
|
||||
# alpha: (batch, num_experts, 1) * contributions: (batch, num_experts, workspace_dim)
|
||||
aggregated = (alpha.unsqueeze(-1) * contributions).sum(dim=1) # (batch, workspace_dim)
|
||||
|
||||
# 动力学: τ_w · dw/dt = -w + aggregated
|
||||
dw = (-w + aggregated) / tau_w
|
||||
w_next = w + dt * dw
|
||||
return w_next, alpha
|
||||
|
||||
|
||||
class JSpaceModel(nn.Module):
|
||||
"""
|
||||
完整模型:N 个专家 + 工作空间 + 输出门控 + 预测头。
|
||||
|
||||
forward 流程(每个时间步):
|
||||
1. 每个专家从 (m_i, w, x) 更新 m_i,产出对工作空间的贡献
|
||||
2. 工作空间从 (w, x, contributions) 更新 w
|
||||
3. (可选)当 ||w|| > threshold 时输出 R(w)
|
||||
4. 预测头 Q(w) 预测下一时刻输入
|
||||
|
||||
时间序列处理:对长度 T 的输入序列,依次跑 T 步,返回每步的预测。
|
||||
"""
|
||||
|
||||
def __init__(self, config: JSpaceConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
self.experts = nn.ModuleList([
|
||||
Expert(
|
||||
expert_dim=config.expert_dim,
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_wells=config.num_wells,
|
||||
sparsity=config.jacobian_sparsity,
|
||||
)
|
||||
for _ in range(config.num_experts)
|
||||
])
|
||||
|
||||
self.workspace = JSpaceWorkspace(
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_experts=config.num_experts,
|
||||
)
|
||||
|
||||
# 输出门控:R(w) → action(这里 action = 预测的下一时刻输入)
|
||||
self.predictor = nn.Sequential(
|
||||
nn.Linear(config.workspace_dim, 32),
|
||||
nn.Tanh(),
|
||||
nn.Linear(32, config.input_dim),
|
||||
)
|
||||
|
||||
def init_state(self, batch_size: int, device: torch.device) -> dict:
|
||||
"""初始化内部状态"""
|
||||
return {
|
||||
'w': torch.zeros(batch_size, self.config.workspace_dim, device=device),
|
||||
'm': [torch.zeros(batch_size, self.config.expert_dim, device=device)
|
||||
for _ in range(self.config.num_experts)],
|
||||
}
|
||||
|
||||
def step(self, state: dict, x: torch.Tensor) -> tuple[dict, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""单时间步前向
|
||||
|
||||
Returns:
|
||||
new_state: 更新后的状态
|
||||
pred: (batch, input_dim) 预测的下一时刻输入
|
||||
alpha: (batch, num_experts) 注意力权重(可解释性)
|
||||
w_norm: (batch,) 工作空间范数(输出门控信号)
|
||||
"""
|
||||
w = state['w']
|
||||
ms = state['m']
|
||||
cfg = self.config
|
||||
|
||||
# ODE 子步积分
|
||||
for _ in range(cfg.ode_steps):
|
||||
# 1. 每个专家更新
|
||||
contributions = []
|
||||
new_ms = []
|
||||
for i, expert in enumerate(self.experts):
|
||||
m_next, contrib = expert(
|
||||
ms[i], w, x,
|
||||
dt=cfg.dt, noise_std=cfg.noise_std,
|
||||
)
|
||||
new_ms.append(m_next)
|
||||
contributions.append(contrib)
|
||||
contributions = torch.stack(contributions, dim=1) # (batch, num_experts, workspace_dim)
|
||||
|
||||
# 2. 工作空间更新
|
||||
w, alpha = self.workspace(
|
||||
w, x, contributions,
|
||||
dt=cfg.dt, tau_w=cfg.tau_w,
|
||||
)
|
||||
ms = new_ms
|
||||
|
||||
# 3. 输出:预测下一时刻输入(船舶涌出,但这里为了训练简化为每步都预测)
|
||||
pred = self.predictor(w)
|
||||
w_norm = w.norm(dim=-1)
|
||||
|
||||
new_state = {'w': w, 'm': ms}
|
||||
return new_state, pred, alpha, w_norm
|
||||
|
||||
def forward(self, xs: torch.Tensor, state: dict | None = None) -> tuple[torch.Tensor, dict]:
|
||||
"""
|
||||
Args:
|
||||
xs: (batch, T, input_dim) 输入序列
|
||||
state: 初始状态,None 则初始化
|
||||
|
||||
Returns:
|
||||
preds: (batch, T, input_dim) 每步对下一时刻的预测
|
||||
info: 包含注意力、w_norm 等可解释性信息
|
||||
"""
|
||||
batch_size, T, _ = xs.shape
|
||||
device = xs.device
|
||||
|
||||
if state is None:
|
||||
state = self.init_state(batch_size, device)
|
||||
|
||||
preds = []
|
||||
alphas = []
|
||||
w_norms = []
|
||||
for t in range(T):
|
||||
state, pred, alpha, w_norm = self.step(state, xs[:, t])
|
||||
preds.append(pred)
|
||||
alphas.append(alpha)
|
||||
w_norms.append(w_norm)
|
||||
|
||||
preds = torch.stack(preds, dim=1) # (batch, T, input_dim)
|
||||
info = {
|
||||
'alpha': torch.stack(alphas, dim=1), # (batch, T, num_experts)
|
||||
'w_norm': torch.stack(w_norms, dim=1), # (batch, T)
|
||||
'final_w': state['w'],
|
||||
'final_m': state['m'],
|
||||
}
|
||||
return preds, info
|
||||
208
jspaceai/core_v2.py
Normal file
208
jspaceai/core_v2.py
Normal file
@@ -0,0 +1,208 @@
|
||||
"""
|
||||
改进版核心架构 v2
|
||||
|
||||
四大改进:
|
||||
1. 异构专家:不同模态用不同架构
|
||||
2. workspace 扩容(256)+ LayerNorm 防衰减
|
||||
3. RK4 积分(比 Euler 稳定)
|
||||
4. 更大容量
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import torch, torch.nn as nn, torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class JSpaceConfigV2:
|
||||
input_dim: int = 32
|
||||
workspace_dim: int = 256
|
||||
expert_dim: int = 64
|
||||
num_experts: int = 12
|
||||
num_wells: int = 8
|
||||
ode_steps: int = 4
|
||||
dt: float = 0.05
|
||||
tau_w: float = 1.0
|
||||
jacobian_sparsity: int = 32
|
||||
noise_std: float = 0.005
|
||||
use_rk4: bool = True
|
||||
use_layer_norm: bool = True
|
||||
|
||||
|
||||
class HeterogeneousExpert(nn.Module):
|
||||
"""异构专家基类——模态特定编码器 + ODE 动力学"""
|
||||
def __init__(self, expert_dim, workspace_dim, input_dim, num_wells, sparsity,
|
||||
use_rk4=True, use_layer_norm=True):
|
||||
super().__init__()
|
||||
self.expert_dim = expert_dim
|
||||
self.workspace_dim = workspace_dim
|
||||
self.use_rk4 = use_rk4
|
||||
|
||||
self.encoder = self._build_encoder(input_dim, expert_dim)
|
||||
self.well_a = nn.Parameter(torch.randn(num_wells, expert_dim) * 0.2)
|
||||
self.well_b = nn.Parameter(torch.zeros(num_wells))
|
||||
self.P_in = nn.Linear(expert_dim, expert_dim, bias=False)
|
||||
self.P_out = nn.Linear(expert_dim, workspace_dim, bias=False)
|
||||
self.J_raw = nn.Parameter(torch.randn(expert_dim, workspace_dim) * 0.05)
|
||||
self.sparsity = sparsity
|
||||
self.use_layer_norm = use_layer_norm
|
||||
if use_layer_norm:
|
||||
self.w_norm = nn.LayerNorm(expert_dim)
|
||||
self.m_norm = nn.LayerNorm(expert_dim)
|
||||
|
||||
def _build_encoder(self, input_dim, expert_dim):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_sparse_J(self):
|
||||
if self.sparsity >= self.workspace_dim:
|
||||
return self.J_raw
|
||||
abs_J = self.J_raw.abs()
|
||||
_, topk_idx = abs_J.topk(self.sparsity, dim=-1)
|
||||
mask = torch.zeros_like(self.J_raw)
|
||||
mask.scatter_(-1, topk_idx, 1.0)
|
||||
return self.J_raw * mask
|
||||
|
||||
def grad_potential(self, m):
|
||||
am = F.linear(m, self.well_a, self.well_b)
|
||||
sig = torch.sigmoid(am)
|
||||
grad_wells = torch.matmul(sig, self.well_a)
|
||||
return m - 0.5 * grad_wells
|
||||
|
||||
def deriv(self, m, w, x_feat):
|
||||
J = self.get_sparse_J()
|
||||
w_proj = F.linear(w, J)
|
||||
x_proj = self.P_in(x_feat)
|
||||
grad_U = self.grad_potential(m)
|
||||
return -grad_U + w_proj + x_proj
|
||||
|
||||
def forward(self, m, w, x, dt, noise_std):
|
||||
x_feat = self.encoder(x)
|
||||
if self.use_layer_norm:
|
||||
x_feat = self.w_norm(x_feat)
|
||||
|
||||
if self.use_rk4:
|
||||
k1 = self.deriv(m, w, x_feat)
|
||||
k2 = self.deriv(m + 0.5*dt*k1, w, x_feat)
|
||||
k3 = self.deriv(m + 0.5*dt*k2, w, x_feat)
|
||||
k4 = self.deriv(m + dt*k3, w, x_feat)
|
||||
m_next = m + (dt/6.0)*(k1 + 2*k2 + 2*k3 + k4)
|
||||
else:
|
||||
dm = self.deriv(m, w, x_feat)
|
||||
m_next = m + dt * dm
|
||||
|
||||
if noise_std > 0:
|
||||
m_next = m_next + torch.randn_like(m_next) * noise_std
|
||||
if self.use_layer_norm:
|
||||
m_next = self.m_norm(m_next)
|
||||
|
||||
contribution = self.P_out(m_next)
|
||||
return m_next, contribution
|
||||
|
||||
|
||||
class VisualExpert(HeterogeneousExpert):
|
||||
def _build_encoder(self, input_dim, expert_dim):
|
||||
return nn.Sequential(
|
||||
nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim))
|
||||
|
||||
class AudioExpert(HeterogeneousExpert):
|
||||
def _build_encoder(self, input_dim, expert_dim):
|
||||
return nn.Sequential(
|
||||
nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim))
|
||||
|
||||
class LanguageExpert(HeterogeneousExpert):
|
||||
def _build_encoder(self, input_dim, expert_dim):
|
||||
return nn.Sequential(
|
||||
nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, expert_dim))
|
||||
|
||||
class CrossModalExpert(HeterogeneousExpert):
|
||||
def _build_encoder(self, input_dim, expert_dim):
|
||||
return nn.Sequential(
|
||||
nn.Linear(input_dim, 64), nn.ReLU(),
|
||||
nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, expert_dim))
|
||||
|
||||
|
||||
class JSpaceWorkspaceV2(nn.Module):
|
||||
"""改进版 workspace——LayerNorm 防衰减"""
|
||||
def __init__(self, workspace_dim, input_dim, num_experts):
|
||||
super().__init__()
|
||||
self.workspace_dim = workspace_dim
|
||||
self.query_gen = nn.Sequential(
|
||||
nn.Linear(input_dim + workspace_dim, 128),
|
||||
nn.ReLU(), nn.Linear(128, workspace_dim))
|
||||
self.ln = nn.LayerNorm(workspace_dim)
|
||||
|
||||
def forward(self, w, x, contributions, dt, tau_w):
|
||||
q = self.query_gen(torch.cat([x, w], dim=-1))
|
||||
scores = (contributions * q.unsqueeze(1)).sum(dim=-1)
|
||||
alpha = F.softmax(scores, dim=-1)
|
||||
aggregated = (alpha.unsqueeze(-1) * contributions).sum(dim=1)
|
||||
dw = (-w + aggregated) / tau_w
|
||||
w_next = w + dt * dw
|
||||
w_next = self.ln(w_next)
|
||||
return w_next, alpha
|
||||
|
||||
|
||||
class JSpaceModelV2(nn.Module):
|
||||
"""改进版 JSpace 模型——异构专家 + 大 workspace + RK4"""
|
||||
def __init__(self, config: JSpaceConfigV2):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
expert_types = (
|
||||
[VisualExpert]*4 + [AudioExpert]*2 +
|
||||
[LanguageExpert]*2 + [CrossModalExpert]*4
|
||||
)[:config.num_experts]
|
||||
self.experts = nn.ModuleList([
|
||||
et(expert_dim=config.expert_dim,
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_wells=config.num_wells,
|
||||
sparsity=config.jacobian_sparsity,
|
||||
use_rk4=config.use_rk4,
|
||||
use_layer_norm=config.use_layer_norm)
|
||||
for et in expert_types
|
||||
])
|
||||
self.expert_modality = (
|
||||
['visual']*2 + ['screen']*2 + ['audio']*2 +
|
||||
['text']*2 + ['mouse']*2 + ['cross']*2
|
||||
)[:config.num_experts]
|
||||
self.workspace = JSpaceWorkspaceV2(
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_experts=config.num_experts,
|
||||
)
|
||||
|
||||
def init_state(self, batch_size, device):
|
||||
return {
|
||||
'w': torch.zeros(batch_size, self.config.workspace_dim, device=device),
|
||||
'm': [torch.zeros(batch_size, self.config.expert_dim, device=device)
|
||||
for _ in range(self.config.num_experts)],
|
||||
}
|
||||
|
||||
def step(self, state, x, record_trajectory=False):
|
||||
w, ms, cfg = state['w'], state['m'], self.config
|
||||
w_traj = []
|
||||
for _ in range(cfg.ode_steps):
|
||||
contributions, new_ms = [], []
|
||||
for i, expert in enumerate(self.experts):
|
||||
m_next, contrib = expert(ms[i], w, x, cfg.dt, cfg.noise_std)
|
||||
new_ms.append(m_next)
|
||||
contributions.append(contrib)
|
||||
contributions = torch.stack(contributions, dim=1)
|
||||
w, alpha = self.workspace(w, x, contributions, cfg.dt, cfg.tau_w)
|
||||
ms = new_ms
|
||||
if record_trajectory:
|
||||
w_traj.append(w.detach())
|
||||
return {'w': w, 'm': ms}, w_traj
|
||||
|
||||
def forward(self, xs, state=None, record_trajectory=False):
|
||||
B, T = xs.shape[0], xs.shape[1]
|
||||
device = xs.device
|
||||
if state is None:
|
||||
state = self.init_state(B, device)
|
||||
outputs, w_norms = [], []
|
||||
for t in range(T):
|
||||
state, _ = self.step(state, xs[:, t], record_trajectory)
|
||||
outputs.append(state['w'])
|
||||
w_norms.append(state['w'].norm(dim=-1))
|
||||
return torch.stack(outputs, dim=1), {
|
||||
'w_norm': torch.stack(w_norms, dim=1),
|
||||
}
|
||||
335
jspaceai/desktop.py
Normal file
335
jspaceai/desktop.py
Normal file
@@ -0,0 +1,335 @@
|
||||
"""
|
||||
屏幕 + 键盘 + 鼠标 I/O 层
|
||||
|
||||
使用 mss(屏幕捕获)+ pynput(键盘鼠标监听)。
|
||||
|
||||
提供 ScreenCapture / KeyboardMonitor / MouseMonitor / DesktopStream / FullSensoryStream
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import mss
|
||||
import numpy as np
|
||||
import torch
|
||||
import threading
|
||||
import queue
|
||||
import time
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, field
|
||||
from pynput import keyboard, mouse
|
||||
from PIL import Image as PILImage
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputEvent:
|
||||
"""统一的输入事件"""
|
||||
timestamp: float
|
||||
device: str
|
||||
event_type: str
|
||||
data: object
|
||||
modifiers: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
class ScreenCapture:
|
||||
"""屏幕捕获流"""
|
||||
|
||||
def __init__(self, target_size: tuple = (32, 32), fps: int = 5):
|
||||
self.target_size = target_size
|
||||
self.fps = fps
|
||||
self.frame_queue: queue.Queue = queue.Queue(maxsize=10)
|
||||
self.running = False
|
||||
self.thread: Optional[threading.Thread] = None
|
||||
self._sct = None
|
||||
|
||||
def start(self):
|
||||
self._sct = mss.MSS()
|
||||
self.running = True
|
||||
self.thread = threading.Thread(target=self._capture_loop, daemon=True)
|
||||
self.thread.start()
|
||||
|
||||
def _capture_loop(self):
|
||||
interval = 1.0 / self.fps
|
||||
while self.running:
|
||||
try:
|
||||
monitor = self._sct.monitors[1]
|
||||
raw = self._sct.grab(monitor)
|
||||
img = np.array(raw)[:, :, :3]
|
||||
pil_img = PILImage.fromarray(img)
|
||||
pil_img = pil_img.resize(self.target_size, PILImage.BILINEAR)
|
||||
img_resized = np.array(pil_img)
|
||||
self.frame_queue.put_nowait(
|
||||
InputEvent(time.time(), 'screen', 'frame', img_resized)
|
||||
)
|
||||
except queue.Full:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(interval)
|
||||
|
||||
def get_frame(self) -> Optional[InputEvent]:
|
||||
try:
|
||||
return self.frame_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
return None
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
if self.thread:
|
||||
self.thread.join(timeout=1.0)
|
||||
if self._sct:
|
||||
self._sct.close()
|
||||
|
||||
|
||||
class KeyboardMonitor:
|
||||
"""键盘监听"""
|
||||
|
||||
def __init__(self):
|
||||
self.event_queue: queue.Queue = queue.Queue(maxsize=200)
|
||||
self.running = False
|
||||
self.listener: Optional[keyboard.Listener] = None
|
||||
self.buffer: list[str] = []
|
||||
|
||||
def start(self):
|
||||
self.running = True
|
||||
self.listener = keyboard.Listener(on_press=self._on_press)
|
||||
self.listener.start()
|
||||
|
||||
def _on_press(self, key):
|
||||
if not self.running:
|
||||
return
|
||||
try:
|
||||
char = key.char
|
||||
self.buffer.append(char)
|
||||
self.event_queue.put_nowait(
|
||||
InputEvent(time.time(), 'keyboard', 'key', char, {'action': 'press'})
|
||||
)
|
||||
except AttributeError:
|
||||
name = str(key).replace('Key.', '')
|
||||
self.event_queue.put_nowait(
|
||||
InputEvent(time.time(), 'keyboard', 'key', name,
|
||||
{'action': 'press', 'special': True})
|
||||
)
|
||||
|
||||
def get_events(self) -> list[InputEvent]:
|
||||
events = []
|
||||
while True:
|
||||
try:
|
||||
events.append(self.event_queue.get_nowait())
|
||||
except queue.Empty:
|
||||
break
|
||||
return events
|
||||
|
||||
def get_buffer_text(self) -> str:
|
||||
return ''.join(self.buffer)
|
||||
|
||||
def clear_buffer(self):
|
||||
self.buffer.clear()
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
if self.listener:
|
||||
self.listener.stop()
|
||||
|
||||
|
||||
class MouseMonitor:
|
||||
"""鼠标监听"""
|
||||
|
||||
def __init__(self):
|
||||
self.event_queue: queue.Queue = queue.Queue(maxsize=200)
|
||||
self.running = False
|
||||
self.listener: Optional[mouse.Listener] = None
|
||||
self.last_pos: tuple[int, int] = (0, 0)
|
||||
|
||||
def start(self):
|
||||
self.running = True
|
||||
self.listener = mouse.Listener(
|
||||
on_move=self._on_move, on_click=self._on_click, on_scroll=self._on_scroll,
|
||||
)
|
||||
self.listener.start()
|
||||
|
||||
def _on_move(self, x, y):
|
||||
if not self.running:
|
||||
return
|
||||
self.last_pos = (x, y)
|
||||
try:
|
||||
self.event_queue.put_nowait(
|
||||
InputEvent(time.time(), 'mouse', 'move', (x, y))
|
||||
)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
def _on_click(self, x, y, button, pressed):
|
||||
if not self.running:
|
||||
return
|
||||
btn = 'left' if button == mouse.Button.left else 'right'
|
||||
action = 'click' if pressed else 'release'
|
||||
try:
|
||||
self.event_queue.put_nowait(
|
||||
InputEvent(time.time(), 'mouse', 'click', (x, y),
|
||||
{'button': btn, 'action': action})
|
||||
)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
def _on_scroll(self, x, y, dx, dy):
|
||||
if not self.running:
|
||||
return
|
||||
try:
|
||||
self.event_queue.put_nowait(
|
||||
InputEvent(time.time(), 'mouse', 'scroll', (x, y),
|
||||
{'dx': dx, 'dy': dy})
|
||||
)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
def get_events(self) -> list[InputEvent]:
|
||||
events = []
|
||||
while True:
|
||||
try:
|
||||
events.append(self.event_queue.get_nowait())
|
||||
except queue.Empty:
|
||||
break
|
||||
return events
|
||||
|
||||
def get_position(self) -> tuple[int, int]:
|
||||
return self.last_pos
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
if self.listener:
|
||||
self.listener.stop()
|
||||
|
||||
|
||||
class DesktopStream:
|
||||
"""统一管理屏幕 + 键盘 + 鼠标"""
|
||||
|
||||
def __init__(self, use_screen: bool = True, use_keyboard: bool = True,
|
||||
use_mouse: bool = True, screen_size: tuple = (32, 32),
|
||||
screen_fps: int = 5):
|
||||
self.use_screen = use_screen
|
||||
self.use_keyboard = use_keyboard
|
||||
self.use_mouse = use_mouse
|
||||
self.screen = ScreenCapture(target_size=screen_size, fps=screen_fps) if use_screen else None
|
||||
self.keyboard = KeyboardMonitor() if use_keyboard else None
|
||||
self.mouse = MouseMonitor() if use_mouse else None
|
||||
|
||||
def start(self):
|
||||
if self.screen:
|
||||
try:
|
||||
self.screen.start()
|
||||
print(" 屏幕捕获已启动")
|
||||
except Exception as e:
|
||||
print(f" 屏幕捕获失败: {e}")
|
||||
self.screen = None
|
||||
|
||||
if self.keyboard:
|
||||
try:
|
||||
self.keyboard.start()
|
||||
print(" 键盘监听已启动")
|
||||
except Exception as e:
|
||||
print(f" 键盘监听失败: {e}")
|
||||
self.keyboard = None
|
||||
|
||||
if self.mouse:
|
||||
try:
|
||||
self.mouse.start()
|
||||
print(" 鼠标监听已启动")
|
||||
except Exception as e:
|
||||
print(f" 鼠标监听失败: {e}")
|
||||
self.mouse = None
|
||||
|
||||
def get_latest(self) -> dict:
|
||||
result = {
|
||||
'screen': None, 'keyboard': [], 'mouse': [],
|
||||
'mouse_pos': (0, 0), 'keyboard_buffer': '',
|
||||
}
|
||||
if self.screen:
|
||||
latest = None
|
||||
while True:
|
||||
f = self.screen.get_frame()
|
||||
if f is None:
|
||||
break
|
||||
latest = f
|
||||
result['screen'] = latest
|
||||
|
||||
if self.keyboard:
|
||||
result['keyboard'] = self.keyboard.get_events()
|
||||
result['keyboard_buffer'] = self.keyboard.get_buffer_text()
|
||||
|
||||
if self.mouse:
|
||||
result['mouse'] = self.mouse.get_events()
|
||||
result['mouse_pos'] = self.mouse.get_position()
|
||||
|
||||
return result
|
||||
|
||||
def stop(self):
|
||||
if self.screen:
|
||||
self.screen.stop()
|
||||
if self.keyboard:
|
||||
self.keyboard.stop()
|
||||
if self.mouse:
|
||||
self.mouse.stop()
|
||||
|
||||
|
||||
class FullSensoryStream:
|
||||
"""
|
||||
完整感知流:摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标。
|
||||
|
||||
五个通道并行采集,这是模型感知世界的全部输入。
|
||||
|
||||
通道优先级(当多个同时有数据时):
|
||||
1. 键盘(用户主动输入,最高)
|
||||
2. 鼠标点击
|
||||
3. 音频(麦克风)
|
||||
4. 屏幕(用户在看什么)
|
||||
5. 摄像头(环境)
|
||||
"""
|
||||
|
||||
def __init__(self, use_camera: bool = True, use_mic: bool = True,
|
||||
use_desktop: bool = True, img_size: tuple = (32, 32),
|
||||
sample_rate: int = 16000, audio_frame_size: int = 1024):
|
||||
from .realtime import MultimodalStream
|
||||
|
||||
self.av_stream = MultimodalStream(
|
||||
use_camera=use_camera, use_mic=use_mic,
|
||||
img_size=img_size, sample_rate=sample_rate,
|
||||
audio_frame_size=audio_frame_size,
|
||||
) if use_camera or use_mic else None
|
||||
|
||||
self.desktop = DesktopStream(
|
||||
use_screen=use_desktop, use_keyboard=use_desktop, use_mouse=use_desktop,
|
||||
screen_size=img_size,
|
||||
) if use_desktop else None
|
||||
|
||||
def start(self):
|
||||
if self.av_stream:
|
||||
self.av_stream.start()
|
||||
if self.desktop:
|
||||
self.desktop.start()
|
||||
|
||||
def get_latest(self) -> dict:
|
||||
"""获取所有通道的最新数据"""
|
||||
result = {
|
||||
'camera': None, 'audio': None,
|
||||
'screen': None, 'keyboard': [], 'mouse': [],
|
||||
'mouse_pos': (0, 0), 'keyboard_buffer': '',
|
||||
}
|
||||
if self.av_stream:
|
||||
av_frames = self.av_stream.get_latest_frames()
|
||||
result['camera'] = av_frames.get('image')
|
||||
result['audio'] = av_frames.get('audio')
|
||||
|
||||
if self.desktop:
|
||||
desk = self.desktop.get_latest()
|
||||
result.update(desk)
|
||||
|
||||
return result
|
||||
|
||||
def play_audio(self, audio: np.ndarray, blocking: bool = False):
|
||||
if self.av_stream:
|
||||
self.av_stream.play_audio(audio, blocking)
|
||||
|
||||
def stop(self):
|
||||
if self.av_stream:
|
||||
self.av_stream.stop()
|
||||
if self.desktop:
|
||||
self.desktop.stop()
|
||||
59
jspaceai/distill_encoder.py
Normal file
59
jspaceai/distill_encoder.py
Normal file
@@ -0,0 +1,59 @@
|
||||
"""小模型感知编码器——把小模型表征投影到 input_dim"""
|
||||
from __future__ import annotations
|
||||
import torch, torch.nn as nn, torch.nn.functional as F
|
||||
|
||||
|
||||
class SmallModelEncoder(nn.Module):
|
||||
def __init__(self, input_dim=32, model_name="gpt2", use_real_model=True):
|
||||
super().__init__()
|
||||
self.input_dim = input_dim
|
||||
self.model_name = model_name
|
||||
self._real_model = None
|
||||
self._tokenizer = None
|
||||
self._hidden_size = 128
|
||||
|
||||
if use_real_model:
|
||||
try:
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
self._tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
self._model = AutoModel.from_pretrained(model_name)
|
||||
self._model.eval()
|
||||
self._hidden_size = self._model.config.hidden_size
|
||||
self._real_model = self._model
|
||||
if self._tokenizer.pad_token is None:
|
||||
self._tokenizer.pad_token = self._tokenizer.eos_token
|
||||
print(f" [编码器] 已加载 {model_name} (hidden={self._hidden_size})")
|
||||
except Exception as e:
|
||||
print(f" [编码器] 降级: {e}")
|
||||
|
||||
if self._real_model is not None:
|
||||
self.proj = nn.Linear(self._hidden_size, input_dim, bias=False)
|
||||
else:
|
||||
self.embed = nn.Embedding(1000, 64)
|
||||
self.proj = nn.Sequential(
|
||||
nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, input_dim))
|
||||
|
||||
def forward(self, input_data):
|
||||
if self._real_model is not None:
|
||||
if isinstance(input_data, str):
|
||||
input_data = [input_data]
|
||||
if isinstance(input_data, list):
|
||||
with torch.no_grad():
|
||||
inputs = self._tokenizer(input_data, return_tensors="pt",
|
||||
truncation=True, max_length=64, padding=True)
|
||||
outputs = self._model(**inputs)
|
||||
hidden = outputs.last_hidden_state.mean(dim=1)
|
||||
return self.proj(hidden)
|
||||
if isinstance(input_data, torch.Tensor) and input_data.dtype == torch.long:
|
||||
with torch.no_grad():
|
||||
outputs = self._model(input_ids=input_data)
|
||||
hidden = outputs.last_hidden_state.mean(dim=1)
|
||||
return self.proj(hidden)
|
||||
if isinstance(input_data, torch.Tensor) and input_data.dtype == torch.long:
|
||||
emb = self.embed(input_data)
|
||||
if emb.dim() == 3: emb = emb.mean(dim=1)
|
||||
elif emb.dim() == 2: emb = emb.mean(dim=0, keepdim=True)
|
||||
return self.proj(emb)
|
||||
if isinstance(input_data, torch.Tensor):
|
||||
return self.proj(input_data)
|
||||
return torch.zeros(1, self.input_dim)
|
||||
82
jspaceai/distill_trainer.py
Normal file
82
jspaceai/distill_trainer.py
Normal file
@@ -0,0 +1,82 @@
|
||||
"""蒸馏训练器——把小模型理解能力迁移到 JspaceAI"""
|
||||
from __future__ import annotations
|
||||
import torch, torch.nn.functional as F, numpy as np
|
||||
|
||||
|
||||
class DistillationTrainer:
|
||||
"""蒸馏训练器
|
||||
|
||||
把小模型的理解能力迁移到 JspaceAI 编码器。
|
||||
训练后编码器离线也有理解能力,小模型可以断开。
|
||||
"""
|
||||
|
||||
def __init__(self, model_v2, encoder, texts, device='cpu'):
|
||||
self.model = model_v2.to(device)
|
||||
self.encoder = encoder.to(device)
|
||||
self.texts = texts
|
||||
self.device = device
|
||||
params = list(model_v2.parameters()) + list(encoder.proj.parameters())
|
||||
self.optimizer = torch.optim.Adam(params, lr=1e-3)
|
||||
self.history = []
|
||||
|
||||
def train_step(self, text_batch):
|
||||
# 1. 小模型表征(target)
|
||||
target = None
|
||||
if self.encoder._real_model is not None:
|
||||
with torch.no_grad():
|
||||
inputs = self.encoder._tokenizer(
|
||||
text_batch, return_tensors="pt",
|
||||
truncation=True, max_length=64, padding=True
|
||||
).to(self.device)
|
||||
outputs = self.encoder._model(**inputs)
|
||||
target = outputs.last_hidden_state.mean(dim=1)
|
||||
|
||||
# 2. 编码器输出
|
||||
encoded = self.encoder(text_batch)
|
||||
|
||||
# 3. workspace 处理
|
||||
xs = encoded.unsqueeze(1) # (batch, 1, input_dim)
|
||||
state = self.model.init_state(xs.shape[0], self.device)
|
||||
w_out, info = self.model(xs, state)
|
||||
|
||||
# 4. 蒸馏 loss
|
||||
if target is not None:
|
||||
target_proj = self.encoder.proj(target.to(self.device)).detach()
|
||||
distill_loss = F.mse_loss(encoded, target_proj)
|
||||
else:
|
||||
# 自监督
|
||||
if encoded.shape[0] > 1:
|
||||
distill_loss = F.mse_loss(encoded[:-1], encoded[1:].detach())
|
||||
else:
|
||||
distill_loss = torch.tensor(0.0, device=self.device)
|
||||
|
||||
# 5. workspace 稳定性(目标 ||w||≈1)
|
||||
w_norm = info['w_norm']
|
||||
stability_loss = F.mse_loss(w_norm, torch.ones_like(w_norm))
|
||||
|
||||
# 6. 总 loss
|
||||
total_loss = distill_loss + 0.1 * stability_loss
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
total_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
|
||||
self.optimizer.step()
|
||||
|
||||
self.history.append({
|
||||
'loss': total_loss.item(),
|
||||
'distill': distill_loss.item(),
|
||||
'stability': stability_loss.item(),
|
||||
'w_norm': w_norm.mean().item(),
|
||||
})
|
||||
return self.history[-1]
|
||||
|
||||
def train(self, n_steps=100, batch_size=4, verbose=True):
|
||||
for step in range(n_steps):
|
||||
batch = np.random.choice(self.texts, size=min(batch_size, len(self.texts)),
|
||||
replace=True).tolist()
|
||||
info = self.train_step(batch)
|
||||
if verbose and (step+1) % 20 == 0:
|
||||
print(f" step {step+1:4d} | loss {info['loss']:.4f} | "
|
||||
f"distill {info['distill']:.4f} | "
|
||||
f"||w|| {info['w_norm']:.3f}")
|
||||
return self.history
|
||||
730
jspaceai/embodied.py
Normal file
730
jspaceai/embodied.py
Normal file
@@ -0,0 +1,730 @@
|
||||
"""
|
||||
输出执行器层 + 神经系统
|
||||
|
||||
对应人类神经系统的各部分:
|
||||
- 大脑皮层: workspace w + 专家池(已在 multimodal.py)
|
||||
- 小脑: 运动控制器(前向模型+逆模型,精细动作)
|
||||
- 中枢神经: 动作调度器(反射弧+决策门控)
|
||||
- 海马体: 外部情景记忆库
|
||||
- 基底神经节: 动作价值学习(习惯化)
|
||||
- 执行器: 鼠标控制 + 键盘输出 + 音频输出 + 屏幕绘制
|
||||
|
||||
核心思想:输出和输入对称。
|
||||
输入:摄像头/麦克风/屏幕/键盘/鼠标 → 编码 → workspace
|
||||
输出:workspace → 解码 → 鼠标移动/键盘按键/音频播放/屏幕绘制
|
||||
|
||||
workspace 是模态无关的"意图空间"。
|
||||
"想点击左上角"这个意图,在 workspace 里是一个向量,
|
||||
解码到鼠标控制器就是移动+点击,解码到键盘就是 Tab+Enter。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import time
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, field
|
||||
from collections import deque
|
||||
from pynput import mouse as pynput_mouse
|
||||
from pynput import keyboard as pynput_keyboard
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 执行器层(对应手脚口)
|
||||
# ============================================================
|
||||
|
||||
class MouseActuator:
|
||||
"""
|
||||
鼠标执行器——对应"手"。
|
||||
|
||||
workspace 解码出动作向量 → 移动鼠标 + 点击。
|
||||
动作空间:
|
||||
(dx, dy, click_left, click_right, scroll)
|
||||
dx, dy: 相对移动量(-1 到 1,乘以灵敏度)
|
||||
click_left/right: 0 或 1
|
||||
scroll: 滚动量
|
||||
"""
|
||||
|
||||
def __init__(self, sensitivity: int = 200, enabled: bool = True):
|
||||
self.sensitivity = sensitivity
|
||||
self.enabled = enabled
|
||||
self.controller = pynput_mouse.Controller() if enabled else None
|
||||
|
||||
def execute(self, action: np.ndarray):
|
||||
"""执行鼠标动作
|
||||
|
||||
Args:
|
||||
action: (5,) = (dx, dy, click_l, click_r, scroll)
|
||||
"""
|
||||
if not self.enabled or self.controller is None:
|
||||
return
|
||||
|
||||
dx, dy, click_l, click_r, scroll = action
|
||||
|
||||
# 移动
|
||||
if abs(dx) > 0.01 or abs(dy) > 0.01:
|
||||
self.controller.move(int(dx * self.sensitivity), int(dy * self.sensitivity))
|
||||
|
||||
# 点击
|
||||
if click_l > 0.5:
|
||||
self.controller.click(pynput_mouse.Button.left)
|
||||
time.sleep(0.05)
|
||||
if click_r > 0.5:
|
||||
self.controller.click(pynput_mouse.Button.right)
|
||||
time.sleep(0.05)
|
||||
|
||||
# 滚动
|
||||
if abs(scroll) > 0.1:
|
||||
self.controller.scroll(0, int(scroll * 5))
|
||||
|
||||
def get_position(self) -> tuple[int, int]:
|
||||
if self.controller:
|
||||
return self.controller.position
|
||||
return (0, 0)
|
||||
|
||||
|
||||
class KeyboardActuator:
|
||||
"""
|
||||
键盘执行器——对应"手"+"口"(打字)。
|
||||
|
||||
workspace 解码出 token → 按键输入。
|
||||
"""
|
||||
|
||||
def __init__(self, enabled: bool = True):
|
||||
self.enabled = enabled
|
||||
self.controller = pynput_keyboard.Controller() if enabled else None
|
||||
|
||||
def type_text(self, text: str):
|
||||
"""输入文本"""
|
||||
if not self.enabled or self.controller is None:
|
||||
return
|
||||
self.controller.type(text)
|
||||
|
||||
def press_key(self, key: str):
|
||||
"""按下单个键"""
|
||||
if not self.enabled or self.controller is None:
|
||||
return
|
||||
try:
|
||||
self.controller.press(key)
|
||||
self.controller.release(key)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class AudioActuator:
|
||||
"""
|
||||
音频执行器——对应"口"(说话)。
|
||||
|
||||
workspace 解码出音频波形 → 扬声器播放。
|
||||
"""
|
||||
|
||||
def __init__(self, sample_rate: int = 16000, enabled: bool = True):
|
||||
self.sample_rate = sample_rate
|
||||
self.enabled = enabled
|
||||
import sounddevice as sd
|
||||
self._sd = sd
|
||||
|
||||
def play(self, audio: np.ndarray, blocking: bool = False):
|
||||
"""播放音频"""
|
||||
if not self.enabled:
|
||||
return
|
||||
audio = np.clip(audio, -1, 1).astype(np.float32)
|
||||
self._sd.play(audio, self.sample_rate)
|
||||
if blocking:
|
||||
self._sd.wait()
|
||||
|
||||
def stop(self):
|
||||
if self.enabled:
|
||||
self._sd.stop()
|
||||
|
||||
|
||||
class ScreenActuator:
|
||||
"""
|
||||
屏幕执行器——对应"手"(绘制)。
|
||||
|
||||
在屏幕上绘制 workspace 解码出的图像。
|
||||
用 OpenCV 显示一个窗口。
|
||||
"""
|
||||
|
||||
def __init__(self, enabled: bool = True, window_name: str = "JspaceAI Output"):
|
||||
self.enabled = enabled
|
||||
self.window_name = window_name
|
||||
import cv2
|
||||
self._cv2 = cv2
|
||||
|
||||
def show_image(self, img: np.ndarray):
|
||||
"""显示图像"""
|
||||
if not self.enabled:
|
||||
return
|
||||
# img: (H, W, 3) RGB 或 (H, W) 灰度
|
||||
if img.ndim == 3 and img.shape[2] == 3:
|
||||
img_bgr = self._cv2.cvtColor(img.astype(np.uint8), self._cv2.COLOR_RGB2BGR)
|
||||
else:
|
||||
img_bgr = img.astype(np.uint8)
|
||||
self._cv2.imshow(self.window_name, img_bgr)
|
||||
self._cv2.waitKey(1)
|
||||
|
||||
def close(self):
|
||||
if self.enabled:
|
||||
self._cv2.destroyWindow(self.window_name)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 小脑:运动控制器(前向模型 + 逆模型)
|
||||
# ============================================================
|
||||
|
||||
class Cerebellum(nn.Module):
|
||||
"""
|
||||
小脑——运动协调与精细控制。
|
||||
|
||||
前向模型:预测"如果执行动作 A,鼠标会到哪里"
|
||||
逆模型:给定"目标位置",计算"需要什么动作"
|
||||
|
||||
人类小脑学习动作的精细映射,让动作平滑准确。
|
||||
我们这里学习 workspace 意图 → 精确动作参数的映射。
|
||||
|
||||
逆模型:workspace → 动作参数
|
||||
前向模型:动作参数 → 预测结果(用于误差反馈学习)
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim: int, action_dim: int = 5):
|
||||
super().__init__()
|
||||
# 逆模型:workspace → action
|
||||
self.inverse_model = nn.Sequential(
|
||||
nn.Linear(workspace_dim, 64),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, 32),
|
||||
nn.ReLU(),
|
||||
nn.Linear(32, action_dim),
|
||||
nn.Tanh(), # 动作在 [-1, 1]
|
||||
)
|
||||
# 前向模型:action + 当前状态 → 预测下一状态
|
||||
self.forward_model = nn.Sequential(
|
||||
nn.Linear(action_dim + workspace_dim, 64),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, workspace_dim),
|
||||
)
|
||||
self.action_dim = action_dim
|
||||
|
||||
def compute_action(self, w: torch.Tensor) -> torch.Tensor:
|
||||
"""逆模型:从 workspace 意图计算动作"""
|
||||
return self.inverse_model(w)
|
||||
|
||||
def predict_next(self, w: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
|
||||
"""前向模型:预测执行动作后的 workspace 状态"""
|
||||
return self.forward_model(torch.cat([action, w], dim=-1))
|
||||
|
||||
def compute_motor_error(self, w: torch.Tensor, action: torch.Tensor,
|
||||
w_actual_next: torch.Tensor) -> torch.Tensor:
|
||||
"""计算运动误差——用于小脑学习"""
|
||||
w_pred = self.predict_next(w, action)
|
||||
return F.mse_loss(w_pred, w_actual_next)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 中枢神经:动作调度器
|
||||
# ============================================================
|
||||
|
||||
@dataclass
|
||||
class ReflexArc:
|
||||
"""反射弧——不经过大脑的快速反应"""
|
||||
trigger: str # 触发条件描述
|
||||
condition: callable # 检查函数
|
||||
action: callable # 执行函数
|
||||
priority: int = 0 # 优先级
|
||||
|
||||
|
||||
class CentralNervousSystem:
|
||||
"""
|
||||
中枢神经系统——动作调度。
|
||||
|
||||
功能:
|
||||
1. 反射弧:快速反应,不经过 workspace
|
||||
- 如:突然大声音 → 退缩
|
||||
- 如:屏幕突然变暗 → 警觉
|
||||
2. 决策门控:决定是否让 workspace 的意图执行
|
||||
- 高风险动作需要"确认"
|
||||
- 习惯化动作直接执行
|
||||
3. 动作序列:把复杂意图拆成动作序列
|
||||
- 如"点击按钮"→ 移动到位置 + 点击
|
||||
|
||||
这是"自由意志"的工程对应——不是所有意图都执行,
|
||||
系统有一个门控机制决定哪些意图变成行动。
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.reflexes: list[ReflexArc] = []
|
||||
self.action_history: deque = deque(maxlen=100)
|
||||
self.inhibit_score: float = 0.0 # 抑制分数,高时阻止动作
|
||||
|
||||
def add_reflex(self, reflex: ReflexArc):
|
||||
self.reflexes.append(reflex)
|
||||
self.reflexes.sort(key=lambda r: -r.priority)
|
||||
|
||||
def check_reflexes(self, sensory_state: dict) -> Optional[callable]:
|
||||
"""检查是否有反射触发"""
|
||||
for reflex in self.reflexes:
|
||||
try:
|
||||
if reflex.condition(sensory_state):
|
||||
return reflex.action
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
|
||||
def should_execute(self, action_strength: float, risk: float = 0.0) -> bool:
|
||||
"""决策门控:是否执行动作
|
||||
|
||||
Args:
|
||||
action_strength: 动作强度(workspace 驱动)
|
||||
risk: 风险评估(0-1)
|
||||
|
||||
Returns:
|
||||
是否执行
|
||||
"""
|
||||
# 抑制分数高时不执行
|
||||
threshold = 0.3 + risk * 0.5 + self.inhibit_score
|
||||
return action_strength > threshold
|
||||
|
||||
def record_action(self, action: np.ndarray, modality: str):
|
||||
"""记录执行的动作"""
|
||||
self.action_history.append({
|
||||
'time': time.time(),
|
||||
'action': action.tolist() if hasattr(action, 'tolist') else action,
|
||||
'modality': modality,
|
||||
})
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 海马体:外部情景记忆库
|
||||
# ============================================================
|
||||
|
||||
class Hippocampus:
|
||||
"""
|
||||
海马体——情景记忆。
|
||||
|
||||
存储历史 workspace 快照 + 时间戳 + 上下文。
|
||||
当前 workspace 可以"回忆"相似的历史状态。
|
||||
|
||||
人类的情景记忆:"我记得昨天在那个房间里说了什么"
|
||||
对应:检索与当前 workspace 相似的历史 workspace。
|
||||
|
||||
实现用简单的向量数据库(numpy + cosine similarity)。
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int = 1000, workspace_dim: int = 64):
|
||||
self.capacity = capacity
|
||||
self.workspace_dim = workspace_dim
|
||||
self.memories: deque = deque(maxlen=capacity)
|
||||
|
||||
def store(self, w: np.ndarray, context: dict = None):
|
||||
"""存储一个 workspace 快照"""
|
||||
self.memories.append({
|
||||
'w': w.copy(),
|
||||
'context': context or {},
|
||||
'timestamp': time.time(),
|
||||
})
|
||||
|
||||
def recall(self, w_query: np.ndarray, top_k: int = 3) -> list[dict]:
|
||||
"""检索相似的历史记忆
|
||||
|
||||
Args:
|
||||
w_query: 当前 workspace
|
||||
top_k: 返回最相似的 k 个
|
||||
|
||||
Returns:
|
||||
list of {w, context, timestamp, similarity}
|
||||
"""
|
||||
if not self.memories:
|
||||
return []
|
||||
|
||||
# 计算相似度
|
||||
similarities = []
|
||||
for mem in self.memories:
|
||||
sim = np.dot(w_query, mem['w']) / (
|
||||
np.linalg.norm(w_query) * np.linalg.norm(mem['w']) + 1e-8
|
||||
)
|
||||
similarities.append(sim)
|
||||
|
||||
# 取 top-k
|
||||
top_idx = np.argsort(similarities)[-top_k:][::-1]
|
||||
results = []
|
||||
for idx in top_idx:
|
||||
mem = self.memories[idx]
|
||||
results.append({
|
||||
'w': mem['w'],
|
||||
'context': mem['context'],
|
||||
'timestamp': mem['timestamp'],
|
||||
'similarity': similarities[idx],
|
||||
})
|
||||
return results
|
||||
|
||||
def size(self) -> int:
|
||||
return len(self.memories)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 基底神经节:动作价值学习
|
||||
# ============================================================
|
||||
|
||||
class BasalGanglia:
|
||||
"""
|
||||
基底神经节——习惯学习与动作选择。
|
||||
|
||||
学习"在什么状态下执行什么动作价值多少"。
|
||||
高频执行的(workspace, action)对会"习惯化"——直接执行不经过思考。
|
||||
|
||||
对应人类的习惯:开车的动作熟练后不需要思考,
|
||||
就是基底神经节接管了动作选择。
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim: int = 64, n_actions: int = 5,
|
||||
learning_rate: float = 0.01):
|
||||
self.workspace_dim = workspace_dim
|
||||
self.n_actions = n_actions
|
||||
self.lr = learning_rate
|
||||
# Q-table 的近似:用线性函数 Q(s, a) = w_a · s
|
||||
self.action_weights = np.zeros((n_actions, workspace_dim))
|
||||
# 习惯化计数
|
||||
self.habit_counts = np.zeros(n_actions)
|
||||
|
||||
def compute_values(self, w: np.ndarray) -> np.ndarray:
|
||||
"""计算各动作的价值 Q(s, a)"""
|
||||
return self.action_weights @ w # (n_actions,)
|
||||
|
||||
def select_action(self, w: np.ndarray, exploration: float = 0.1) -> int:
|
||||
"""选择动作(ε-greedy)"""
|
||||
values = self.compute_values(w)
|
||||
if np.random.random() < exploration:
|
||||
return np.random.randint(self.n_actions)
|
||||
return np.argmax(values)
|
||||
|
||||
def update(self, w: np.ndarray, action: int, reward: float):
|
||||
"""更新动作价值(TD learning 简化版)"""
|
||||
values = self.compute_values(w)
|
||||
td_error = reward - values[action]
|
||||
self.action_weights[action] += self.lr * td_error * w
|
||||
self.habit_counts[action] += 1
|
||||
|
||||
def is_habitual(self, action: int, threshold: int = 10) -> bool:
|
||||
"""判断动作是否已习惯化"""
|
||||
return self.habit_counts[action] >= threshold
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 完整的具身 Agent
|
||||
# ============================================================
|
||||
|
||||
class EmbodiedAgent:
|
||||
"""
|
||||
完整的具身 Agent——感知-思考-行动闭环。
|
||||
|
||||
结构(对应人类神经系统):
|
||||
感知层(眼耳皮肤)→ FullSensoryStream
|
||||
↓ 编码
|
||||
大脑皮层(思考)→ MultimodalJSpaceModel
|
||||
↓ workspace w
|
||||
海马体(记忆)→ Hippocampus 存储和检索
|
||||
↓
|
||||
基底神经节(动作选择)→ BasalGanglia 选动作
|
||||
↓
|
||||
小脑(运动控制)→ Cerebellum 精细化动作
|
||||
↓
|
||||
中枢神经(门控)→ CentralNervousSystem 决定是否执行
|
||||
↓
|
||||
执行器(手足口)→ MouseActuator + KeyboardActuator + AudioActuator
|
||||
|
||||
循环:
|
||||
1. 感知:从五感获取输入
|
||||
2. 思考:workspace 更新
|
||||
3. 回忆:海马体检索相关记忆
|
||||
4. 决策:基底神经节选动作
|
||||
5. 精化:小脑计算动作参数
|
||||
6. 门控:中枢神经决定执行
|
||||
7. 行动:执行器执行
|
||||
8. 学习:更新基底神经节、小脑、海马体
|
||||
"""
|
||||
|
||||
def __init__(self, model, device: str = 'cpu',
|
||||
enable_mouse_output: bool = False,
|
||||
enable_keyboard_output: bool = False,
|
||||
enable_audio_output: bool = True,
|
||||
enable_screen_output: bool = True,
|
||||
enable_memory: bool = True,
|
||||
risk_threshold: float = 0.3):
|
||||
self.model = model
|
||||
self.device = device
|
||||
self.config = model.config
|
||||
|
||||
# 感知层
|
||||
from .desktop import FullSensoryStream
|
||||
from .platform import get_screen_size
|
||||
self.screen_w, self.screen_h = get_screen_size()
|
||||
self.senses = FullSensoryStream(
|
||||
use_camera=True, use_mic=True, use_desktop=True,
|
||||
img_size=(self.config.img_size, self.config.img_size),
|
||||
)
|
||||
|
||||
# 执行器层
|
||||
self.mouse_actuator = MouseActuator(enabled=enable_mouse_output)
|
||||
self.keyboard_actuator = KeyboardActuator(enabled=enable_keyboard_output)
|
||||
self.audio_actuator = AudioActuator(enabled=enable_audio_output)
|
||||
self.screen_actuator = ScreenActuator(enabled=enable_screen_output)
|
||||
|
||||
# 神经系统
|
||||
self.cerebellum = Cerebellum(
|
||||
workspace_dim=self.config.workspace_dim,
|
||||
action_dim=5, # (dx, dy, click_l, click_r, scroll)
|
||||
).to(device)
|
||||
self.cns = CentralNervousSystem()
|
||||
self.hippocampus = Hippocampus(
|
||||
workspace_dim=self.config.workspace_dim,
|
||||
) if enable_memory else None
|
||||
self.basal_ganglia = BasalGanglia(
|
||||
workspace_dim=self.config.workspace_dim,
|
||||
n_actions=5,
|
||||
)
|
||||
|
||||
# 内部状态
|
||||
self.state = model.init_state(1, torch.device(device))
|
||||
self.step_count = 0
|
||||
self.running = False
|
||||
self.risk_threshold = risk_threshold
|
||||
|
||||
# 设置反射弧
|
||||
self._setup_reflexes()
|
||||
|
||||
def _setup_reflexes(self):
|
||||
"""设置基本反射"""
|
||||
# 反射1:大声音 → 抑制动作
|
||||
def loud_noise_condition(state):
|
||||
audio = state.get('audio')
|
||||
if audio and hasattr(audio, 'data'):
|
||||
return np.abs(audio.data).mean() > 0.5
|
||||
return False
|
||||
|
||||
def loud_noise_response():
|
||||
self.cns.inhibit_score = 0.5
|
||||
|
||||
self.cns.add_reflex(ReflexArc(
|
||||
trigger='loud_noise',
|
||||
condition=loud_noise_condition,
|
||||
action=loud_noise_response,
|
||||
priority=10,
|
||||
))
|
||||
|
||||
def perceive(self) -> dict:
|
||||
"""感知:从五感获取输入"""
|
||||
return self.senses.get_latest()
|
||||
|
||||
def think(self, sensory_data: dict) -> tuple[torch.Tensor, str]:
|
||||
"""思考:更新 workspace
|
||||
|
||||
Returns:
|
||||
w: workspace 状态
|
||||
modality: 本次输入的模态
|
||||
"""
|
||||
# 按优先级选模态
|
||||
modality = None
|
||||
input_tensor = None
|
||||
|
||||
if sensory_data['keyboard']:
|
||||
key_ids = [ord(ev.data[0]) if ev.data and len(ev.data) == 1 else 0
|
||||
for ev in sensory_data['keyboard'][:8]]
|
||||
if key_ids:
|
||||
modality = 'keyboard'
|
||||
input_tensor = torch.tensor([key_ids], dtype=torch.long).to(self.device)
|
||||
|
||||
elif sensory_data['mouse']:
|
||||
ev = sensory_data['mouse'][0]
|
||||
x, y = ev.data
|
||||
modality = 'mouse'
|
||||
input_tensor = torch.tensor([[
|
||||
x / self.screen_w, y / self.screen_h,
|
||||
1.0 if ev.modifiers.get('button') == 'left' else 0.0,
|
||||
1.0 if ev.modifiers.get('button') == 'right' else 0.0,
|
||||
]], dtype=torch.float32).to(self.device)
|
||||
|
||||
elif sensory_data['audio']:
|
||||
modality = 'audio'
|
||||
input_tensor = torch.tensor([sensory_data['audio'].data],
|
||||
dtype=torch.float32).to(self.device)
|
||||
|
||||
elif sensory_data['screen']:
|
||||
img = sensory_data['screen'].data
|
||||
modality = 'screen'
|
||||
input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0
|
||||
|
||||
elif sensory_data['camera']:
|
||||
img = sensory_data['camera'].data
|
||||
modality = 'image'
|
||||
input_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0
|
||||
|
||||
if modality is None or input_tensor is None:
|
||||
return self.state['w'], 'idle'
|
||||
|
||||
# 编码 + forward
|
||||
with torch.no_grad():
|
||||
x = self.model.encode_modality(modality, input_tensor)
|
||||
if x.dim() == 1: x = x.unsqueeze(0)
|
||||
if x.dim() == 3: x = x[:, -1, :]
|
||||
if x.shape[0] != 1: x = x[-1:]
|
||||
self.state, _ = self.model.step(self.state, x)
|
||||
|
||||
return self.state['w'], modality
|
||||
|
||||
def remember(self, w: torch.Tensor, context: dict):
|
||||
"""存储到海马体"""
|
||||
if self.hippocampus:
|
||||
self.hippocampus.store(w[0].cpu().numpy(), context)
|
||||
|
||||
def recall_memories(self, w: torch.Tensor, top_k: int = 3) -> list:
|
||||
"""从海马体回忆"""
|
||||
if self.hippocampus:
|
||||
return self.hippocampus.recall(w[0].cpu().numpy(), top_k)
|
||||
return []
|
||||
|
||||
def decide_and_act(self, w: torch.Tensor, modality: str) -> dict:
|
||||
"""决策和行动
|
||||
|
||||
1. 基底神经节选动作
|
||||
2. 小脑计算动作参数
|
||||
3. 中枢神经门控
|
||||
4. 执行器执行
|
||||
"""
|
||||
w_np = w[0].cpu().numpy()
|
||||
|
||||
# 1. 基底神经节:选动作类别
|
||||
action_idx = self.basal_ganglia.select_action(w_np, exploration=0.2)
|
||||
|
||||
# 2. 小脑:计算精确动作参数
|
||||
with torch.no_grad():
|
||||
action_params = self.cerebellum.compute_action(w)[0].cpu().numpy()
|
||||
|
||||
# 3. 中枢神经:门控
|
||||
action_strength = np.abs(action_params).max()
|
||||
risk = 0.0
|
||||
# 鼠标点击风险较高
|
||||
if action_params[2] > 0.5 or action_params[3] > 0.5:
|
||||
risk = 0.5
|
||||
|
||||
execute = self.cns.should_execute(action_strength, risk)
|
||||
|
||||
# 4. 执行
|
||||
action_taken = None
|
||||
if execute:
|
||||
# 鼠标动作
|
||||
self.mouse_actuator.execute(action_params)
|
||||
# 音频输出(从 workspace 解码)
|
||||
if self.step_count % 10 == 0:
|
||||
with torch.no_grad():
|
||||
audio_out = self.model.audio_decoder(w)[0].cpu().numpy()
|
||||
self.audio_actuator.play(audio_out)
|
||||
# 屏幕输出
|
||||
if self.screen_actuator.enabled:
|
||||
with torch.no_grad():
|
||||
img_out = self.model.visual_decoder(w)[0].cpu().permute(1, 2, 0).numpy()
|
||||
img_out = ((img_out + 1) / 2 * 255).clip(0, 255).astype(np.uint8)
|
||||
self.screen_actuator.show_image(img_out)
|
||||
|
||||
action_taken = action_params.tolist()
|
||||
self.cns.record_action(action_params, modality)
|
||||
|
||||
return {
|
||||
'action_idx': action_idx,
|
||||
'action_params': action_params.tolist(),
|
||||
'executed': execute,
|
||||
'action_strength': float(action_strength),
|
||||
'risk': float(risk),
|
||||
}
|
||||
|
||||
def learn(self, w: torch.Tensor, action_params: np.ndarray, reward: float = 0.0):
|
||||
"""学习——更新基底神经节和小脑"""
|
||||
w_np = w[0].cpu().numpy()
|
||||
|
||||
# 基底神经节:更新动作价值
|
||||
action_idx = self.basal_ganglia.select_action(w_np, exploration=0.0)
|
||||
self.basal_ganglia.update(w_np, action_idx, reward)
|
||||
|
||||
def step_once(self) -> dict:
|
||||
"""执行一步完整的感知-思考-行动循环"""
|
||||
# 1. 感知
|
||||
sensory_data = self.perceive()
|
||||
|
||||
# 2. 检查反射
|
||||
reflex_action = self.cns.check_reflexes(sensory_data)
|
||||
if reflex_action:
|
||||
reflex_action()
|
||||
|
||||
# 3. 思考
|
||||
w, modality = self.think(sensory_data)
|
||||
|
||||
# 4. 回忆
|
||||
memories = self.recall_memories(w)
|
||||
|
||||
# 5. 决策和行动
|
||||
action_info = self.decide_and_act(w, modality)
|
||||
|
||||
# 6. 学习(自监督:预测误差作为 reward)
|
||||
reward = -action_info['risk'] # 简化:风险越低 reward 越高
|
||||
self.learn(w, np.array(action_info['action_params']), reward)
|
||||
|
||||
# 7. 记忆存储
|
||||
self.remember(w, {
|
||||
'modality': modality,
|
||||
'action': action_info,
|
||||
'step': self.step_count,
|
||||
})
|
||||
|
||||
self.step_count += 1
|
||||
|
||||
return {
|
||||
'step': self.step_count,
|
||||
'modality': modality,
|
||||
'w_norm': w.norm().item(),
|
||||
'action': action_info,
|
||||
'memories_count': self.hippocampus.size() if self.hippocampus else 0,
|
||||
}
|
||||
|
||||
def run(self, n_steps: int = 100, interval: float = 0.2,
|
||||
on_step: callable = None):
|
||||
"""运行感知-思考-行动循环"""
|
||||
self.running = True
|
||||
self.senses.start()
|
||||
print(f"具身 Agent 启动,{n_steps} 步")
|
||||
print(f"执行器: mouse={self.mouse_actuator.enabled}, "
|
||||
f"keyboard={self.keyboard_actuator.enabled}, "
|
||||
f"audio={self.audio_actuator.enabled}, "
|
||||
f"screen={self.screen_actuator.enabled}")
|
||||
print("=" * 60)
|
||||
|
||||
step_log = []
|
||||
try:
|
||||
for _ in range(n_steps):
|
||||
if not self.running:
|
||||
break
|
||||
info = self.step_once()
|
||||
step_log.append(info)
|
||||
|
||||
if on_step:
|
||||
on_step(info)
|
||||
elif info['step'] % 5 == 0:
|
||||
print(f" step {info['step']:3d} | mod {info['modality']:8s} | "
|
||||
f"||w|| {info['w_norm']:.3f} | "
|
||||
f"action {info['action']['action_idx']} | "
|
||||
f"executed {info['action']['executed']}")
|
||||
|
||||
time.sleep(interval)
|
||||
except KeyboardInterrupt:
|
||||
print("\n用户中断")
|
||||
finally:
|
||||
self.running = False
|
||||
self.senses.stop()
|
||||
self.audio_actuator.stop()
|
||||
if self.screen_actuator.enabled:
|
||||
self.screen_actuator.close()
|
||||
|
||||
return step_log
|
||||
216
jspaceai/evolution.py
Normal file
216
jspaceai/evolution.py
Normal file
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
自主进化训练器
|
||||
|
||||
核心:模型在推理的同时持续学习。每处理一段文本,参数就更新一次。
|
||||
|
||||
进化循环:
|
||||
1. 喂入新文本片段
|
||||
2. forward + 计算 next-token loss
|
||||
3. EWC 优化器更新参数(保护旧知识)
|
||||
4. 经验回放:当前片段存入 buffer,定期回放旧片段
|
||||
5. 周期性 consolidate EWC(更新 Fisher 信息和锚点)
|
||||
6. 追踪专家可塑性统计
|
||||
7. 定期生成样本,观察进化效果
|
||||
|
||||
这个循环可以无限运行——模型永远不会"训练完成",它一直在进化。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from typing import Callable
|
||||
from .language_model import (
|
||||
JSpaceLanguageModel, LanguageConfig,
|
||||
ExperienceReplay, EWCOptimizer, ExpertPlasticity,
|
||||
)
|
||||
|
||||
|
||||
class EvolutionTrainer:
|
||||
"""自主进化训练器"""
|
||||
|
||||
def __init__(self, model: JSpaceLanguageModel, config: LanguageConfig,
|
||||
lr: float = 1e-3, ewc_lambda: float = 0.1,
|
||||
device: str = 'cpu'):
|
||||
self.model = model.to(device)
|
||||
self.config = config
|
||||
self.device = device
|
||||
|
||||
# 三大自主进化机制
|
||||
self.ewc_optimizer = EWCOptimizer(model, lr=lr, ewc_lambda=ewc_lambda)
|
||||
self.replay_buffer = ExperienceReplay(capacity=500, seq_len=64)
|
||||
self.plasticity = ExpertPlasticity(num_experts=config.num_experts)
|
||||
|
||||
# 进化历史追踪
|
||||
self.history: list[dict] = []
|
||||
|
||||
def learn_step(self, token_seq: torch.Tensor) -> dict:
|
||||
"""单步学习
|
||||
|
||||
Args:
|
||||
token_seq: (batch, T) token indices
|
||||
|
||||
Returns:
|
||||
stats: 包含 loss、注意力、专家统计等
|
||||
"""
|
||||
token_seq = token_seq.to(self.device)
|
||||
|
||||
# 1. Forward
|
||||
logits, info = self.model(token_seq)
|
||||
|
||||
# 2. Next-token prediction loss
|
||||
# logits[:, t] 预测 token_seq[:, t+1]
|
||||
pred_logits = logits[:, :-1] # (batch, T-1, vocab)
|
||||
targets = token_seq[:, 1:] # (batch, T-1)
|
||||
loss = F.cross_entropy(
|
||||
pred_logits.reshape(-1, self.config.vocab_size),
|
||||
targets.reshape(-1),
|
||||
)
|
||||
|
||||
# 3. 经验回放:如果有足够样本,混入旧数据
|
||||
replay_loss = torch.tensor(0.0, device=self.device)
|
||||
if len(self.replay_buffer.buffer) >= 8:
|
||||
replay_seq = self.replay_buffer.sample(4)
|
||||
if replay_seq is not None:
|
||||
replay_seq = replay_seq.to(self.device)
|
||||
replay_logits, _ = self.model(replay_seq)
|
||||
replay_pred = replay_logits[:, :-1]
|
||||
replay_targets = replay_seq[:, 1:]
|
||||
replay_loss = F.cross_entropy(
|
||||
replay_pred.reshape(-1, self.config.vocab_size),
|
||||
replay_targets.reshape(-1),
|
||||
)
|
||||
|
||||
# 4. 总 loss + EWC 正则 + 经验回放
|
||||
total_task_loss = loss + 0.5 * replay_loss
|
||||
total_loss = self.ewc_optimizer.step(total_task_loss)
|
||||
|
||||
# 5. 更新专家可塑性统计
|
||||
self.plasticity.update(info['alpha'].detach(), token_seq.detach())
|
||||
|
||||
# 6. 存入经验回放
|
||||
self.replay_buffer.push(token_seq.detach().cpu())
|
||||
|
||||
stats = {
|
||||
'loss': loss.item(),
|
||||
'replay_loss': replay_loss.item() if isinstance(replay_loss, torch.Tensor) else replay_loss,
|
||||
'total_loss': total_loss,
|
||||
'w_norm_mean': info['w_norm'].mean().item(),
|
||||
'alpha_mean': info['alpha'].mean(dim=(0, 1)).detach().cpu().tolist(),
|
||||
}
|
||||
return stats
|
||||
|
||||
def consolidate(self, data_sample: torch.Tensor):
|
||||
"""周期性 consolidate EWC——更新参数重要性锚点
|
||||
|
||||
在学完一段文本后调用,把当前知识"固化"
|
||||
"""
|
||||
self.ewc_optimizer.consolidate(data_sample.to(self.device), n_samples=20)
|
||||
|
||||
def evolve(self, text_stream: list[str], tokenizer,
|
||||
seq_len: int = 64, batch_size: int = 4,
|
||||
consolidate_every: int = 20,
|
||||
generate_every: int = 50,
|
||||
max_steps: int | None = None,
|
||||
prompt_text: str = "To be",
|
||||
on_progress: Callable | None = None) -> list[dict]:
|
||||
"""
|
||||
持续进化主循环
|
||||
|
||||
Args:
|
||||
text_stream: 文本片段列表(模拟持续到来的数据流)
|
||||
tokenizer: CharTokenizer
|
||||
seq_len: 序列长度
|
||||
batch_size: 每次喂入的 batch 大小
|
||||
consolidate_every: 每隔多少步 consolidate EWC
|
||||
generate_every: 每隔多少步生成样本观察
|
||||
prompt_text: 生成样本的提示词
|
||||
on_progress: 回调函数,返回当前进度
|
||||
|
||||
Returns:
|
||||
history: 进化历史
|
||||
"""
|
||||
step = 0
|
||||
all_tokens = []
|
||||
|
||||
# 把所有文本编码成 token 流
|
||||
for text in text_stream:
|
||||
tokens = tokenizer.encode(text)
|
||||
all_tokens.extend(tokens)
|
||||
|
||||
# 用滑动窗口在完整 token 流上采样 batch
|
||||
# 每个 batch 包含 batch_size 条序列,每条长 seq_len
|
||||
# 相邻 batch 之间步进 stride 个序列
|
||||
stride = batch_size # 每个 batch 用 batch_size 个新起点
|
||||
n_possible_starts = max(0, len(all_tokens) - seq_len - 1)
|
||||
n_batches = max(0, n_possible_starts // stride)
|
||||
|
||||
print(f"自主进化开始:{len(all_tokens)} tokens, {len(text_stream)} 段文本")
|
||||
print(f"配置: seq_len={seq_len}, batch_size={batch_size}, "
|
||||
f"stride={stride}, n_batches={n_batches}")
|
||||
print("=" * 70)
|
||||
|
||||
for batch_idx in range(n_batches):
|
||||
# 取一个 batch 的序列(滑动窗口)
|
||||
batch_tokens = []
|
||||
for i in range(batch_size):
|
||||
start = batch_idx * stride + i
|
||||
if start + seq_len >= len(all_tokens):
|
||||
batch_tokens.append(all_tokens[-seq_len:])
|
||||
else:
|
||||
batch_tokens.append(all_tokens[start:start + seq_len])
|
||||
|
||||
if any(len(b) < seq_len for b in batch_tokens):
|
||||
continue
|
||||
|
||||
token_seq = torch.tensor(batch_tokens, dtype=torch.long)
|
||||
stats = self.learn_step(token_seq)
|
||||
stats['step'] = step
|
||||
stats['batch_idx'] = batch_idx
|
||||
self.history.append(stats)
|
||||
|
||||
# 周期性 consolidate
|
||||
if step > 0 and step % consolidate_every == 0:
|
||||
self.consolidate(token_seq)
|
||||
|
||||
# 周期性生成 + 报告
|
||||
if step % generate_every == 0 or step == n_batches - 1:
|
||||
prompt_ids = tokenizer.encode(prompt_text)
|
||||
generated = self.model.generate(
|
||||
prompt_ids, n_new=80, temperature=0.8, top_k=5
|
||||
)
|
||||
sample = prompt_text + tokenizer.decode(generated)
|
||||
stats['sample'] = sample
|
||||
|
||||
print(f"\n[step {step:4d}] loss={stats['loss']:.4f} "
|
||||
f"replay={stats['replay_loss']:.4f} "
|
||||
f"||w||={stats['w_norm_mean']:.3f}")
|
||||
print(f" 专家使用率: {[f'{u:.2f}' for u in self.plasticity.usage.tolist()]}")
|
||||
print(f" 生成样本: {repr(sample[:120])}...")
|
||||
|
||||
if on_progress:
|
||||
on_progress(stats)
|
||||
|
||||
step += 1
|
||||
|
||||
if max_steps is not None and step >= max_steps:
|
||||
break
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("自主进化完成")
|
||||
return self.history
|
||||
|
||||
def get_evolution_summary(self) -> dict:
|
||||
"""获取进化总结"""
|
||||
if not self.history:
|
||||
return {}
|
||||
|
||||
losses = [h['loss'] for h in self.history]
|
||||
return {
|
||||
'steps': len(self.history),
|
||||
'final_loss': losses[-1],
|
||||
'initial_loss': losses[0],
|
||||
'min_loss': min(losses),
|
||||
'expert_usage': self.plasticity.usage.tolist(),
|
||||
'expert_specialization': self.plasticity.get_stats()['top_specialization'],
|
||||
'samples': [h.get('sample', '') for h in self.history if 'sample' in h],
|
||||
}
|
||||
264
jspaceai/jlens.py
Normal file
264
jspaceai/jlens.py
Normal file
@@ -0,0 +1,264 @@
|
||||
"""
|
||||
Jacobian Lens (J-lens) —— 观测模型内部 J-space 的可解释性工具
|
||||
|
||||
灵感来自 Anthropic 2026 论文 "Verbalizable Representations Form a Global
|
||||
Workspace in Language Models"。
|
||||
|
||||
核心思想:
|
||||
J-lens 计算中间层激活对最终输出的平均因果效应。
|
||||
J_ℓ = E[∂h_final / ∂h_ℓ] —— 跨大量 context 平均的 Jacobian
|
||||
lens(h_ℓ) = softmax(W_U · norm(J_ℓ · h_ℓ))
|
||||
|
||||
J-lens 向量 = W_U · J_ℓ 的行,每个向量对应词汇表中的一个 token。
|
||||
一个激活向量在 J-lens 下的 top tokens = 模型"准备要说"的概念。
|
||||
|
||||
与 logit lens 的区别:
|
||||
logit lens 直接用 W_U 投影(假设 J_ℓ = I)。
|
||||
J-lens 修正了层间表征变化,能在更早的层揭示可解释内容。
|
||||
|
||||
在我们的 ODE 架构中:
|
||||
- 每个 ODE 子步是一个"层"
|
||||
- workspace w 在每个子步演化
|
||||
- J-lens 可以在任意子步读 w,揭示模型在该时刻"在想什么"
|
||||
|
||||
实现简化:
|
||||
- 完整 J-lens 需要 backprop 从输出到中间层,在我们的 ODE 模型中代价高
|
||||
- 我们用近似:直接训练一个 lens matrix L_ℓ,让 lens(h) ≈ output
|
||||
- L_ℓ 通过最小化 ||output - L_ℓ · h||² 在数据上学习
|
||||
- 这等价于 tuned lens,但在我们的架构中更高效
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class JLensConfig:
|
||||
"""J-lens 配置"""
|
||||
n_substeps: int = 4 # ODE 子步数(对应"层"数)
|
||||
workspace_dim: int = 32 # 工作空间维度
|
||||
vocab_size: int = 100 # 词汇表大小
|
||||
top_k: int = 10 # 默认 top-k 读出
|
||||
|
||||
|
||||
class JLensProbe(nn.Module):
|
||||
"""
|
||||
单个子步的 J-lens 探针。
|
||||
|
||||
学习一个线性映射 L: workspace_dim → vocab_size,
|
||||
使得 lens(w) ≈ model_output。
|
||||
|
||||
这近似了 J_ℓ · W_U(Jacobian 与 unembedding 的复合)。
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim: int, vocab_size: int):
|
||||
super().__init__()
|
||||
self.lens = nn.Linear(workspace_dim, vocab_size, bias=False)
|
||||
# 用 output_head 的权重初始化(如果可用)
|
||||
nn.init.normal_(self.lens.weight, std=0.02)
|
||||
|
||||
def forward(self, w: torch.Tensor) -> torch.Tensor:
|
||||
"""w → logits (近似 J-lens 读出)"""
|
||||
return self.lens(w)
|
||||
|
||||
def top_tokens(self, w: torch.Tensor, idx_to_char: dict,
|
||||
top_k: int = 10) -> list[tuple[str, float]]:
|
||||
"""获取 top-k token 及其概率"""
|
||||
logits = self.forward(w)
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
topk_probs, topk_idx = probs.topk(top_k)
|
||||
return [(idx_to_char.get(i.item(), '?'), p.item())
|
||||
for i, p in zip(topk_idx[0], topk_probs[0])]
|
||||
|
||||
|
||||
class JLensSuite(nn.Module):
|
||||
"""
|
||||
J-lens 套件:每个 ODE 子步一个探针。
|
||||
|
||||
在 forward 时记录每个子步的 w,用对应探针读出。
|
||||
训练时让每个探针预测最终输出。
|
||||
|
||||
结构对应 Anthropic 论文的三层分层:
|
||||
子步 0-0: sensory(输入处理,J-lens 噪声大)
|
||||
子步 1-2: workspace(抽象思考,J-lens 可解释)
|
||||
子步 3: motor(输出准备,J-lens ≈ output)
|
||||
"""
|
||||
|
||||
def __init__(self, config: JLensConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.probes = nn.ModuleList([
|
||||
JLensProbe(config.workspace_dim, config.vocab_size)
|
||||
for _ in range(config.n_substeps)
|
||||
])
|
||||
|
||||
def forward(self, w_trajectory: list[torch.Tensor]) -> list[torch.Tensor]:
|
||||
"""
|
||||
对轨迹中每个 w 读出 logits
|
||||
|
||||
Args:
|
||||
w_trajectory: list of (batch, workspace_dim),每个子步的 w
|
||||
|
||||
Returns:
|
||||
list of (batch, vocab_size),每个子步的 lens 读出
|
||||
"""
|
||||
return [probe(w) for probe, w in zip(self.probes, w_trajectory)]
|
||||
|
||||
def train_on_trajectory(self, w_trajectory: list[torch.Tensor],
|
||||
target_ids: torch.Tensor) -> float:
|
||||
"""
|
||||
训练探针:让每个子步的 lens 读出预测最终 target token
|
||||
|
||||
Args:
|
||||
w_trajectory: 每个子步的 w
|
||||
target_ids: 目标 token ids (batch,)
|
||||
|
||||
Returns:
|
||||
平均 loss
|
||||
"""
|
||||
total_loss = 0.0
|
||||
for probe, w in zip(self.probes, w_trajectory):
|
||||
pred = probe(w)
|
||||
loss = F.cross_entropy(pred, target_ids)
|
||||
total_loss += loss
|
||||
return total_loss / len(w_trajectory)
|
||||
|
||||
|
||||
class WorkspaceAblator:
|
||||
"""
|
||||
Workspace ablation 工具——验证 selectivity。
|
||||
|
||||
ablate workspace 后看哪些能力受损:
|
||||
- 自动任务(续写、分类)应该不受影响
|
||||
- 灵活推理(多跳、规划)应该受损
|
||||
|
||||
实现:在 forward 时把 w 的 top-k J-lens 方向投影掉。
|
||||
"""
|
||||
|
||||
def __init__(self, model, lens_suite: JLensSuite):
|
||||
self.model = model
|
||||
self.lens_suite = lens_suite
|
||||
|
||||
@torch.no_grad()
|
||||
def get_top_lens_directions(self, w: torch.Tensor, k: int = 5) -> torch.Tensor:
|
||||
"""获取 w 当前激活最强的 k 个 J-lens 方向"""
|
||||
# 用第一个 workspace 探针(中间子步)
|
||||
probe = self.lens_suite.probes[len(self.lens_suite.probes) // 2]
|
||||
logits = probe(w) # (batch, vocab)
|
||||
# top-k token 的 lens 向量(probe.lens.weight 的行)
|
||||
topk_vals, topk_idx = logits.topk(k, dim=-1) # (batch, k)
|
||||
# 获取这些 token 对应的 lens 方向
|
||||
# probe.lens.weight: (vocab, workspace_dim)
|
||||
directions = probe.lens.weight[topk_idx] # (batch, k, workspace_dim)
|
||||
return directions
|
||||
|
||||
def ablate_workspace(self, w: torch.Tensor, k: int = 5) -> torch.Tensor:
|
||||
"""
|
||||
Ablate workspace 的 top-k J-lens 方向
|
||||
|
||||
把 w 在这些方向上的投影去掉,保留正交分量。
|
||||
"""
|
||||
directions = self.get_top_lens_directions(w, k) # (batch, k, workspace_dim)
|
||||
# 对每个方向,投影掉
|
||||
w_ablated = w.clone()
|
||||
for b in range(w.shape[0]):
|
||||
for d in directions[b]: # (workspace_dim,)
|
||||
d_norm = d / (d.norm() + 1e-8)
|
||||
proj = (w_ablated[b] @ d_norm) * d_norm
|
||||
w_ablated[b] = w_ablated[b] - proj
|
||||
return w_ablated
|
||||
|
||||
|
||||
class DirectedModulation:
|
||||
"""
|
||||
Directed Modulation——让模型被指令"想某概念"。
|
||||
|
||||
实现:在 forward 时给 workspace w 注入一个概念向量。
|
||||
这个向量从 J-lens 的某个 token 方向获取。
|
||||
|
||||
对应 Anthropic 论文实验:
|
||||
"concentrate on citrus fruits" → orange 出现在 J-lens
|
||||
即使输出在抄无关文本,workspace 里装的是被指令的概念。
|
||||
|
||||
机制:
|
||||
1. 获取概念 token 的 J-lens 向量 v_concept
|
||||
2. 在 forward 时给 w 加上 α · v_concept
|
||||
3. 模型输出会受这个注入影响
|
||||
"""
|
||||
|
||||
def __init__(self, model, lens_suite: JLensSuite):
|
||||
self.model = model
|
||||
self.lens_suite = lens_suite
|
||||
|
||||
def get_concept_vector(self, token_id: int, substep: int = None) -> torch.Tensor:
|
||||
"""获取某个 token 在 J-lens 中的方向向量"""
|
||||
if substep is None:
|
||||
substep = len(self.lens_suite.probes) // 2
|
||||
probe = self.lens_suite.probes[substep]
|
||||
# probe.lens.weight: (vocab, workspace_dim)
|
||||
return probe.lens.weight[token_id].clone()
|
||||
|
||||
def modulate_state(self, state: dict, token_id: int,
|
||||
strength: float = 1.0, substep: int = None) -> dict:
|
||||
"""给 state 的 workspace 注入概念向量"""
|
||||
concept_vec = self.get_concept_vector(token_id, substep) # (workspace_dim,)
|
||||
new_state = {
|
||||
'w': state['w'] + strength * concept_vec.unsqueeze(0), # (1, workspace_dim)
|
||||
'm': state['m'],
|
||||
}
|
||||
return new_state
|
||||
|
||||
|
||||
class CounterfactualReflection:
|
||||
"""
|
||||
Counterfactual Reflection Training——通过塑造 J-space 来塑造行为。
|
||||
|
||||
灵感:Anthropic 论文发现,训练模型"如果被打断要反思什么原则",
|
||||
会让它在正常工作时也遵守这些原则——因为训练塑造了 J-space 内容。
|
||||
|
||||
实现:
|
||||
1. 正常 forward 产生输出
|
||||
2. 在输出后追加"反思提示"(如"反思:我应该...")
|
||||
3. 让模型在反思提示下生成反思内容
|
||||
4. 用反思内容的 loss 反向传播,更新模型
|
||||
|
||||
这让模型的 J-space 在相关 context 下自然装载这些原则。
|
||||
"""
|
||||
|
||||
def __init__(self, model, tokenizer, reflection_prompt: str = "\nReflect: "):
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
self.reflection_prompt_ids = tokenizer.encode(reflection_prompt)
|
||||
|
||||
def create_reflection_sequence(self, input_ids: list[int],
|
||||
reflection_target_ids: list[int]) -> list[int]:
|
||||
"""创建反思训练序列:input + reflection_prompt + reflection_target"""
|
||||
return input_ids + self.reflection_prompt_ids + reflection_target_ids
|
||||
|
||||
def compute_reflection_loss(self, input_seq: torch.Tensor,
|
||||
reflection_target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
计算反思训练 loss
|
||||
|
||||
Args:
|
||||
input_seq: (batch, T) 包含 input + reflection_prompt
|
||||
reflection_target: (batch, T_reflect) 期望的反思内容
|
||||
|
||||
Returns:
|
||||
loss
|
||||
"""
|
||||
# forward 整个序列
|
||||
logits, _ = self.model(input_seq)
|
||||
|
||||
# 反思部分是序列末尾的 reflection_target 长度
|
||||
T_reflect = reflection_target.shape[1]
|
||||
reflect_logits = logits[:, -T_reflect:] # (batch, T_reflect, vocab)
|
||||
|
||||
return F.cross_entropy(
|
||||
reflect_logits.reshape(-1, self.model.config.vocab_size),
|
||||
reflection_target.reshape(-1),
|
||||
)
|
||||
147
jspaceai/language_data.py
Normal file
147
jspaceai/language_data.py
Normal file
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
字符级 tokenizer + 数据集
|
||||
|
||||
为什么字符级而不是 subword:
|
||||
- 词汇表小(~100),模型可以小,验证架构用
|
||||
- 字符级有明确的"风格"信号(拼写、标点、节奏)
|
||||
- 持续学习场景下,subword 词汇表会变,字符级稳定
|
||||
|
||||
数据:用经典 Shakespeare 文本作为持续学习的语料。
|
||||
也可以换成任何 UTF-8 文本。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
@dataclass
|
||||
class CharTokenizer:
|
||||
"""字符级 tokenizer,支持训练时新字符的动态加入"""
|
||||
chars: list[str]
|
||||
char_to_idx: dict[str, int]
|
||||
idx_to_char: dict[int, str]
|
||||
|
||||
@classmethod
|
||||
def from_text(cls, text: str) -> "CharTokenizer":
|
||||
chars = sorted(set(text))
|
||||
char_to_idx = {c: i for i, c in enumerate(chars)}
|
||||
idx_to_char = {i: c for i, c in enumerate(chars)}
|
||||
return cls(chars, char_to_idx, idx_to_char)
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return len(self.chars)
|
||||
|
||||
def encode(self, text: str) -> list[int]:
|
||||
# 未知字符用 0(假设第一个字符是常见的,或后续加 unk token)
|
||||
return [self.char_to_idx.get(c, 0) for c in text]
|
||||
|
||||
def decode(self, ids: list[int]) -> str:
|
||||
return "".join(self.idx_to_char.get(i, "") for i in ids)
|
||||
|
||||
|
||||
class CharDataset(Dataset):
|
||||
"""字符级 next-char 预测数据集
|
||||
|
||||
每个样本:(input_seq, target_seq) 长度均为 seq_len
|
||||
target[i] = input[i+1]
|
||||
"""
|
||||
|
||||
def __init__(self, text: str, seq_len: int = 64, tokenizer: CharTokenizer | None = None):
|
||||
self.seq_len = seq_len
|
||||
if tokenizer is None:
|
||||
self.tokenizer = CharTokenizer.from_text(text)
|
||||
else:
|
||||
self.tokenizer = tokenizer
|
||||
self.data = self.tokenizer.encode(text)
|
||||
|
||||
def __len__(self):
|
||||
return max(0, len(self.data) - self.seq_len - 1)
|
||||
|
||||
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
chunk = self.data[idx:idx + self.seq_len + 1]
|
||||
x = torch.tensor(chunk[:-1], dtype=torch.long)
|
||||
y = torch.tensor(chunk[1:], dtype=torch.long)
|
||||
return x, y
|
||||
|
||||
|
||||
def load_shakespeare(data_dir: Path | None = None) -> str:
|
||||
"""加载 Shakespeare 文本。
|
||||
|
||||
如果 data_dir 不存在或没有文本,返回一个内嵌的小样本。
|
||||
"""
|
||||
if data_dir is not None:
|
||||
text_path = data_dir / "shakespeare.txt"
|
||||
if text_path.exists():
|
||||
return text_path.read_text(encoding="utf-8")
|
||||
|
||||
# 内嵌样本(足够小但能展现语言结构)
|
||||
return """To be, or not to be, that is the question:
|
||||
Whether 'tis nobler in the mind to suffer
|
||||
The slings and arrows of outrageous fortune,
|
||||
Or to take arms against a sea of troubles
|
||||
And by opposing end them. To die—to sleep,
|
||||
No more; and by a sleep to say we end
|
||||
The heart-ache and the thousand natural shocks
|
||||
That flesh is heir to: 'tis a consummation
|
||||
Devoutly to be wish'd. To die, to sleep;
|
||||
To sleep, perchance to dream—ay, there's the rub:
|
||||
For in that sleep of death what dreams may come,
|
||||
When we have shuffled off this mortal coil,
|
||||
Must give us pause—there's the respect
|
||||
That makes calamity of so long life.
|
||||
|
||||
Romeo, Romeo! wherefore art thou Romeo?
|
||||
Deny thy father and refuse thy name;
|
||||
Or, if thou wilt not, be but sworn my love,
|
||||
And I'll no longer be a Capulet.
|
||||
|
||||
O Romeo, Romeo! wherefore art thou Romeo?
|
||||
'Tis but thy name that is my enemy;
|
||||
Thou art thyself, though not a Montague.
|
||||
What's Montague? It is nor hand, nor foot,
|
||||
Nor arm, nor face, nor any other part
|
||||
Belonging to a man. O, be some other name!
|
||||
What's in a name? that which we call a rose
|
||||
By any other name would smell as sweet;
|
||||
So Romeo would, were he not Romeo call'd,
|
||||
Retain that dear perfection which he owes
|
||||
Without that title. Romeo, doff thy name,
|
||||
And for that name which is no part of thee
|
||||
Take all myself.
|
||||
|
||||
Friends, Romans, countrymen, lend me your ears;
|
||||
I come to bury Caesar, not to praise him.
|
||||
The evil that men do lives after them;
|
||||
The good is oft interred with their bones;
|
||||
So let it be with Caesar. The noble Brutus
|
||||
Hath told you Caesar was ambitious:
|
||||
If it were so, it was a grievous fault,
|
||||
And grievously hath Caesar answer'd it.
|
||||
|
||||
The course of true love never did run smooth;
|
||||
But, either it was different in blood,
|
||||
Or else misgraffed in respect of years.
|
||||
|
||||
If music be the food of love, play on,
|
||||
Give me excess of it, that, surfeiting,
|
||||
The appetite may sicken, and so die.
|
||||
That strain again! it had a dying fall:
|
||||
O, it came o'er my ear like the sweet sound
|
||||
That breathes upon a bank of violets,
|
||||
Stealing and giving odour.
|
||||
|
||||
All the world's a stage,
|
||||
And all the men and women merely players:
|
||||
They have their exits and their entrances;
|
||||
And one man in his time plays many parts,
|
||||
His acts being seven ages.
|
||||
|
||||
Now is the winter of our discontent
|
||||
Made glorious summer by this sun of York;
|
||||
And all the clouds that lour'd upon our house
|
||||
In the deep bosom of the ocean buried.
|
||||
"""
|
||||
425
jspaceai/language_model.py
Normal file
425
jspaceai/language_model.py
Normal file
@@ -0,0 +1,425 @@
|
||||
"""
|
||||
语言版 JSpace 模型 + 自主进化机制
|
||||
|
||||
核心扩展:
|
||||
1. JSpaceLanguageModel: 在 JSpaceModel 基础上加 token embedding + logit 输出
|
||||
2. 在线学习:每个 forward 累积梯度并更新参数(边推理边学习)
|
||||
3. EWC(Elastic Weight Consolidation):保护重要参数,防灾难性遗忘
|
||||
4. ExperienceReplay:经验回放缓冲区
|
||||
5. ExpertPlasticity:专家专业化追踪,新知识优先路由到"空闲"专家
|
||||
|
||||
数学形式:
|
||||
|
||||
专家动力学(同 core.py):
|
||||
dm_i/dt = -∇U_i(m_i) + J_i · w + P_i_in · embed(x)
|
||||
|
||||
工作空间动力学(同 core.py):
|
||||
τ_w · dw/dt = -w + Σ α_i · P_i_out(m_i)
|
||||
|
||||
输出:
|
||||
logits = Q(w) # 从工作空间投影到词汇表
|
||||
|
||||
在线学习目标:
|
||||
L = -log p(x_{t+1} | x_{0:t}) + λ_EWC · Σ_i F_i · (θ_i - θ*_i)²
|
||||
|
||||
EWC 中 F_i 是 Fisher 信息矩阵对角线,衡量参数重要性。
|
||||
重要参数被"锚定"在旧值附近,新知识只能修改不重要的参数。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass, field
|
||||
from collections import deque
|
||||
import random
|
||||
|
||||
from .core import JSpaceConfig, Expert, JSpaceWorkspace
|
||||
|
||||
|
||||
@dataclass
|
||||
class LanguageConfig(JSpaceConfig):
|
||||
"""语言模型配置,继承 JSpaceConfig"""
|
||||
vocab_size: int = 100 # 词汇表大小
|
||||
embed_dim: int = 16 # token embedding 维度(≠ input_dim,会投影)
|
||||
# input_dim 仍用 JSpaceConfig 的,作为工作空间接收的输入维度
|
||||
|
||||
|
||||
class JSpaceLanguageModel(nn.Module):
|
||||
"""
|
||||
语言版 JSpace 模型。
|
||||
|
||||
流程:
|
||||
token → embedding → 投影到 input_dim → 喂入 JSpace 动力学
|
||||
工作空间 w → 投影到 vocab_size → logits → 采样 token
|
||||
|
||||
每步可训练(在线学习):
|
||||
forward 后用 cross-entropy loss 更新参数
|
||||
"""
|
||||
|
||||
def __init__(self, config: LanguageConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
# Token embedding
|
||||
self.embedding = nn.Embedding(config.vocab_size, config.embed_dim)
|
||||
# 投影 embedding → input_dim(JSpace 期望的输入维度)
|
||||
self.input_proj = nn.Linear(config.embed_dim, config.input_dim, bias=False)
|
||||
|
||||
# 专家池
|
||||
self.experts = nn.ModuleList([
|
||||
Expert(
|
||||
expert_dim=config.expert_dim,
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_wells=config.num_wells,
|
||||
sparsity=config.jacobian_sparsity,
|
||||
)
|
||||
for _ in range(config.num_experts)
|
||||
])
|
||||
|
||||
# 工作空间
|
||||
self.workspace = JSpaceWorkspace(
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_experts=config.num_experts,
|
||||
)
|
||||
|
||||
# 输出头:w → logits
|
||||
self.output_head = nn.Sequential(
|
||||
nn.Linear(config.workspace_dim, 64),
|
||||
nn.Tanh(),
|
||||
nn.Linear(64, config.vocab_size),
|
||||
)
|
||||
|
||||
def init_state(self, batch_size: int, device: torch.device) -> dict:
|
||||
return {
|
||||
'w': torch.zeros(batch_size, self.config.workspace_dim, device=device),
|
||||
'm': [torch.zeros(batch_size, self.config.expert_dim, device=device)
|
||||
for _ in range(self.config.num_experts)],
|
||||
}
|
||||
|
||||
def step(self, state: dict, token_ids: torch.Tensor,
|
||||
record_trajectory: bool = False) -> tuple[dict, torch.Tensor, torch.Tensor, list]:
|
||||
"""单时间步前向
|
||||
|
||||
Args:
|
||||
state: {'w': ..., 'm': [...]}
|
||||
token_ids: (batch,) token indices
|
||||
record_trajectory: 是否记录 w 轨迹(J-lens 用)
|
||||
|
||||
Returns:
|
||||
new_state, logits (batch, vocab_size), alpha (batch, num_experts),
|
||||
w_trajectory (list of (batch, workspace_dim)) 或空 list
|
||||
"""
|
||||
w = state['w']
|
||||
ms = state['m']
|
||||
cfg = self.config
|
||||
w_trajectory = []
|
||||
|
||||
# token → embedding → input projection
|
||||
emb = self.embedding(token_ids) # (batch, embed_dim)
|
||||
x = self.input_proj(emb) # (batch, input_dim)
|
||||
|
||||
# ODE 子步积分(对应 Anthropic 论文的"层")
|
||||
for substep in range(cfg.ode_steps):
|
||||
contributions = []
|
||||
new_ms = []
|
||||
for i, expert in enumerate(self.experts):
|
||||
m_next, contrib = expert(
|
||||
ms[i], w, x,
|
||||
dt=cfg.dt, noise_std=cfg.noise_std,
|
||||
)
|
||||
new_ms.append(m_next)
|
||||
contributions.append(contrib)
|
||||
contributions = torch.stack(contributions, dim=1)
|
||||
|
||||
w, alpha = self.workspace(
|
||||
w, x, contributions,
|
||||
dt=cfg.dt, tau_w=cfg.tau_w,
|
||||
)
|
||||
ms = new_ms
|
||||
|
||||
if record_trajectory:
|
||||
w_trajectory.append(w.detach())
|
||||
|
||||
logits = self.output_head(w) # (batch, vocab_size)
|
||||
new_state = {'w': w, 'm': ms}
|
||||
return new_state, logits, alpha, w_trajectory
|
||||
|
||||
def forward(self, token_seqs: torch.Tensor, state: dict | None = None,
|
||||
record_trajectory: bool = False) -> tuple[torch.Tensor, dict]:
|
||||
"""
|
||||
Args:
|
||||
token_seqs: (batch, T) token indices
|
||||
record_trajectory: 是否记录 w 轨迹
|
||||
|
||||
Returns:
|
||||
logits: (batch, T, vocab_size)
|
||||
info: {'alpha': (batch, T, num_experts), 'w_norm': (batch, T),
|
||||
'w_trajectory': list of (batch, T, workspace_dim) 或空}
|
||||
"""
|
||||
batch_size, T = token_seqs.shape
|
||||
device = token_seqs.device
|
||||
|
||||
if state is None:
|
||||
state = self.init_state(batch_size, device)
|
||||
|
||||
logits_list = []
|
||||
alphas = []
|
||||
w_norms = []
|
||||
w_traj_per_step = [] # list of (list of substep w)
|
||||
for t in range(T):
|
||||
state, logits, alpha, w_traj = self.step(
|
||||
state, token_seqs[:, t], record_trajectory=record_trajectory
|
||||
)
|
||||
logits_list.append(logits)
|
||||
alphas.append(alpha)
|
||||
w_norms.append(state['w'].norm(dim=-1))
|
||||
if record_trajectory and w_traj:
|
||||
w_traj_per_step.append(torch.stack(w_traj, dim=1)) # (batch, n_substeps, workspace_dim)
|
||||
|
||||
logits = torch.stack(logits_list, dim=1) # (batch, T, vocab_size)
|
||||
info = {
|
||||
'alpha': torch.stack(alphas, dim=1), # (batch, T, num_experts)
|
||||
'w_norm': torch.stack(w_norms, dim=1), # (batch, T)
|
||||
}
|
||||
if record_trajectory and w_traj_per_step:
|
||||
info['w_trajectory'] = torch.stack(w_traj_per_step, dim=1) # (batch, T, n_substeps, workspace_dim)
|
||||
return logits, info
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, prompt: list[int], n_new: int = 50, temperature: float = 1.0,
|
||||
top_k: int = 5) -> list[int]:
|
||||
"""自回归生成
|
||||
|
||||
Args:
|
||||
prompt: 起始 token ids
|
||||
n_new: 生成的新 token 数
|
||||
temperature: 采样温度
|
||||
top_k: top-k 采样
|
||||
"""
|
||||
self.eval()
|
||||
device = next(self.parameters()).device
|
||||
state = self.init_state(1, device)
|
||||
|
||||
# 预热 state with prompt
|
||||
tokens = list(prompt)
|
||||
for tok in tokens:
|
||||
state, _, _, _ = self.step(state, torch.tensor([tok], device=device))
|
||||
|
||||
# 生成
|
||||
generated = []
|
||||
for _ in range(n_new):
|
||||
state, logits, _, _ = self.step(state, torch.tensor([tokens[-1]], device=device))
|
||||
logits = logits[0] / max(temperature, 1e-6)
|
||||
|
||||
if top_k > 0:
|
||||
top_k = min(top_k, logits.size(-1))
|
||||
vals, idxs = logits.topk(top_k)
|
||||
probs = F.softmax(vals, dim=-1)
|
||||
next_tok = idxs[torch.multinomial(probs, 1)].item()
|
||||
else:
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
next_tok = torch.multinomial(probs, 1).item()
|
||||
|
||||
generated.append(next_tok)
|
||||
tokens.append(next_tok)
|
||||
|
||||
self.train()
|
||||
return generated
|
||||
|
||||
|
||||
class ExperienceReplay:
|
||||
"""
|
||||
经验回放缓冲区。
|
||||
|
||||
存储见过的序列片段,训练时随机采样混入当前 batch。
|
||||
防止灾难性遗忘——旧知识被定期"复习"。
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int = 1000, seq_len: int = 64):
|
||||
self.capacity = capacity
|
||||
self.seq_len = seq_len
|
||||
self.buffer: deque = deque(maxlen=capacity)
|
||||
|
||||
def push(self, token_seq: torch.Tensor):
|
||||
"""push 一个序列 (T,) 或 (batch, T)"""
|
||||
if token_seq.dim() == 1:
|
||||
token_seq = token_seq.unsqueeze(0)
|
||||
for seq in token_seq:
|
||||
if len(seq) >= self.seq_len:
|
||||
self.buffer.append(seq.clone())
|
||||
|
||||
def sample(self, batch_size: int) -> torch.Tensor | None:
|
||||
"""采样 (batch_size, seq_len)"""
|
||||
if len(self.buffer) < batch_size:
|
||||
return None
|
||||
samples = random.sample(list(self.buffer), batch_size)
|
||||
# 随机裁剪到 seq_len
|
||||
result = []
|
||||
for s in samples:
|
||||
if len(s) > self.seq_len:
|
||||
start = random.randint(0, len(s) - self.seq_len - 1)
|
||||
result.append(s[start:start + self.seq_len])
|
||||
else:
|
||||
result.append(s)
|
||||
return torch.stack(result)
|
||||
|
||||
|
||||
class EWCOptimizer:
|
||||
"""
|
||||
Elastic Weight Consolidation 优化器包装。
|
||||
|
||||
核心思想:参数 θ 有"重要性" F(Fisher 信息)。
|
||||
重要参数偏离原值 θ* 会被惩罚。
|
||||
新知识只能修改不重要的参数。
|
||||
|
||||
L_total = L_task + λ · Σ_i F_i · (θ_i - θ*_i)²
|
||||
|
||||
工作流:
|
||||
1. 正常训练一段时间
|
||||
2. 调用 consolidate():计算 Fisher 信息,锚定当前参数
|
||||
3. 继续训练——loss 中加入 EWC 正则
|
||||
4. 周期性 consolidate(更新锚点和重要性)
|
||||
"""
|
||||
|
||||
def __init__(self, model: nn.Module, lr: float = 1e-3,
|
||||
ewc_lambda: float = 1.0, max_grad_norm: float = 1.0):
|
||||
self.model = model
|
||||
self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
|
||||
self.ewc_lambda = ewc_lambda
|
||||
self.max_grad_norm = max_grad_norm
|
||||
|
||||
# Fisher 信息和锚定参数
|
||||
self.fisher: dict[str, torch.Tensor] = {}
|
||||
self.anchored_params: dict[str, torch.Tensor] = {}
|
||||
|
||||
def consolidate(self, data_sample: torch.Tensor, n_samples: int = 50):
|
||||
"""计算 Fisher 信息并锚定当前参数
|
||||
|
||||
Args:
|
||||
data_sample: (batch, T) 用于计算 Fisher 的数据样本
|
||||
n_samples: 采样次数(Fisher 信息的 Monte Carlo 估计)
|
||||
"""
|
||||
# 保存当前参数作为锚点
|
||||
self.anchored_params = {
|
||||
name: param.data.clone()
|
||||
for name, param in self.model.named_parameters()
|
||||
}
|
||||
|
||||
# 计算 Fisher 信息(对角近似)
|
||||
fisher = {
|
||||
name: torch.zeros_like(param)
|
||||
for name, param in self.model.named_parameters()
|
||||
}
|
||||
|
||||
self.model.eval()
|
||||
for _ in range(n_samples):
|
||||
self.model.zero_grad()
|
||||
logits, _ = self.model(data_sample)
|
||||
# 只用 logits[:, :-1] 对应 targets[:, 1:] 的部分
|
||||
logits_pred = logits[:, :-1] # (batch, T-1, vocab)
|
||||
probs = F.softmax(logits_pred, dim=-1) # (batch, T-1, vocab)
|
||||
# 采样 token 计算 Fisher
|
||||
sampled_tokens = torch.multinomial(
|
||||
probs.reshape(-1, probs.size(-1)), 1
|
||||
).view_as(logits_pred[:, :, 0]) # (batch, T-1)
|
||||
log_probs = F.log_softmax(logits_pred, dim=-1)
|
||||
loss = -log_probs.gather(-1, sampled_tokens.unsqueeze(-1)).mean()
|
||||
loss.backward()
|
||||
|
||||
for name, param in self.model.named_parameters():
|
||||
if param.grad is not None:
|
||||
fisher[name] += param.grad.data ** 2
|
||||
|
||||
# 平均
|
||||
for name in fisher:
|
||||
fisher[name] /= n_samples
|
||||
|
||||
self.fisher = fisher
|
||||
self.model.zero_grad()
|
||||
self.model.train()
|
||||
|
||||
def ewc_penalty(self) -> torch.Tensor:
|
||||
"""计算 EWC 正则项"""
|
||||
if not self.fisher:
|
||||
return torch.tensor(0.0, device=next(self.model.parameters()).device)
|
||||
|
||||
penalty = 0.0
|
||||
for name, param in self.model.named_parameters():
|
||||
if name in self.fisher:
|
||||
penalty = penalty + (self.fisher[name] * (param - self.anchored_params[name]) ** 2).sum()
|
||||
return penalty
|
||||
|
||||
def step(self, loss: torch.Tensor):
|
||||
"""一步优化:task loss + EWC 正则"""
|
||||
total_loss = loss + self.ewc_lambda * self.ewc_penalty()
|
||||
self.optimizer.zero_grad()
|
||||
total_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
return total_loss.item()
|
||||
|
||||
|
||||
class ExpertPlasticity:
|
||||
"""
|
||||
专家结构可塑性追踪。
|
||||
|
||||
追踪每个专家的"专业化程度":
|
||||
- 哪些专家在处理哪些模式
|
||||
- 哪些专家"负载过重"(应该分裂或新增)
|
||||
- 哪些专家"空闲"(可以接收新知识)
|
||||
|
||||
这不是真正的动态增删专家(实现复杂),而是:
|
||||
- 统计专家使用率
|
||||
- 在路由时给空闲专家加权(鼓励新知识流向空闲专家)
|
||||
"""
|
||||
|
||||
def __init__(self, num_experts: int, ema_alpha: float = 0.99):
|
||||
self.num_experts = num_experts
|
||||
self.ema_alpha = ema_alpha
|
||||
# 每个专家的使用率(EMA)
|
||||
self.usage = torch.ones(num_experts) / num_experts
|
||||
# 每个专家的"领地"——它擅长的 token 分布
|
||||
self.expert_specialization: list[dict[int, float]] = [{} for _ in range(num_experts)]
|
||||
|
||||
def update(self, alpha: torch.Tensor, tokens: torch.Tensor):
|
||||
"""更新专家统计
|
||||
|
||||
Args:
|
||||
alpha: (batch, T, num_experts) 注意力权重
|
||||
tokens: (batch, T) 对应的 token
|
||||
"""
|
||||
# 使用率(时间维度平均)
|
||||
usage_batch = alpha.mean(dim=(0, 1)).detach().cpu() # (num_experts,)
|
||||
self.usage = self.ema_alpha * self.usage + (1 - self.ema_alpha) * usage_batch
|
||||
|
||||
# 专业化:每个专家最常处理哪些 token
|
||||
alpha_flat = alpha.reshape(-1, self.num_experts).detach().cpu() # (batch*T, num_experts)
|
||||
tokens_flat = tokens.reshape(-1).detach().cpu().tolist()
|
||||
for tok, weights in zip(tokens_flat, alpha_flat):
|
||||
for i, w in enumerate(weights.tolist()):
|
||||
if w > 0.1: # 只记录显著激活
|
||||
self.expert_specialization[i][tok] = \
|
||||
self.expert_specialization[i].get(tok, 0) + w
|
||||
|
||||
def get_diversity_bonus(self) -> torch.Tensor:
|
||||
"""返回多样性奖励——给使用率低的专家加权
|
||||
|
||||
在路由注意力上加上这个 bonus,鼓励新知识流向空闲专家
|
||||
"""
|
||||
# 使用率越低,bonus 越高
|
||||
bonus = (1.0 - self.usage) / self.usage.clamp(min=1e-4)
|
||||
bonus = bonus / bonus.sum() # 归一化
|
||||
return bonus
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""返回可解释性统计"""
|
||||
return {
|
||||
'usage': self.usage.tolist(),
|
||||
'top_specialization': [
|
||||
sorted(s.items(), key=lambda x: -x[1])[:5]
|
||||
for s in self.expert_specialization
|
||||
],
|
||||
}
|
||||
255
jspaceai/modules.py
Normal file
255
jspaceai/modules.py
Normal file
@@ -0,0 +1,255 @@
|
||||
"""
|
||||
外挂模块系统 —— 可热插拔的外部能力
|
||||
|
||||
设计:
|
||||
- 核心心智不依赖外挂,断开后继续工作
|
||||
- 标准接口,任何模块都能插入
|
||||
- 运行时热插拔,不需要重启
|
||||
- 心智知道外挂状态
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import torch, torch.nn as nn, numpy as np, time, math
|
||||
from typing import Optional, Any
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class ExternalModule(ABC):
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str: ...
|
||||
@abstractmethod
|
||||
def connect(self) -> bool: ...
|
||||
@abstractmethod
|
||||
def disconnect(self): ...
|
||||
@abstractmethod
|
||||
def is_connected(self) -> bool: ...
|
||||
@abstractmethod
|
||||
def query(self, input_data: Any) -> Any: ...
|
||||
@abstractmethod
|
||||
def describe(self) -> str: ...
|
||||
|
||||
|
||||
class SmallModelModule(ExternalModule):
|
||||
"""小模型外挂——提供语言/知识能力。可热插拔。"""
|
||||
|
||||
def __init__(self, workspace_dim=64, model_name="placeholder"):
|
||||
self._name = f"small_model:{model_name}"
|
||||
self.workspace_dim = workspace_dim
|
||||
self.model_name = model_name
|
||||
self._connected = False
|
||||
self._model = None
|
||||
self._tokenizer = None
|
||||
self._proj = None
|
||||
|
||||
@property
|
||||
def name(self): return self._name
|
||||
|
||||
def connect(self) -> bool:
|
||||
try:
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
print(f" [外挂] 加载 {self.model_name}...")
|
||||
self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
||||
self._model = AutoModel.from_pretrained(self.model_name)
|
||||
self._model.eval()
|
||||
hs = self._model.config.hidden_size
|
||||
self._proj = nn.Linear(hs, self.workspace_dim, bias=False)
|
||||
self._connected = True
|
||||
print(f" [外挂] {self.model_name} 已连接 (hidden={hs})")
|
||||
return True
|
||||
except ImportError:
|
||||
print(f" [外挂] transformers 未装,占位模式")
|
||||
self._connected = True
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f" [外挂] 连接失败: {e}")
|
||||
return False
|
||||
|
||||
def disconnect(self):
|
||||
if self._model is not None:
|
||||
del self._model, self._tokenizer
|
||||
self._model = self._tokenizer = None
|
||||
self._proj = None
|
||||
self._connected = False
|
||||
print(f" [外挂] {self.name} 已断开")
|
||||
|
||||
def is_connected(self): return self._connected
|
||||
|
||||
def query(self, input_data):
|
||||
if not self._connected:
|
||||
return None
|
||||
if self._model is None:
|
||||
n = input_data.shape[0] if isinstance(input_data, torch.Tensor) and input_data.dim() > 0 else 1
|
||||
return torch.randn(n, self.workspace_dim)
|
||||
try:
|
||||
text = str(input_data) if not isinstance(input_data, torch.Tensor) \
|
||||
else f"state_{input_data.mean().item():.3f}"
|
||||
with torch.no_grad():
|
||||
inputs = self._tokenizer(text, return_tensors="pt",
|
||||
truncation=True, max_length=128)
|
||||
outputs = self._model(**inputs)
|
||||
hidden = outputs.last_hidden_state.mean(dim=1)
|
||||
return self._proj(hidden)
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def describe(self):
|
||||
if self._model is None and self._connected:
|
||||
return "占位模式"
|
||||
return f"语言模型 {self.model_name}"
|
||||
|
||||
|
||||
class KnowledgeBaseModule(ExternalModule):
|
||||
"""知识库外挂——向量检索"""
|
||||
|
||||
def __init__(self, workspace_dim=64, capacity=10000):
|
||||
self._name = "knowledge_base"
|
||||
self.workspace_dim = workspace_dim
|
||||
self.capacity = capacity
|
||||
self._connected = False
|
||||
self.entries = []
|
||||
self.vectors = None
|
||||
|
||||
@property
|
||||
def name(self): return self._name
|
||||
|
||||
def connect(self):
|
||||
self._connected = True
|
||||
print(f" [外挂] 知识库已连接 ({len(self.entries)} 条)")
|
||||
return True
|
||||
|
||||
def disconnect(self):
|
||||
self._connected = False
|
||||
print(f" [外挂] 知识库已断开(数据保留)")
|
||||
|
||||
def is_connected(self): return self._connected
|
||||
|
||||
def add_entry(self, vector, text, metadata=None):
|
||||
if len(self.entries) >= self.capacity:
|
||||
self.entries.pop(0)
|
||||
self.entries.append({'vector': np.array(vector), 'text': text,
|
||||
'metadata': metadata or {}})
|
||||
self.vectors = np.array([e['vector'] for e in self.entries])
|
||||
|
||||
def query(self, input_data):
|
||||
if not self._connected or not self.entries:
|
||||
return []
|
||||
if not isinstance(input_data, torch.Tensor):
|
||||
return []
|
||||
q = input_data[0].cpu().numpy() if input_data.dim() > 1 else input_data.cpu().numpy()
|
||||
if self.vectors is None or len(self.vectors) == 0:
|
||||
return []
|
||||
norms = np.linalg.norm(self.vectors, axis=1) * np.linalg.norm(q)
|
||||
sims = self.vectors @ q / (norms + 1e-8)
|
||||
top = np.argsort(sims)[-5:][::-1]
|
||||
return [{'text': self.entries[i]['text'], 'sim': float(sims[i])} for i in top]
|
||||
|
||||
def describe(self):
|
||||
return f"知识库({len(self.entries)}/{self.capacity})"
|
||||
|
||||
|
||||
class ToolModule(ExternalModule):
|
||||
"""工具外挂——可调用的外部工具"""
|
||||
|
||||
def __init__(self):
|
||||
self._name = "tools"
|
||||
self._connected = False
|
||||
self.tools = {}
|
||||
|
||||
@property
|
||||
def name(self): return self._name
|
||||
|
||||
def connect(self):
|
||||
self._connected = True
|
||||
self.register('calc', lambda e: eval(e, {'__builtins__': {}}, {'math': math}), "计算")
|
||||
self.register('time', lambda: time.time(), "时间戳")
|
||||
print(f" [外挂] 工具箱已连接 ({len(self.tools)} 工具)")
|
||||
return True
|
||||
|
||||
def disconnect(self):
|
||||
self._connected = False
|
||||
self.tools.clear()
|
||||
print(f" [外挂] 工具箱已断开")
|
||||
|
||||
def is_connected(self): return self._connected
|
||||
|
||||
def register(self, name, func, desc=""):
|
||||
self.tools[name] = {'func': func, 'desc': desc}
|
||||
|
||||
def query(self, input_data):
|
||||
if not self._connected:
|
||||
return None
|
||||
if isinstance(input_data, dict) and 'tool' in input_data:
|
||||
tn = input_data['tool']
|
||||
args = input_data.get('args', [])
|
||||
if tn in self.tools:
|
||||
try:
|
||||
return self.tools[tn]['func'](*args)
|
||||
except Exception as e:
|
||||
return f"错误: {e}"
|
||||
return None
|
||||
|
||||
def describe(self):
|
||||
return f"工具箱({list(self.tools.keys())})"
|
||||
|
||||
|
||||
class ModuleDock:
|
||||
"""外挂坞——USB hub 式管理可热插拔模块。
|
||||
|
||||
用法:
|
||||
dock = ModuleDock()
|
||||
dock.register('llm', SmallModelModule())
|
||||
dock.connect('llm') # 热插上
|
||||
result = dock.query('llm', "hello")
|
||||
dock.disconnect('llm') # 拔掉,心智不停
|
||||
"""
|
||||
|
||||
def __init__(self, workspace_dim=64):
|
||||
self.workspace_dim = workspace_dim
|
||||
self.slots = {}
|
||||
self.log = []
|
||||
|
||||
def register(self, name, module):
|
||||
self.slots[name] = module
|
||||
|
||||
def connect(self, name) -> bool:
|
||||
if name not in self.slots:
|
||||
return False
|
||||
if self.slots[name].is_connected():
|
||||
return True
|
||||
ok = self.slots[name].connect()
|
||||
self.log.append({'time': time.time(), 'slot': name,
|
||||
'action': 'connect', 'ok': ok})
|
||||
return ok
|
||||
|
||||
def disconnect(self, name):
|
||||
if name not in self.slots:
|
||||
return
|
||||
self.slots[name].disconnect()
|
||||
self.log.append({'time': time.time(), 'slot': name,
|
||||
'action': 'disconnect', 'ok': True})
|
||||
|
||||
def disconnect_all(self):
|
||||
for n in list(self.slots):
|
||||
if self.slots[n].is_connected():
|
||||
self.disconnect(n)
|
||||
|
||||
def is_connected(self, name) -> bool:
|
||||
m = self.slots.get(name)
|
||||
return m.is_connected() if m else False
|
||||
|
||||
def query(self, name, input_data):
|
||||
"""查询某个外挂。未连接返回 None。"""
|
||||
m = self.slots.get(name)
|
||||
if m and m.is_connected():
|
||||
return m.query(input_data)
|
||||
return None
|
||||
|
||||
def status(self) -> dict:
|
||||
"""所有外挂状态"""
|
||||
return {name: {
|
||||
'connected': m.is_connected(),
|
||||
'description': m.describe(),
|
||||
} for name, m in self.slots.items()}
|
||||
|
||||
def connected_count(self) -> int:
|
||||
return sum(1 for m in self.slots.values() if m.is_connected())
|
||||
464
jspaceai/multimodal.py
Normal file
464
jspaceai/multimodal.py
Normal file
@@ -0,0 +1,464 @@
|
||||
"""
|
||||
多模态感知-行动系统
|
||||
|
||||
原生支持图片、音频、视频、文本四种模态。
|
||||
所有模态编码到统一的 workspace 向量空间,从 workspace 解码到任意模态。
|
||||
|
||||
设计原则:
|
||||
1. 原生多模态——不经过文本中间表示
|
||||
2. 流式处理——支持实时麦克风/摄像头输入
|
||||
3. 闭环——输出反馈到输入
|
||||
4. 在线学习——持续进化
|
||||
|
||||
模态专家分工(对应 Anthropic 论文的"专家分工"):
|
||||
- 视觉专家:处理图片/视频帧
|
||||
- 听觉专家:处理音频频谱
|
||||
- 语言专家:处理文本 token
|
||||
- 跨模态专家:对齐不同模态的概念
|
||||
|
||||
每个专家有自己的内部状态 m_i,通过 J-space 在 workspace w 中广播。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
import numpy as np
|
||||
|
||||
from .core import JSpaceConfig, Expert, JSpaceWorkspace
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultimodalConfig(JSpaceConfig):
|
||||
"""多模态配置"""
|
||||
# 文本
|
||||
vocab_size: int = 100
|
||||
embed_dim: int = 16
|
||||
# 视觉
|
||||
img_channels: int = 3
|
||||
img_size: int = 32 # 编码后的图像尺寸(原始会 resize)
|
||||
visual_feature_dim: int = 32
|
||||
# 音频
|
||||
audio_sample_rate: int = 16000
|
||||
audio_frame_size: int = 1024 # 每帧采样数
|
||||
audio_feature_dim: int = 32
|
||||
# 共享
|
||||
workspace_dim: int = 64 # 比 unimodal 大,承载多模态
|
||||
num_experts: int = 12 # 12 个专家:2视觉 + 2屏幕 + 2听觉 + 2语言 + 2鼠标 + 2跨模态
|
||||
expert_dim: int = 24
|
||||
# 键盘/鼠标
|
||||
keyboard_vocab: int = 128 # 键盘字符词汇表
|
||||
mouse_feature_dim: int = 16 # 鼠标特征维度
|
||||
|
||||
|
||||
class VisualEncoder(nn.Module):
|
||||
"""
|
||||
视觉编码器:图片 → workspace 向量
|
||||
|
||||
轻量 CNN,不用预训练。从零学。
|
||||
输入:(batch, 3, H, W)
|
||||
输出:(batch, input_dim) 投影到 workspace 输入空间
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(3, 16, 3, stride=2, padding=1), # 32→16
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(16, 32, 3, stride=2, padding=1), # 16→8
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(32, 32, 3, stride=2, padding=1), # 8→4
|
||||
nn.ReLU(),
|
||||
nn.AdaptiveAvgPool2d(1), # 全局池化 → (32, 1, 1)
|
||||
)
|
||||
self.proj = nn.Linear(32, input_dim)
|
||||
|
||||
def forward(self, img: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
img: (batch, 3, H, W) 或 (batch, T, 3, H, W) 视频序列
|
||||
returns: (batch, input_dim) 或 (batch, T, input_dim)
|
||||
"""
|
||||
if img.dim() == 5: # 视频序列
|
||||
B, T, C, H, W = img.shape
|
||||
img = img.reshape(B * T, C, H, W)
|
||||
feat = self.conv(img).squeeze(-1).squeeze(-1) # (B*T, 32)
|
||||
feat = self.proj(feat) # (B*T, input_dim)
|
||||
return feat.reshape(B, T, -1)
|
||||
else:
|
||||
feat = self.conv(img).squeeze(-1).squeeze(-1)
|
||||
return self.proj(feat)
|
||||
|
||||
|
||||
class AudioEncoder(nn.Module):
|
||||
"""
|
||||
音频编码器:音频波形 → workspace 向量
|
||||
|
||||
轻量 1D CNN 处理原始波形。
|
||||
输入:(batch, audio_frame_size) 原始音频采样
|
||||
输出:(batch, input_dim)
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim: int, audio_frame_size: int = 1024):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv1d(1, 16, 64, stride=4, padding=32), # 下采样
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(16, 32, 32, stride=4, padding=16),
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(32, 32, 16, stride=2, padding=8),
|
||||
nn.ReLU(),
|
||||
nn.AdaptiveAvgPool1d(1),
|
||||
)
|
||||
self.proj = nn.Linear(32, input_dim)
|
||||
|
||||
def forward(self, audio: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
audio: (batch, frame_size) 或 (batch, T, frame_size)
|
||||
returns: (batch, input_dim) 或 (batch, T, input_dim)
|
||||
"""
|
||||
if audio.dim() == 3:
|
||||
B, T, F = audio.shape
|
||||
audio = audio.reshape(B * T, F)
|
||||
x = audio.unsqueeze(1) # (B*T, 1, F)
|
||||
feat = self.conv(x).squeeze(-1)
|
||||
feat = self.proj(feat)
|
||||
return feat.reshape(B, T, -1)
|
||||
else:
|
||||
x = audio.unsqueeze(1) # (B, 1, F)
|
||||
feat = self.conv(x).squeeze(-1)
|
||||
return self.proj(feat)
|
||||
|
||||
|
||||
class TextEncoder(nn.Module):
|
||||
"""文本编码器:token → workspace 向量"""
|
||||
|
||||
def __init__(self, vocab_size: int, embed_dim: int, input_dim: int):
|
||||
super().__init__()
|
||||
self.embedding = nn.Embedding(vocab_size, embed_dim)
|
||||
self.proj = nn.Linear(embed_dim, input_dim, bias=False)
|
||||
|
||||
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""token_ids: (batch,) 或 (batch, T)"""
|
||||
if token_ids.dim() == 1:
|
||||
emb = self.embedding(token_ids)
|
||||
return self.proj(emb)
|
||||
else:
|
||||
emb = self.embedding(token_ids) # (B, T, embed)
|
||||
return self.proj(emb)
|
||||
|
||||
|
||||
class KeyboardEncoder(nn.Module):
|
||||
"""
|
||||
键盘编码器:按键序列 → workspace 向量
|
||||
|
||||
把键盘按键序列编码成语义向量。
|
||||
和 TextEncoder 类似,但字符集不同(含特殊键)。
|
||||
"""
|
||||
|
||||
def __init__(self, vocab_size: int, embed_dim: int, input_dim: int):
|
||||
super().__init__()
|
||||
self.embedding = nn.Embedding(vocab_size, embed_dim)
|
||||
self.proj = nn.Linear(embed_dim, input_dim, bias=False)
|
||||
|
||||
def forward(self, key_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""key_ids: (batch,) 或 (batch, T)"""
|
||||
if key_ids.dim() == 1:
|
||||
return self.proj(self.embedding(key_ids))
|
||||
else:
|
||||
return self.proj(self.embedding(key_ids))
|
||||
|
||||
|
||||
class MouseEncoder(nn.Module):
|
||||
"""
|
||||
鼠标编码器:鼠标位置 + 事件 → workspace 向量
|
||||
|
||||
输入:(batch, 4) = (x_norm, y_norm, click_left, click_right)
|
||||
x_norm, y_norm: 归一化到 [0, 1] 的鼠标坐标
|
||||
click_left, click_right: 0 或 1
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim: int):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(4, 32),
|
||||
nn.ReLU(),
|
||||
nn.Linear(32, input_dim),
|
||||
)
|
||||
|
||||
def forward(self, mouse_data: torch.Tensor) -> torch.Tensor:
|
||||
"""mouse_data: (batch, 4) → (batch, input_dim)"""
|
||||
return self.net(mouse_data)
|
||||
|
||||
|
||||
class VisualDecoder(nn.Module):
|
||||
"""视觉解码器:workspace 向量 → 图片"""
|
||||
|
||||
def __init__(self, workspace_dim: int, img_size: int = 32):
|
||||
super().__init__()
|
||||
self.img_size = img_size
|
||||
self.fc = nn.Linear(workspace_dim, 32 * 4 * 4)
|
||||
self.deconv = nn.Sequential(
|
||||
nn.ConvTranspose2d(32, 32, 3, stride=2, padding=1, output_padding=1), # 4→8
|
||||
nn.ReLU(),
|
||||
nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1), # 8→16
|
||||
nn.ReLU(),
|
||||
nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1), # 16→32
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, w: torch.Tensor) -> torch.Tensor:
|
||||
"""w: (batch, workspace_dim) → (batch, 3, 32, 32)"""
|
||||
x = self.fc(w).reshape(-1, 32, 4, 4)
|
||||
return self.deconv(x)
|
||||
|
||||
|
||||
class AudioDecoder(nn.Module):
|
||||
"""音频解码器:workspace 向量 → 音频波形"""
|
||||
|
||||
def __init__(self, workspace_dim: int, audio_frame_size: int = 1024):
|
||||
super().__init__()
|
||||
self.audio_frame_size = audio_frame_size
|
||||
self.fc = nn.Linear(workspace_dim, 32 * 64)
|
||||
self.deconv = nn.Sequential(
|
||||
nn.ConvTranspose1d(32, 32, 32, stride=4, padding=14, output_padding=2),
|
||||
nn.ReLU(),
|
||||
nn.ConvTranspose1d(32, 16, 16, stride=4, padding=6, output_padding=2),
|
||||
nn.ReLU(),
|
||||
nn.ConvTranspose1d(16, 1, 64, stride=4, padding=30, output_padding=2),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, w: torch.Tensor) -> torch.Tensor:
|
||||
"""w: (batch, workspace_dim) → (batch, audio_frame_size)"""
|
||||
x = self.fc(w).reshape(-1, 32, 64)
|
||||
audio = self.deconv(x).squeeze(1) # (batch, L)
|
||||
# 裁剪或填充到目标长度
|
||||
if audio.shape[-1] > self.audio_frame_size:
|
||||
audio = audio[:, :self.audio_frame_size]
|
||||
else:
|
||||
audio = F.pad(audio, (0, self.audio_frame_size - audio.shape[-1]))
|
||||
return audio
|
||||
|
||||
|
||||
class TextDecoder(nn.Module):
|
||||
"""文本解码器:workspace 向量 → logits"""
|
||||
|
||||
def __init__(self, workspace_dim: int, vocab_size: int):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(workspace_dim, 64),
|
||||
nn.Tanh(),
|
||||
nn.Linear(64, vocab_size),
|
||||
)
|
||||
|
||||
def forward(self, w: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(w)
|
||||
|
||||
|
||||
class MultimodalJSpaceModel(nn.Module):
|
||||
"""
|
||||
多模态 JSpace 模型。
|
||||
|
||||
12 个专家分工:
|
||||
专家 0,1: 视觉(处理摄像头图片)
|
||||
专家 2,3: 屏幕(处理屏幕截图)
|
||||
专家 4,5: 听觉(处理音频帧)
|
||||
专家 6,7: 语言(处理文本 token)
|
||||
专家 8,9: 鼠标(处理鼠标位置/点击)
|
||||
专家 10,11: 跨模态(对齐不同模态的概念)
|
||||
|
||||
所有专家共享 workspace w,通过 J-space 广播。
|
||||
|
||||
支持的输入模态:
|
||||
'image' - 摄像头图片 (batch, 3, H, W)
|
||||
'screen' - 屏幕截图 (batch, 3, H, W)(用视觉编码器)
|
||||
'audio' - 音频帧 (batch, frame_size)
|
||||
'text' - 文本 token (batch,)
|
||||
'keyboard' - 键盘按键 (batch,) 或 (batch, T)
|
||||
'mouse' - 鼠标数据 (batch, 4) = (x, y, click_l, click_r)
|
||||
"""
|
||||
|
||||
def __init__(self, config: MultimodalConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
# 编码器(各模态 → input_dim)
|
||||
self.visual_encoder = VisualEncoder(config.input_dim) # 摄像头 + 屏幕
|
||||
self.audio_encoder = AudioEncoder(config.input_dim, config.audio_frame_size)
|
||||
self.text_encoder = TextEncoder(config.vocab_size, config.embed_dim, config.input_dim)
|
||||
self.keyboard_encoder = KeyboardEncoder(config.keyboard_vocab, config.embed_dim, config.input_dim)
|
||||
self.mouse_encoder = MouseEncoder(config.input_dim)
|
||||
|
||||
# 解码器(workspace → 各模态)
|
||||
self.visual_decoder = VisualDecoder(config.workspace_dim, config.img_size)
|
||||
self.audio_decoder = AudioDecoder(config.workspace_dim, config.audio_frame_size)
|
||||
self.text_decoder = TextDecoder(config.workspace_dim, config.vocab_size)
|
||||
|
||||
# 专家池
|
||||
self.experts = nn.ModuleList([
|
||||
Expert(
|
||||
expert_dim=config.expert_dim,
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_wells=config.num_wells,
|
||||
sparsity=config.jacobian_sparsity,
|
||||
)
|
||||
for _ in range(config.num_experts)
|
||||
])
|
||||
|
||||
# 工作空间
|
||||
self.workspace = JSpaceWorkspace(
|
||||
workspace_dim=config.workspace_dim,
|
||||
input_dim=config.input_dim,
|
||||
num_experts=config.num_experts,
|
||||
)
|
||||
|
||||
# 模态类型标记(12 个专家的分工)
|
||||
self.expert_modality = (
|
||||
['visual'] * 2 + # 0,1: 摄像头
|
||||
['screen'] * 2 + # 2,3: 屏幕
|
||||
['audio'] * 2 + # 4,5: 听觉
|
||||
['text'] * 2 + # 6,7: 语言
|
||||
['mouse'] * 2 + # 8,9: 鼠标
|
||||
['cross'] * 2 # 10,11: 跨模态
|
||||
)[:config.num_experts]
|
||||
|
||||
def init_state(self, batch_size: int, device: torch.device) -> dict:
|
||||
return {
|
||||
'w': torch.zeros(batch_size, self.config.workspace_dim, device=device),
|
||||
'm': [torch.zeros(batch_size, self.config.expert_dim, device=device)
|
||||
for _ in range(self.config.num_experts)],
|
||||
}
|
||||
|
||||
def step(self, state: dict, x: torch.Tensor,
|
||||
record_trajectory: bool = False) -> tuple[dict, list]:
|
||||
"""
|
||||
单步前向
|
||||
|
||||
Args:
|
||||
state: {'w': ..., 'm': [...]}
|
||||
x: (batch, input_dim) 已编码的输入(任意模态)
|
||||
record_trajectory: 是否记录 w 轨迹
|
||||
|
||||
Returns:
|
||||
new_state, w_trajectory (list)
|
||||
"""
|
||||
w = state['w']
|
||||
ms = state['m']
|
||||
cfg = self.config
|
||||
w_trajectory = []
|
||||
|
||||
for substep in range(cfg.ode_steps):
|
||||
contributions = []
|
||||
new_ms = []
|
||||
for i, expert in enumerate(self.experts):
|
||||
m_next, contrib = expert(
|
||||
ms[i], w, x,
|
||||
dt=cfg.dt, noise_std=cfg.noise_std,
|
||||
)
|
||||
new_ms.append(m_next)
|
||||
contributions.append(contrib)
|
||||
contributions = torch.stack(contributions, dim=1)
|
||||
|
||||
w, alpha = self.workspace(
|
||||
w, x, contributions,
|
||||
dt=cfg.dt, tau_w=cfg.tau_w,
|
||||
)
|
||||
ms = new_ms
|
||||
|
||||
if record_trajectory:
|
||||
w_trajectory.append(w.detach())
|
||||
|
||||
new_state = {'w': w, 'm': ms}
|
||||
return new_state, w_trajectory
|
||||
|
||||
def encode_modality(self, modality: str, data: torch.Tensor) -> torch.Tensor:
|
||||
"""编码任意模态到 input_dim
|
||||
|
||||
支持的模态:image, screen, audio, text, keyboard, mouse
|
||||
- image/screen: 共用 visual_encoder(都是 RGB 图像)
|
||||
- keyboard: 键盘按键 id
|
||||
- mouse: (batch, 4) = (x_norm, y_norm, click_left, click_right)
|
||||
"""
|
||||
if modality in ('image', 'screen'):
|
||||
return self.visual_encoder(data)
|
||||
elif modality == 'audio':
|
||||
return self.audio_encoder(data)
|
||||
elif modality == 'text':
|
||||
return self.text_encoder(data)
|
||||
elif modality == 'keyboard':
|
||||
return self.keyboard_encoder(data)
|
||||
elif modality == 'mouse':
|
||||
return self.mouse_encoder(data)
|
||||
else:
|
||||
raise ValueError(f"Unknown modality: {modality}")
|
||||
|
||||
def decode_modality(self, modality: str, w: torch.Tensor) -> torch.Tensor:
|
||||
"""从 workspace 解码到任意模态
|
||||
|
||||
screen 用 visual_decoder(和 image 共享)
|
||||
keyboard 用 text_decoder(都是 token)
|
||||
"""
|
||||
if modality in ('image', 'screen'):
|
||||
return self.visual_decoder(w)
|
||||
elif modality == 'audio':
|
||||
return self.audio_decoder(w)
|
||||
elif modality in ('text', 'keyboard'):
|
||||
return self.text_decoder(w)
|
||||
else:
|
||||
raise ValueError(f"Unknown modality: {modality}")
|
||||
|
||||
def forward_multimodal(self, modality: str, data: torch.Tensor,
|
||||
state: dict | None = None,
|
||||
record_trajectory: bool = False) -> tuple[dict, dict]:
|
||||
"""
|
||||
多模态前向
|
||||
|
||||
Args:
|
||||
modality: 'image' / 'audio' / 'text'
|
||||
data: 模态原始数据
|
||||
state: 初始状态
|
||||
record_trajectory: 是否记录 w 轨迹
|
||||
|
||||
Returns:
|
||||
outputs: {'w': workspace, 'logits/img/audio': 各模态解码}
|
||||
info: {'alpha', 'w_norm', 'w_trajectory'}
|
||||
"""
|
||||
if state is None:
|
||||
# 先编码确定 batch size
|
||||
x = self.encode_modality(modality, data)
|
||||
batch_size = x.shape[0]
|
||||
device = x.device
|
||||
state = self.init_state(batch_size, device)
|
||||
else:
|
||||
x = self.encode_modality(modality, data)
|
||||
|
||||
# 处理序列或单步
|
||||
if x.dim() == 2: # 单步 (batch, input_dim)
|
||||
state, w_traj = self.step(state, x, record_trajectory)
|
||||
w = state['w']
|
||||
else: # 序列 (batch, T, input_dim)
|
||||
w_traj_all = []
|
||||
for t in range(x.shape[1]):
|
||||
state, w_traj = self.step(state, x[:, t], record_trajectory)
|
||||
if record_trajectory:
|
||||
w_traj_all.append(w_traj)
|
||||
w = state['w']
|
||||
|
||||
# 解码到所有模态(workspace 是模态无关的,可解码到任意模态)
|
||||
outputs = {
|
||||
'w': w,
|
||||
'image': self.visual_decoder(w),
|
||||
'audio': self.audio_decoder(w),
|
||||
'text_logits': self.text_decoder(w),
|
||||
}
|
||||
|
||||
info = {
|
||||
'w_norm': w.norm(dim=-1),
|
||||
}
|
||||
if record_trajectory and w_traj_all:
|
||||
info['w_trajectory'] = w_traj_all
|
||||
|
||||
return outputs, info
|
||||
128
jspaceai/platform.py
Normal file
128
jspaceai/platform.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
平台抽象层——跨平台支持(macOS / Windows / Linux)
|
||||
|
||||
统一封装各平台差异:屏幕尺寸、权限检查、路径等。
|
||||
所有平台相关代码集中在这里。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import platform
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlatformInfo:
|
||||
"""平台信息"""
|
||||
system: str
|
||||
machine: str
|
||||
python_version: str
|
||||
is_macos: bool
|
||||
is_windows: bool
|
||||
is_linux: bool
|
||||
|
||||
@classmethod
|
||||
def detect(cls) -> "PlatformInfo":
|
||||
s = platform.system()
|
||||
return cls(
|
||||
system=s,
|
||||
machine=platform.machine(),
|
||||
python_version=sys.version,
|
||||
is_macos=(s == 'Darwin'),
|
||||
is_windows=(s == 'Windows'),
|
||||
is_linux=(s == 'Linux'),
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.system}/{self.machine} (Python {self.python_version.split()[0]})"
|
||||
|
||||
|
||||
PLATFORM = PlatformInfo.detect()
|
||||
|
||||
|
||||
def get_screen_size() -> tuple[int, int]:
|
||||
"""获取主屏幕尺寸(跨平台)"""
|
||||
try:
|
||||
import mss
|
||||
with mss.MSS() as sct:
|
||||
mon = sct.monitors[1]
|
||||
return (mon['width'], mon['height'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if PLATFORM.is_windows:
|
||||
try:
|
||||
import ctypes
|
||||
user32 = ctypes.windll.user32
|
||||
return (user32.GetSystemMetrics(0), user32.GetSystemMetrics(1))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return (1920, 1080)
|
||||
|
||||
|
||||
def check_camera_permission() -> bool:
|
||||
"""检查摄像头权限"""
|
||||
if not PLATFORM.is_macos:
|
||||
return True
|
||||
try:
|
||||
import cv2
|
||||
cap = cv2.VideoCapture(0)
|
||||
ok = cap.isOpened()
|
||||
cap.release()
|
||||
return ok
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def check_microphone_permission() -> bool:
|
||||
"""检查麦克风权限"""
|
||||
try:
|
||||
import sounddevice as sd
|
||||
sd.query_devices(kind='input')
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def check_input_monitoring_permission() -> bool:
|
||||
"""检查键盘/鼠标监听权限"""
|
||||
if PLATFORM.is_windows:
|
||||
return True
|
||||
if PLATFORM.is_linux:
|
||||
return True
|
||||
if PLATFORM.is_macos:
|
||||
try:
|
||||
import subprocess
|
||||
result = subprocess.run(
|
||||
['osascript', '-e',
|
||||
'tell application "System Events" to keystroke ""'],
|
||||
capture_output=True, timeout=2,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def print_permission_guide():
|
||||
"""打印权限配置指南"""
|
||||
print("\n权限配置指南:")
|
||||
print("-" * 50)
|
||||
if PLATFORM.is_macos:
|
||||
print("macOS 权限设置:")
|
||||
print(" 1. 摄像头: 系统设置 → 隐私与安全 → 摄像头")
|
||||
print(" 2. 麦克风: 系统设置 → 隐私与安全 → 麦克风")
|
||||
print(" 3. 键盘/鼠标监听: 系统设置 → 隐私与安全 → 辅助功能")
|
||||
print(" 4. 屏幕录制: 系统设置 → 隐私与安全 → 屏幕录制")
|
||||
elif PLATFORM.is_linux:
|
||||
print("Linux 权限设置:")
|
||||
print(" 1. 摄像头: 确保用户在 video 组 (sudo usermod -aG video $USER)")
|
||||
print(" 2. 音频: 确保用户在 audio 组")
|
||||
print(" 3. 键盘/鼠标: 需要 X11 或 Wayland 输入权限")
|
||||
print(" 4. 屏幕捕获: 需要 X11 或安装 grim/slurp (Wayland)")
|
||||
elif PLATFORM.is_windows:
|
||||
print("Windows 权限设置:")
|
||||
print(" 1. 摄像头/麦克风: 设置 → 隐私 → 摄像头/麦克风")
|
||||
print(" 2. 键盘/鼠标: 通常不需要额外权限")
|
||||
print("-" * 50)
|
||||
374
jspaceai/realtime.py
Normal file
374
jspaceai/realtime.py
Normal file
@@ -0,0 +1,374 @@
|
||||
"""
|
||||
实时多模态 I/O 层
|
||||
|
||||
使用 OpenCV(视频/图像)+ sounddevice(音频 I/O,基于 PortAudio 开源库)
|
||||
+ macOS afplay(音频播放)实现实时感知-行动循环。
|
||||
|
||||
提供:
|
||||
- CameraStream: 实时摄像头采集
|
||||
- MicrophoneStream: 实时麦克风采集
|
||||
- AudioPlayer: 音频播放(扬声器)
|
||||
- MultimodalStream: 统一的实时多模态流管理
|
||||
|
||||
所有流都是非阻塞的,用队列缓冲。可以同时采集视频和音频。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import sounddevice as sd
|
||||
import torch
|
||||
import threading
|
||||
import queue
|
||||
import time
|
||||
from typing import Optional, Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Frame:
|
||||
"""统一的帧数据结构"""
|
||||
timestamp: float
|
||||
modality: str # 'image' / 'audio' / 'text'
|
||||
data: np.ndarray # 图像 (H,W,3) / 音频 (frame_size,) / 文本 str
|
||||
|
||||
|
||||
class CameraStream:
|
||||
"""
|
||||
OpenCV 摄像头流。
|
||||
|
||||
非阻塞采集,帧放入队列。
|
||||
"""
|
||||
|
||||
def __init__(self, camera_index: int = 0, frame_size: tuple = (32, 32),
|
||||
fps: int = 10):
|
||||
self.camera_index = camera_index
|
||||
self.target_size = frame_size # resize 到这个尺寸(给模型用)
|
||||
self.fps = fps
|
||||
self.frame_queue: queue.Queue = queue.Queue(maxsize=30)
|
||||
self.running = False
|
||||
self.thread: Optional[threading.Thread] = None
|
||||
self.cap: Optional[cv2.VideoCapture] = None
|
||||
|
||||
def start(self):
|
||||
"""启动采集"""
|
||||
self.cap = cv2.VideoCapture(self.camera_index)
|
||||
if not self.cap.isOpened():
|
||||
raise RuntimeError(f"无法打开摄像头 {self.camera_index}(需要 macOS 权限)")
|
||||
|
||||
self.running = True
|
||||
self.thread = threading.Thread(target=self._capture_loop, daemon=True)
|
||||
self.thread.start()
|
||||
|
||||
def _capture_loop(self):
|
||||
frame_interval = 1.0 / self.fps
|
||||
while self.running:
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
continue
|
||||
|
||||
# BGR → RGB
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
# resize 到目标尺寸
|
||||
frame_resized = cv2.resize(frame_rgb, self.target_size)
|
||||
|
||||
try:
|
||||
self.frame_queue.put_nowait(
|
||||
Frame(time.time(), 'image', frame_resized)
|
||||
)
|
||||
except queue.Full:
|
||||
pass # 丢帧
|
||||
|
||||
time.sleep(frame_interval)
|
||||
|
||||
def get_frame(self) -> Optional[Frame]:
|
||||
"""获取最新帧(非阻塞)"""
|
||||
try:
|
||||
return self.frame_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
return None
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
if self.thread:
|
||||
self.thread.join(timeout=1.0)
|
||||
if self.cap:
|
||||
self.cap.release()
|
||||
|
||||
|
||||
class MicrophoneStream:
|
||||
"""
|
||||
sounddevice 麦克风流。
|
||||
|
||||
持续采集音频,分帧放入队列。
|
||||
"""
|
||||
|
||||
def __init__(self, sample_rate: int = 16000, frame_size: int = 1024,
|
||||
channels: int = 1):
|
||||
self.sample_rate = sample_rate
|
||||
self.frame_size = frame_size
|
||||
self.channels = channels
|
||||
self.frame_queue: queue.Queue = queue.Queue(maxsize=100)
|
||||
self.running = False
|
||||
self.stream: Optional[sd.InputStream] = None
|
||||
|
||||
def start(self):
|
||||
"""启动采集"""
|
||||
self.running = True
|
||||
self.stream = sd.InputStream(
|
||||
samplerate=self.sample_rate,
|
||||
blocksize=self.frame_size,
|
||||
channels=self.channels,
|
||||
dtype='float32',
|
||||
callback=self._audio_callback,
|
||||
)
|
||||
self.stream.start()
|
||||
|
||||
def _audio_callback(self, indata, frames, time_info, status):
|
||||
if not self.running:
|
||||
return
|
||||
# indata: (frame_size, channels)
|
||||
audio = indata[:, 0] if self.channels > 1 else indata.flatten()
|
||||
try:
|
||||
self.frame_queue.put_nowait(
|
||||
Frame(time.time(), 'audio', audio.copy())
|
||||
)
|
||||
except queue.Full:
|
||||
pass
|
||||
|
||||
def get_frame(self) -> Optional[Frame]:
|
||||
try:
|
||||
return self.frame_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
return None
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
if self.stream:
|
||||
self.stream.stop()
|
||||
self.stream.close()
|
||||
|
||||
|
||||
class AudioPlayer:
|
||||
"""
|
||||
音频播放器。用 sounddevice 输出到扬声器。
|
||||
"""
|
||||
|
||||
def __init__(self, sample_rate: int = 16000):
|
||||
self.sample_rate = sample_rate
|
||||
self.playing = False
|
||||
|
||||
def play(self, audio: np.ndarray, blocking: bool = False):
|
||||
"""播放音频波形
|
||||
|
||||
Args:
|
||||
audio: (n_samples,) float32, 范围 [-1, 1]
|
||||
blocking: 是否阻塞等待播放完成
|
||||
"""
|
||||
audio = np.clip(audio, -1, 1).astype(np.float32)
|
||||
sd.play(audio, self.sample_rate)
|
||||
if blocking:
|
||||
sd.wait()
|
||||
|
||||
def stop(self):
|
||||
sd.stop()
|
||||
|
||||
|
||||
class MultimodalStream:
|
||||
"""
|
||||
统一管理摄像头 + 麦克风 + 扬声器的实时流。
|
||||
|
||||
提供统一的接口:
|
||||
- start(): 启动所有采集
|
||||
- get_frames(): 获取当前所有可用帧(按时间戳对齐)
|
||||
- respond_audio(): 播放音频回应
|
||||
- stop(): 停止所有
|
||||
|
||||
这是感知-行动循环的物理层。
|
||||
"""
|
||||
|
||||
def __init__(self, use_camera: bool = True, use_mic: bool = True,
|
||||
img_size: tuple = (32, 32), sample_rate: int = 16000,
|
||||
audio_frame_size: int = 1024):
|
||||
self.use_camera = use_camera
|
||||
self.use_mic = use_mic
|
||||
|
||||
self.camera = CameraStream(frame_size=img_size) if use_camera else None
|
||||
self.mic = MicrophoneStream(
|
||||
sample_rate=sample_rate, frame_size=audio_frame_size
|
||||
) if use_mic else None
|
||||
self.player = AudioPlayer(sample_rate=sample_rate)
|
||||
|
||||
def start(self):
|
||||
"""启动所有采集"""
|
||||
if self.camera:
|
||||
try:
|
||||
self.camera.start()
|
||||
print(" 摄像头已启动")
|
||||
except RuntimeError as e:
|
||||
print(f" 摄像头启动失败: {e}")
|
||||
self.camera = None
|
||||
|
||||
if self.mic:
|
||||
try:
|
||||
self.mic.start()
|
||||
print(" 麦克风已启动")
|
||||
except Exception as e:
|
||||
print(f" 麦克风启动失败: {e}")
|
||||
self.mic = None
|
||||
|
||||
def get_latest_frames(self) -> dict[str, Optional[Frame]]:
|
||||
"""获取最新的各模态帧"""
|
||||
result = {
|
||||
'image': None,
|
||||
'audio': None,
|
||||
}
|
||||
if self.camera:
|
||||
# 清空队列,只保留最新
|
||||
latest = None
|
||||
while True:
|
||||
f = self.camera.get_frame()
|
||||
if f is None:
|
||||
break
|
||||
latest = f
|
||||
result['image'] = latest
|
||||
|
||||
if self.mic:
|
||||
latest = None
|
||||
while True:
|
||||
f = self.mic.get_frame()
|
||||
if f is None:
|
||||
break
|
||||
latest = f
|
||||
result['audio'] = latest
|
||||
|
||||
return result
|
||||
|
||||
def play_audio(self, audio: np.ndarray, blocking: bool = False):
|
||||
"""播放音频"""
|
||||
self.player.play(audio, blocking)
|
||||
|
||||
def stop(self):
|
||||
if self.camera:
|
||||
self.camera.stop()
|
||||
if self.mic:
|
||||
self.mic.stop()
|
||||
self.player.stop()
|
||||
|
||||
|
||||
class SensoryMotorLoop:
|
||||
"""
|
||||
感知-行动循环:传感器 → 模型 → 执行器 → 传感器...
|
||||
|
||||
这是自进化语言网络的核心运行时。
|
||||
|
||||
循环:
|
||||
1. 从传感器读帧(摄像头/麦克风)
|
||||
2. 编码到 workspace 输入
|
||||
3. 模型 forward 一步(workspace 更新)
|
||||
4. 从 workspace 解码输出(音频/图像/文本)
|
||||
5. 执行器输出(扬声器播放/显示图像)
|
||||
6. 输出反馈到输入(自监督:预测自己的输出)
|
||||
7. 在线学习(更新参数)
|
||||
|
||||
这个循环永不停止——模型持续感知、思考、行动、学习。
|
||||
"""
|
||||
|
||||
def __init__(self, model, stream: MultimodalStream, device: str = 'cpu'):
|
||||
self.model = model
|
||||
self.stream = stream
|
||||
self.device = device
|
||||
self.state = model.init_state(1, torch.device(device))
|
||||
self.running = False
|
||||
self.step_count = 0
|
||||
|
||||
# 自监督学习记录
|
||||
self.last_output = None # 上一步的输出(用于反馈)
|
||||
self.online_loss_history = []
|
||||
|
||||
def step_once(self) -> dict:
|
||||
"""执行一步感知-行动循环
|
||||
|
||||
Returns:
|
||||
info: 包含输入模态、输出、loss 等
|
||||
"""
|
||||
# 1. 读取传感器
|
||||
frames = self.stream.get_latest_frames()
|
||||
|
||||
# 选择活跃模态(优先音频,其次图像)
|
||||
modality = None
|
||||
input_data = None
|
||||
if frames['audio'] is not None:
|
||||
modality = 'audio'
|
||||
input_data = torch.tensor(
|
||||
frames['audio'].data, dtype=torch.float32
|
||||
).unsqueeze(0).to(self.device)
|
||||
elif frames['image'] is not None:
|
||||
modality = 'image'
|
||||
img = frames['image'].data # (H, W, 3) uint8
|
||||
img_tensor = torch.tensor(
|
||||
img, dtype=torch.float32
|
||||
).permute(2, 0, 1).unsqueeze(0).to(self.device) / 127.5 - 1.0 # 归一化到 [-1, 1]
|
||||
input_data = img_tensor
|
||||
|
||||
if modality is None:
|
||||
return {'status': 'no_input'}
|
||||
|
||||
# 2. 编码 + forward
|
||||
with torch.no_grad():
|
||||
outputs, info = self.model.forward_multimodal(
|
||||
modality, input_data, state=self.state
|
||||
)
|
||||
self.state = self.model.init_state(1, torch.device(self.device))
|
||||
# 保留 state(forward_multimodal 修改了 state)
|
||||
# 实际上 forward_multimodal 内部更新了 state,但接口设计问题
|
||||
# 这里简化:每次用上一步的 state
|
||||
|
||||
# 3. 输出
|
||||
info['modality'] = modality
|
||||
info['step'] = self.step_count
|
||||
self.step_count += 1
|
||||
|
||||
# 4. 自监督:记录输出用于下一步反馈
|
||||
self.last_output = outputs
|
||||
|
||||
return info
|
||||
|
||||
def run(self, n_steps: int = 100, interval: float = 0.1,
|
||||
on_step: Optional[Callable] = None):
|
||||
"""运行感知-行动循环
|
||||
|
||||
Args:
|
||||
n_steps: 总步数(None = 无限)
|
||||
interval: 每步间隔(秒)
|
||||
on_step: 每步回调
|
||||
"""
|
||||
self.running = True
|
||||
self.stream.start()
|
||||
print(f"感知-行动循环启动,{n_steps} 步,间隔 {interval}s")
|
||||
print("(首次运行 macOS 会请求摄像头/麦克风权限)")
|
||||
|
||||
try:
|
||||
step = 0
|
||||
while self.running and (n_steps is None or step < n_steps):
|
||||
info = self.step_once()
|
||||
|
||||
if info.get('status') != 'no_input':
|
||||
if on_step:
|
||||
on_step(info)
|
||||
|
||||
# 如果有音频输出,播放
|
||||
if self.last_output and 'audio' in self.last_output:
|
||||
audio_out = self.last_output['audio'][0].cpu().numpy()
|
||||
# 每 10 步播放一次(避免太频繁)
|
||||
if step % 10 == 0:
|
||||
self.stream.play_audio(audio_out)
|
||||
|
||||
step += 1
|
||||
time.sleep(interval)
|
||||
except KeyboardInterrupt:
|
||||
print("\n用户中断")
|
||||
finally:
|
||||
self.running = False
|
||||
self.stream.stop()
|
||||
print(f"循环结束,共 {step} 步")
|
||||
101
jspaceai/task.py
Normal file
101
jspaceai/task.py
Normal file
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
玩具任务:连续时间序列预测
|
||||
|
||||
生成多模态时间序列——不同时段由不同的潜在规则驱动:
|
||||
时段 A(t mod 4 ∈ [0,1)):正弦波 + 慢漂移
|
||||
时段 B(t mod 4 ∈ [1,2)):快速振荡
|
||||
时段 C(t mod 4 ∈ [2,3)):脉冲信号
|
||||
时段 D(t mod 4 ∈ [3,4)):组合信号
|
||||
|
||||
任务:给定 x_t,预测 x_{t+1}。
|
||||
|
||||
设计意图:
|
||||
- 不同时段需要不同的"专家"处理——验证工作空间能否学到分工
|
||||
- 模式切换时需要工作空间广播给正确的专家——验证 J-space 路由
|
||||
- 扁平 MLP 在模式切换时会糊掉,工作空间架构应该更鲁棒
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class ContinuousSequenceTask:
|
||||
"""生成连续多模态时间序列"""
|
||||
|
||||
def __init__(self, input_dim: int = 8, seq_len: int = 64, seed: int = 42):
|
||||
self.input_dim = input_dim
|
||||
self.seq_len = seq_len
|
||||
self.rng = np.random.default_rng(seed)
|
||||
|
||||
# 每个时段有自己的参数
|
||||
self.phase_len = 16 # 每个模式持续 16 步
|
||||
self.phase_params = self._gen_phase_params()
|
||||
|
||||
def _gen_phase_params(self) -> list[dict]:
|
||||
"""4 个时段,每段不同参数"""
|
||||
params = []
|
||||
for i in range(4):
|
||||
params.append({
|
||||
'freq': self.rng.uniform(0.1, 0.5), # 振荡频率
|
||||
'phase': self.rng.uniform(0, 2 * np.pi), # 相位
|
||||
'drift': self.rng.uniform(-0.02, 0.02), # 慢漂移
|
||||
'amp': self.rng.uniform(0.3, 1.0), # 振幅
|
||||
'noise': self.rng.uniform(0.01, 0.05), # 噪声
|
||||
})
|
||||
return params
|
||||
|
||||
def generate_sequence(self, n_steps: int) -> np.ndarray:
|
||||
"""生成长度为 n_steps 的序列"""
|
||||
x = np.zeros((n_steps, self.input_dim), dtype=np.float32)
|
||||
state = np.zeros(self.input_dim, dtype=np.float32)
|
||||
|
||||
for t in range(n_steps):
|
||||
phase_idx = (t // self.phase_len) % 4
|
||||
p = self.phase_params[phase_idx]
|
||||
|
||||
# 每个维度独立演化,但共享相位
|
||||
for d in range(self.input_dim):
|
||||
d_offset = d * 0.3
|
||||
if phase_idx == 0:
|
||||
# 正弦 + 漂移
|
||||
state[d] = p['amp'] * np.sin(p['freq'] * t + p['phase'] + d_offset) + p['drift'] * t
|
||||
elif phase_idx == 1:
|
||||
# 快速振荡
|
||||
state[d] = p['amp'] * np.sin(p['freq'] * 3 * t + p['phase'] + d_offset)
|
||||
elif phase_idx == 2:
|
||||
# 脉冲
|
||||
pulse = 1.0 if (t % 5 == 0) else 0.0
|
||||
state[d] = p['amp'] * pulse + 0.1 * np.sin(p['phase'] + d_offset)
|
||||
else:
|
||||
# 组合
|
||||
state[d] = (p['amp'] * 0.5 * np.sin(p['freq'] * t + p['phase'] + d_offset)
|
||||
+ p['amp'] * 0.3 * np.cos(p['freq'] * 2 * t + d_offset))
|
||||
|
||||
state[d] += self.rng.normal(0, p['noise'])
|
||||
|
||||
x[t] = state
|
||||
|
||||
# 归一化到 [-1, 1]
|
||||
x = np.tanh(x)
|
||||
return x
|
||||
|
||||
def generate_batch(self, batch_size: int, n_steps: int | None = None) -> torch.Tensor:
|
||||
"""生成一个 batch 的序列"""
|
||||
n_steps = n_steps or self.seq_len
|
||||
batch = np.stack([self.generate_sequence(n_steps) for _ in range(batch_size)])
|
||||
return torch.from_numpy(batch)
|
||||
|
||||
|
||||
class SequenceDataset(Dataset):
|
||||
"""用于 DataLoader 的数据集"""
|
||||
|
||||
def __init__(self, task: ContinuousSequenceTask, n_samples: int, seq_len: int):
|
||||
self.data = [task.generate_sequence(seq_len) for _ in range(n_samples)]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return torch.from_numpy(self.data[idx])
|
||||
92
jspaceai/trainer.py
Normal file
92
jspaceai/trainer.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
训练器:预测学习目标 + 对比训练
|
||||
|
||||
学习目标(第一性原理):
|
||||
智慧系统的核心是"内部建模世界"——即预测下一时刻世界状态。
|
||||
所以 loss = ||x_{t+1} - pred_t||²
|
||||
|
||||
pred_t = model(x_{0:t})
|
||||
|
||||
这是自监督的——不需要标签,只需要时间序列本身。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
from .task import ContinuousSequenceTask
|
||||
|
||||
|
||||
class Trainer:
|
||||
"""通用训练器,适用于 JSpaceModel 和 FlatBaseline"""
|
||||
|
||||
def __init__(self, model: nn.Module, task: ContinuousSequenceTask,
|
||||
lr: float = 1e-3, device: str = 'cpu'):
|
||||
self.model = model.to(device)
|
||||
self.task = task
|
||||
self.device = device
|
||||
self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
|
||||
self.loss_fn = nn.MSELoss()
|
||||
self.history: list[float] = []
|
||||
|
||||
def train_step(self, xs: torch.Tensor) -> float:
|
||||
"""
|
||||
xs: (batch, T, input_dim)
|
||||
returns: loss value
|
||||
"""
|
||||
xs = xs.to(self.device)
|
||||
# 预测目标:x_{t+1} = xs[:, 1:], 输入 xs[:, :-1]
|
||||
# 模型输出 preds[:, t] 应该预测 xs[:, t+1]
|
||||
preds, _ = self.model(xs)
|
||||
# preds: (batch, T, input_dim) —— preds[:, t] 是从 xs[:, :t+1] 预测的下一时刻
|
||||
# 对齐:preds[:, :-1] 预测 xs[:, 1:]
|
||||
pred = preds[:, :-1]
|
||||
target = xs[:, 1:]
|
||||
|
||||
loss = self.loss_fn(pred, target)
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
# 梯度裁剪(ODE 训练容易爆炸)
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
|
||||
self.optimizer.step()
|
||||
|
||||
return loss.item()
|
||||
|
||||
def evaluate(self, xs: torch.Tensor) -> float:
|
||||
"""评估,不更新参数"""
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
xs = xs.to(self.device)
|
||||
preds, _ = self.model(xs)
|
||||
pred = preds[:, :-1]
|
||||
target = xs[:, 1:]
|
||||
loss = self.loss_fn(pred, target).item()
|
||||
self.model.train()
|
||||
return loss
|
||||
|
||||
def train(self, n_steps: int = 500, batch_size: int = 32,
|
||||
eval_interval: int = 50, verbose: bool = True) -> list[float]:
|
||||
"""完整训练循环"""
|
||||
self.model.train()
|
||||
for step in range(n_steps):
|
||||
xs = self.task.generate_batch(batch_size)
|
||||
loss = self.train_step(xs)
|
||||
self.history.append(loss)
|
||||
|
||||
if verbose and (step + 1) % eval_interval == 0:
|
||||
eval_xs = self.task.generate_batch(64)
|
||||
eval_loss = self.evaluate(eval_xs)
|
||||
print(f" step {step+1:4d} | train_loss {loss:.4f} | eval_loss {eval_loss:.4f}")
|
||||
|
||||
return self.history
|
||||
|
||||
def get_attention(self, xs: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
"""获取注意力权重(仅 JSpaceModel 有,用于可解释性)"""
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
xs = xs.to(self.device)
|
||||
_, info = self.model(xs)
|
||||
self.model.train()
|
||||
return info.get('alpha', None)
|
||||
246
main.py
Normal file
246
main.py
Normal file
@@ -0,0 +1,246 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
JspaceAI —— 主程序入口
|
||||
|
||||
对比实验:
|
||||
1. JSpaceModel(工作空间 + J-space 广播 + ODE 动力学)
|
||||
2. FlatBaseline(同样参数量的扁平 MLP)
|
||||
|
||||
任务:多模态连续时间序列预测(4 种模式切换)
|
||||
|
||||
输出:
|
||||
- 训练/评估 loss 曲线对比
|
||||
- 注意力热力图(哪个专家在哪个时段被激活)
|
||||
- 工作空间 ||w|| 演化曲线(输出门控信号)
|
||||
- 预测 vs 真实序列对比
|
||||
|
||||
运行:
|
||||
python main.py
|
||||
python main.py --steps 1000 --device cuda
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
|
||||
from jspaceai import (
|
||||
JSpaceConfig, JSpaceModel, FlatBaseline,
|
||||
ContinuousSequenceTask, Trainer,
|
||||
)
|
||||
|
||||
|
||||
def count_params(model: torch.nn.Module) -> int:
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
def run_experiment(n_steps: int, device: str, outdir: Path):
|
||||
outdir.mkdir(parents=True, exist_ok=True)
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
# 任务
|
||||
task = ContinuousSequenceTask(input_dim=8, seq_len=64, seed=42)
|
||||
|
||||
# 模型 1: JSpaceModel
|
||||
config = JSpaceConfig(
|
||||
input_dim=8,
|
||||
workspace_dim=32,
|
||||
expert_dim=16,
|
||||
num_experts=5,
|
||||
num_wells=4,
|
||||
ode_steps=4,
|
||||
dt=0.1,
|
||||
tau_w=0.3,
|
||||
output_threshold=0.5,
|
||||
jacobian_sparsity=8,
|
||||
noise_std=0.01,
|
||||
)
|
||||
jspace_model = JSpaceModel(config)
|
||||
|
||||
# 模型 2: FlatBaseline(参数量对齐)
|
||||
# JSpaceModel 大约的参数量:
|
||||
# 专家: 5 × (4×16 + 4 + 16×8 + 16×32 + 16×32) ≈ 5 × 1108 = 5540
|
||||
# 工作空间: (8+32)×32 + 32 + 32×32 + 32 ≈ 2144
|
||||
# 预测头: 32×32 + 32 + 32×8 + 8 ≈ 1384
|
||||
# 总计 ≈ 9000
|
||||
flat_model = FlatBaseline(input_dim=8, hidden_dim=90, num_layers=2)
|
||||
|
||||
print("=" * 60)
|
||||
print("JspaceAI 对比实验")
|
||||
print("=" * 60)
|
||||
print(f"任务: 多模态连续序列预测 (input_dim=8, seq_len=64)")
|
||||
print(f"训练步数: {n_steps}")
|
||||
print(f"设备: {device}")
|
||||
print()
|
||||
print(f"JSpaceModel 参数量: {count_params(jspace_model):,}")
|
||||
print(f"FlatBaseline 参数量: {count_params(flat_model):,}")
|
||||
print()
|
||||
|
||||
# 训练
|
||||
print("-" * 60)
|
||||
print("训练 JSpaceModel...")
|
||||
trainer_js = Trainer(jspace_model, task, lr=1e-3, device=device)
|
||||
history_js = trainer_js.train(n_steps=n_steps, batch_size=32, eval_interval=n_steps // 10)
|
||||
|
||||
print()
|
||||
print("-" * 60)
|
||||
print("训练 FlatBaseline...")
|
||||
trainer_flat = Trainer(flat_model, task, lr=1e-3, device=device)
|
||||
history_flat = trainer_flat.train(n_steps=n_steps, batch_size=32, eval_interval=n_steps // 10)
|
||||
|
||||
# 评估
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("最终评估")
|
||||
print("=" * 60)
|
||||
eval_xs = task.generate_batch(128)
|
||||
eval_loss_js = trainer_js.evaluate(eval_xs)
|
||||
eval_loss_flat = trainer_flat.evaluate(eval_xs)
|
||||
print(f"JSpaceModel eval MSE: {eval_loss_js:.6f}")
|
||||
print(f"FlatBaseline eval MSE: {eval_loss_flat:.6f}")
|
||||
winner = "JSpaceModel" if eval_loss_js < eval_loss_flat else "FlatBaseline"
|
||||
improvement = abs(eval_loss_js - eval_loss_flat) / max(eval_loss_js, eval_loss_flat) * 100
|
||||
print(f"胜者: {winner} (相对优势 {improvement:.1f}%)")
|
||||
|
||||
# 可视化
|
||||
print()
|
||||
print("生成可视化...")
|
||||
|
||||
# 1. Loss 曲线对比
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
|
||||
ax = axes[0, 0]
|
||||
ax.plot(history_js, label='JSpaceModel', alpha=0.7, linewidth=0.8)
|
||||
ax.plot(history_flat, label='FlatBaseline', alpha=0.7, linewidth=0.8)
|
||||
# 平滑曲线
|
||||
if len(history_js) > 20:
|
||||
smooth_js = np.convolve(history_js, np.ones(20)/20, mode='valid')
|
||||
smooth_flat = np.convolve(history_flat, np.ones(20)/20, mode='valid')
|
||||
ax.plot(smooth_js, label='JSpace (smoothed)', linewidth=2)
|
||||
ax.plot(smooth_flat, label='Flat (smoothed)', linewidth=2)
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('MSE Loss')
|
||||
ax.set_title('Training Loss')
|
||||
ax.legend()
|
||||
ax.set_yscale('log')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# 2. 注意力热力图
|
||||
test_xs = task.generate_batch(1) # 单条序列
|
||||
alpha = trainer_js.get_attention(test_xs) # (1, T, num_experts)
|
||||
if alpha is not None:
|
||||
ax = axes[0, 1]
|
||||
alpha_np = alpha[0].cpu().numpy() # (T, num_experts)
|
||||
im = ax.imshow(alpha_np.T, aspect='auto', cmap='hot', interpolation='nearest')
|
||||
ax.set_xlabel('Time step')
|
||||
ax.set_ylabel('Expert index')
|
||||
ax.set_title('J-space Attention (which expert is active)')
|
||||
plt.colorbar(im, ax=ax)
|
||||
# 标注时段边界
|
||||
for i in range(1, 64 // 16):
|
||||
ax.axvline(x=i * 16, color='cyan', linestyle='--', alpha=0.7)
|
||||
ax.set_xticks(range(0, 64, 16))
|
||||
|
||||
# 3. 工作空间 ||w|| 演化
|
||||
jspace_model.eval()
|
||||
with torch.no_grad():
|
||||
xs = test_xs.to(device)
|
||||
_, info = jspace_model(xs)
|
||||
w_norm = info['w_norm'][0].cpu().numpy()
|
||||
|
||||
ax = axes[1, 0]
|
||||
ax.plot(w_norm, label='||w||', linewidth=2)
|
||||
ax.axhline(y=config.output_threshold, color='r', linestyle='--',
|
||||
label=f'threshold={config.output_threshold}', alpha=0.7)
|
||||
ax.set_xlabel('Time step')
|
||||
ax.set_ylabel('||w||')
|
||||
ax.set_title('Workspace norm (output gating signal)')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
for i in range(1, 64 // 16):
|
||||
ax.axvline(x=i * 16, color='gray', linestyle='--', alpha=0.5)
|
||||
|
||||
# 4. 预测 vs 真实(取第 0 维)
|
||||
ax = axes[1, 1]
|
||||
with torch.no_grad():
|
||||
xs = test_xs.to(device)
|
||||
preds_js, _ = jspace_model(xs)
|
||||
preds_flat, _ = flat_model(xs)
|
||||
true_seq = test_xs[0, 1:, 0].numpy()
|
||||
pred_js = preds_js[0, :-1, 0].cpu().numpy()
|
||||
pred_flat = preds_flat[0, :-1, 0].cpu().numpy()
|
||||
|
||||
ax.plot(true_seq, label='True', linewidth=2, color='black')
|
||||
ax.plot(pred_js, label='JSpace', alpha=0.8)
|
||||
ax.plot(pred_flat, label='Flat', alpha=0.8)
|
||||
ax.set_xlabel('Time step')
|
||||
ax.set_ylabel('Value (dim 0)')
|
||||
ax.set_title('Prediction vs True (dim 0)')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
for i in range(1, 64 // 16):
|
||||
ax.axvline(x=i * 16, color='gray', linestyle='--', alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
fig_path = outdir / 'experiment.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"可视化已保存: {fig_path}")
|
||||
plt.close()
|
||||
|
||||
# 保存模型
|
||||
torch.save({
|
||||
'jspace_model': jspace_model.state_dict(),
|
||||
'flat_model': flat_model.state_dict(),
|
||||
'config': config,
|
||||
'eval_loss_js': eval_loss_js,
|
||||
'eval_loss_flat': eval_loss_flat,
|
||||
}, outdir / 'models.pt')
|
||||
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("实验完成")
|
||||
print("=" * 60)
|
||||
print(f"JSpaceModel MSE: {eval_loss_js:.6f}")
|
||||
print(f"FlatBaseline MSE: {eval_loss_flat:.6f}")
|
||||
print(f"胜者: {winner}")
|
||||
print()
|
||||
print("关键观察点(看 experiment.png):")
|
||||
print(" 1. Loss 曲线: JSpace 是否收敛更快/更低?")
|
||||
print(" 2. 注意力热力图: 不同时段是否激活不同专家?")
|
||||
print(" 3. ||w|| 曲线: 模式切换时是否有明显尖峰?")
|
||||
print(" 4. 预测对比: 模式切换处哪个模型更鲁棒?")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='JspaceAI 对比实验')
|
||||
parser.add_argument('--steps', type=int, default=500,
|
||||
help='训练步数 (default: 500)')
|
||||
parser.add_argument('--device', type=str, default='cpu',
|
||||
help='设备 (cpu / cuda / mps)')
|
||||
parser.add_argument('--outdir', type=str, default='outputs',
|
||||
help='输出目录')
|
||||
args = parser.parse_args()
|
||||
|
||||
# 自动检测设备
|
||||
if args.device == 'auto':
|
||||
if torch.cuda.is_available():
|
||||
device = 'cuda'
|
||||
elif torch.backends.mps.is_available():
|
||||
device = 'mps'
|
||||
else:
|
||||
device = 'cpu'
|
||||
else:
|
||||
device = args.device
|
||||
|
||||
run_experiment(
|
||||
n_steps=args.steps,
|
||||
device=device,
|
||||
outdir=Path(args.outdir),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
97
main_autonomous.py
Normal file
97
main_autonomous.py
Normal file
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env python3
|
||||
"""JspaceAI 自主心智 demo —— 永不停止的自主进化
|
||||
|
||||
运行:
|
||||
python main_autonomous.py --steps 100
|
||||
# 再次运行会从上次状态继续
|
||||
python main_autonomous.py --steps 100
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse, torch, numpy as np, matplotlib.pyplot as plt
|
||||
from jspaceai import (
|
||||
MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent,
|
||||
AutonomousMind, PLATFORM,
|
||||
)
|
||||
|
||||
|
||||
def get_config():
|
||||
return MultimodalConfig(
|
||||
vocab_size=50, embed_dim=16, input_dim=8, workspace_dim=64,
|
||||
expert_dim=24, num_experts=12, num_wells=4, ode_steps=3,
|
||||
dt=0.1, tau_w=0.3, jacobian_sparsity=16, noise_std=0.01,
|
||||
img_size=32, audio_frame_size=1024, keyboard_vocab=128,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description='JspaceAI 自主心智')
|
||||
p.add_argument('--steps', type=int, default=100)
|
||||
p.add_argument('--device', default='cpu')
|
||||
p.add_argument('--unsafe', action='store_true')
|
||||
args = p.parse_args()
|
||||
|
||||
dev = args.device
|
||||
if dev == 'auto':
|
||||
dev = 'cuda' if torch.cuda.is_available() else (
|
||||
'mps' if torch.backends.mps.is_available() else 'cpu')
|
||||
|
||||
print(f"平台: {PLATFORM}")
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(dev)
|
||||
model.eval()
|
||||
|
||||
agent = EmbodiedAgent(
|
||||
model, device=dev,
|
||||
enable_mouse_output=args.unsafe,
|
||||
enable_keyboard_output=args.unsafe,
|
||||
enable_audio_output=True, enable_screen_output=False,
|
||||
)
|
||||
mind = AutonomousMind(agent, save_dir='outputs/mind', device=dev)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("自主心智 - 永不停止的进化")
|
||||
print("=" * 60)
|
||||
print("1. 好奇心驱动 2. 状态持久化 3. 自我模型 4. 元学习")
|
||||
print(f"运行 {args.steps} 步(Ctrl+C 中断,状态自动保存)\n")
|
||||
|
||||
log = []
|
||||
|
||||
def on_step(info):
|
||||
log.append(info)
|
||||
if info['step'] % 10 == 0:
|
||||
print(f" step {info['step']:4d} | mod {info['modality']:8s} | "
|
||||
f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | "
|
||||
f"success {info['success']:.2f} | weak={info['weakness']}")
|
||||
|
||||
mind.run(n_steps=args.steps, interval=0.2, save_every=30, on_step=on_step)
|
||||
print("\n" + mind.introspect())
|
||||
|
||||
if log:
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
steps = [s['step'] for s in log]
|
||||
|
||||
axes[0, 0].plot(steps, [s['w_norm'] for s in log], 'b-')
|
||||
axes[0, 0].set_title('Workspace ||w||'); axes[0, 0].grid(True, alpha=0.3)
|
||||
|
||||
axes[0, 1].plot(steps, [s['curiosity'] for s in log], 'r-', label='curiosity')
|
||||
axes[0, 1].plot(steps, [s['success'] for s in log], 'g-', label='success')
|
||||
axes[0, 1].legend(); axes[0, 1].set_title('Curiosity & Success')
|
||||
axes[0, 1].grid(True, alpha=0.3)
|
||||
|
||||
fc = log[-1]['self_confidence']
|
||||
axes[1, 0].barh(list(fc.keys()), list(fc.values()),
|
||||
color=plt.cm.RdYlGn(list(fc.values())))
|
||||
axes[1, 0].set_xlim(0, 1); axes[1, 0].set_title('Self Model')
|
||||
|
||||
axes[1, 1].plot(steps, [s['world_loss'] for s in log], 'orange')
|
||||
axes[1, 1].set_title('World Model Loss'); axes[1, 1].grid(True, alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
Path('outputs').mkdir(exist_ok=True)
|
||||
plt.savefig('outputs/autonomous_mind.png', dpi=150, bbox_inches='tight')
|
||||
print(f"\n可视化: outputs/autonomous_mind.png")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from pathlib import Path
|
||||
main()
|
||||
212
main_embodied.py
Normal file
212
main_embodied.py
Normal file
@@ -0,0 +1,212 @@
|
||||
#!/usr/bin/env python3
|
||||
"""JspaceAI 具身 Agent - 完整神经系统 + 跨平台"""
|
||||
from __future__ import annotations
|
||||
import argparse, torch, numpy as np, matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
import time
|
||||
from jspaceai import (
|
||||
MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent,
|
||||
PLATFORM, get_screen_size, print_permission_guide,
|
||||
check_camera_permission, check_microphone_permission,
|
||||
check_input_monitoring_permission,
|
||||
)
|
||||
|
||||
|
||||
def get_config():
|
||||
return MultimodalConfig(
|
||||
vocab_size=50, embed_dim=16, input_dim=8, workspace_dim=64,
|
||||
expert_dim=24, num_experts=12, num_wells=4, ode_steps=3,
|
||||
dt=0.1, tau_w=0.3, jacobian_sparsity=16, noise_std=0.01,
|
||||
img_size=32, audio_frame_size=1024, keyboard_vocab=128,
|
||||
)
|
||||
|
||||
|
||||
def test_subsystems():
|
||||
print("=" * 60)
|
||||
print(f"具身 Agent 子系统测试 | 平台: {PLATFORM}")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n1. 权限检查:")
|
||||
checks = {
|
||||
'摄像头': check_camera_permission(),
|
||||
'麦克风': check_microphone_permission(),
|
||||
'键盘/鼠标监听': check_input_monitoring_permission(),
|
||||
}
|
||||
for name, ok in checks.items():
|
||||
print(f" {name}: {'OK' if ok else '需要权限'}")
|
||||
if not all(checks.values()):
|
||||
print_permission_guide()
|
||||
|
||||
sw, sh = get_screen_size()
|
||||
print(f" 屏幕尺寸: {sw}x{sh}")
|
||||
|
||||
print("\n2. 模型 + Agent 初始化:")
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config)
|
||||
print(f" 模型参数: {sum(p.numel() for p in model.parameters()):,}")
|
||||
print(f" 专家分工: {model.expert_modality}")
|
||||
|
||||
agent = EmbodiedAgent(
|
||||
model, device='cpu',
|
||||
enable_mouse_output=False, enable_keyboard_output=False,
|
||||
enable_audio_output=True, enable_screen_output=False,
|
||||
)
|
||||
print(f" 小脑参数: {sum(p.numel() for p in agent.cerebellum.parameters()):,}")
|
||||
print(f" 海马体容量: {agent.hippocampus.capacity}")
|
||||
print(f" 反射弧数: {len(agent.cns.reflexes)}")
|
||||
|
||||
print("\n3. 单步循环测试(2秒采集):")
|
||||
agent.senses.start()
|
||||
time.sleep(2)
|
||||
for i in range(5):
|
||||
info = agent.step_once()
|
||||
print(f" step {info['step']:2d} | mod {info['modality']:8s} | "
|
||||
f"||w|| {info['w_norm']:.3f} | action {info['action']['action_idx']} | "
|
||||
f"executed {info['action']['executed']} | mem {info['memories_count']}")
|
||||
time.sleep(0.5)
|
||||
agent.senses.stop()
|
||||
agent.audio_actuator.stop()
|
||||
|
||||
print("\n4. 海马体回忆测试:")
|
||||
if agent.hippocampus and agent.hippocampus.size() > 0:
|
||||
w = agent.state['w'][0].cpu().numpy()
|
||||
for i, m in enumerate(agent.hippocampus.recall(w, top_k=3)):
|
||||
print(f" 记忆 {i}: sim={m['similarity']:.3f} ctx={m['context']}")
|
||||
|
||||
print("\n所有子系统测试完成")
|
||||
|
||||
|
||||
def live_embodied(n_steps, device, safe_mode=False):
|
||||
print(f"平台: {PLATFORM}")
|
||||
sw, sh = get_screen_size()
|
||||
print(f"屏幕: {sw}x{sh}")
|
||||
if safe_mode:
|
||||
print("\n安全模式:不执行鼠标/键盘动作")
|
||||
print_permission_guide()
|
||||
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(device)
|
||||
mp = Path('outputs/multimodal_model.pt')
|
||||
if mp.exists():
|
||||
try:
|
||||
ckpt = torch.load(mp, map_location=device, weights_only=False)
|
||||
model.load_state_dict(ckpt['model'], strict=False)
|
||||
print(f"已加载模型: {mp}")
|
||||
except Exception:
|
||||
print("模型加载失败,随机初始化")
|
||||
else:
|
||||
print("未找到训练模型,随机初始化")
|
||||
model.eval()
|
||||
|
||||
agent = EmbodiedAgent(
|
||||
model, device=device,
|
||||
enable_mouse_output=not safe_mode,
|
||||
enable_keyboard_output=not safe_mode,
|
||||
enable_audio_output=True,
|
||||
enable_screen_output=not safe_mode,
|
||||
risk_threshold=0.5 if safe_mode else 0.3,
|
||||
)
|
||||
|
||||
log = []
|
||||
|
||||
def on_step(info):
|
||||
log.append(info)
|
||||
if info['step'] % 5 == 0:
|
||||
a = info['action']
|
||||
print(f" step {info['step']:3d} | mod {info['modality']:8s} | "
|
||||
f"||w|| {info['w_norm']:.3f} | action {a['action_idx']} "
|
||||
f"str {a['action_strength']:.2f} "
|
||||
f"exec {'Y' if a['executed'] else 'N'} | "
|
||||
f"mem {info['memories_count']}")
|
||||
|
||||
print(f"\n运行 {n_steps} 步具身循环...")
|
||||
agent.run(n_steps=n_steps, interval=0.2, on_step=on_step)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("总结")
|
||||
print("=" * 60)
|
||||
print(f"总步数: {len(log)}")
|
||||
if log:
|
||||
wn = [s['w_norm'] for s in log]
|
||||
print(f"||w||: [{min(wn):.3f}, {max(wn):.3f}] mean={np.mean(wn):.3f}")
|
||||
mc = {}
|
||||
for s in log:
|
||||
mc[s['modality']] = mc.get(s['modality'], 0) + 1
|
||||
print("\n模态分布:")
|
||||
for m, c in sorted(mc.items(), key=lambda x: -x[1]):
|
||||
print(f" {m:10s}: {c:3d} ({c/len(log)*100:.0f}%)")
|
||||
ex = sum(1 for s in log if s['action']['executed'])
|
||||
print(f"\n动作执行: {ex}/{len(log)} ({ex/len(log)*100:.0f}%)")
|
||||
print("\n基底神经节:")
|
||||
for i, c in enumerate(agent.basal_ganglia.habit_counts):
|
||||
h = " (习惯化)" if agent.basal_ganglia.is_habitual(i) else ""
|
||||
print(f" 动作{i}: {int(c):4d}{h}")
|
||||
if agent.hippocampus:
|
||||
print(f"\n海马体: {agent.hippocampus.size()} 条记忆")
|
||||
|
||||
if log:
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
steps = [s['step'] for s in log]
|
||||
wn = [s['w_norm'] for s in log]
|
||||
mods = [s['modality'] for s in log]
|
||||
cm = {'image': 'green', 'screen': 'purple', 'audio': 'blue',
|
||||
'keyboard': 'orange', 'mouse': 'red', 'idle': 'gray'}
|
||||
colors = [cm.get(m, 'gray') for m in mods]
|
||||
|
||||
ax = axes[0, 0]
|
||||
ax.scatter(steps, wn, c=colors, alpha=0.7, s=30)
|
||||
ax.set_xlabel('Step'); ax.set_ylabel('||w||')
|
||||
ax.set_title('Workspace Norm'); ax.grid(True, alpha=0.3)
|
||||
from matplotlib.patches import Patch
|
||||
ax.legend(handles=[Patch(facecolor=c, label=m) for m, c in cm.items()],
|
||||
fontsize=7)
|
||||
|
||||
ax = axes[0, 1]
|
||||
mc2 = {}
|
||||
for m in mods: mc2[m] = mc2.get(m, 0) + 1
|
||||
ax.bar(mc2.keys(), mc2.values(),
|
||||
color=[cm.get(m, 'gray') for m in mc2.keys()])
|
||||
ax.set_title('Modality Distribution')
|
||||
|
||||
ax = axes[1, 0]
|
||||
ast = [s['action']['action_strength'] for s in log]
|
||||
exf = [1 if s['action']['executed'] else 0 for s in log]
|
||||
ax.plot(steps, ast, label='strength', alpha=0.7)
|
||||
ax.scatter(steps, exf, c=['green' if e else 'red' for e in exf],
|
||||
label='executed', alpha=0.5, s=20)
|
||||
ax.set_title('Action Strength & Execution'); ax.legend(); ax.grid(True, alpha=0.3)
|
||||
|
||||
ax = axes[1, 1]
|
||||
ai = [s['action']['action_idx'] for s in log]
|
||||
ax.hist(ai, bins=range(6), align='left', rwidth=0.8, color='steelblue', alpha=0.7)
|
||||
ax.set_title('Basal Ganglia Action Selection')
|
||||
ax.set_xticks(range(5))
|
||||
|
||||
plt.tight_layout()
|
||||
Path('outputs').mkdir(exist_ok=True)
|
||||
plt.savefig('outputs/embodied_live.png', dpi=150, bbox_inches='tight')
|
||||
print(f"\n可视化: outputs/embodied_live.png")
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description='JspaceAI 具身 Agent')
|
||||
p.add_argument('--mode', default='test', choices=['test', 'live', 'safe'])
|
||||
p.add_argument('--steps', type=int, default=50)
|
||||
p.add_argument('--device', default='cpu')
|
||||
args = p.parse_args()
|
||||
|
||||
dev = args.device
|
||||
if dev == 'auto':
|
||||
dev = 'cuda' if torch.cuda.is_available() else (
|
||||
'mps' if torch.backends.mps.is_available() else 'cpu')
|
||||
|
||||
if args.mode == 'test':
|
||||
test_subsystems()
|
||||
elif args.mode == 'live':
|
||||
live_embodied(args.steps, dev, safe_mode=False)
|
||||
elif args.mode == 'safe':
|
||||
live_embodied(args.steps, dev, safe_mode=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
295
main_full_sensory.py
Normal file
295
main_full_sensory.py
Normal file
@@ -0,0 +1,295 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
JspaceAI 全感官交互 demo
|
||||
|
||||
接入全部 5 个输入通道:摄像头 + 麦克风 + 屏幕 + 键盘 + 鼠标
|
||||
|
||||
模式:
|
||||
--mode test: 测试各 I/O 通道是否工作
|
||||
--mode live: 实时全感官感知循环
|
||||
|
||||
运行:
|
||||
python main_full_sensory.py --mode test
|
||||
python main_full_sensory.py --mode live --steps 50
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
import time
|
||||
|
||||
from jspaceai import (
|
||||
MultimodalConfig, MultimodalJSpaceModel,
|
||||
FullSensoryStream,
|
||||
)
|
||||
from jspaceai.platform import get_screen_size, PLATFORM, print_permission_guide
|
||||
|
||||
|
||||
def get_config() -> MultimodalConfig:
|
||||
return MultimodalConfig(
|
||||
vocab_size=50, embed_dim=16, input_dim=8,
|
||||
workspace_dim=64, expert_dim=24, num_experts=12,
|
||||
num_wells=4, ode_steps=3, dt=0.1, tau_w=0.3,
|
||||
jacobian_sparsity=16, noise_std=0.01,
|
||||
img_size=32, audio_frame_size=1024, keyboard_vocab=128,
|
||||
)
|
||||
|
||||
|
||||
def test_io():
|
||||
"""测试所有 I/O 通道"""
|
||||
print("=" * 60)
|
||||
print("测试全感官 I/O 通道")
|
||||
print("=" * 60)
|
||||
|
||||
stream = FullSensoryStream(
|
||||
use_camera=True, use_mic=True, use_desktop=True, img_size=(32, 32),
|
||||
)
|
||||
stream.start()
|
||||
|
||||
print("\n采集 5 秒数据...(请移动鼠标、按键、对摄像头说话)")
|
||||
log = {'camera': 0, 'audio': 0, 'screen': 0, 'keyboard': 0, 'mouse': 0}
|
||||
mouse_positions = []
|
||||
|
||||
for i in range(25):
|
||||
data = stream.get_latest()
|
||||
if data['camera']:
|
||||
log['camera'] += 1
|
||||
if data['audio']:
|
||||
log['audio'] += 1
|
||||
if data['screen']:
|
||||
log['screen'] += 1
|
||||
if data['keyboard']:
|
||||
log['keyboard'] += len(data['keyboard'])
|
||||
if data['mouse']:
|
||||
log['mouse'] += len(data['mouse'])
|
||||
mouse_positions.append(data['mouse_pos'])
|
||||
time.sleep(0.2)
|
||||
|
||||
stream.stop()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("采集结果")
|
||||
print("=" * 60)
|
||||
for ch, count in log.items():
|
||||
status = "OK" if count > 0 else "无数据"
|
||||
print(f" {ch:10s}: {count:4d} 帧/事件 [{status}]")
|
||||
|
||||
if mouse_positions:
|
||||
xs = [p[0] for p in mouse_positions]
|
||||
ys = [p[1] for p in mouse_positions]
|
||||
print(f" 鼠标位置范围: x=[{min(xs)}, {max(xs)}], y=[{min(ys)}, {max(ys)}]")
|
||||
|
||||
# 可视化
|
||||
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
|
||||
|
||||
ax = axes[0]
|
||||
if mouse_positions:
|
||||
xs = [p[0] for p in mouse_positions]
|
||||
ys = [p[1] for p in mouse_positions]
|
||||
ax.plot(xs, ys, 'b.-', alpha=0.5)
|
||||
ax.set_xlabel('X')
|
||||
ax.set_ylabel('Y')
|
||||
ax.set_title(f'Mouse Trajectory ({len(mouse_positions)} points)')
|
||||
ax.invert_yaxis()
|
||||
else:
|
||||
ax.text(0.5, 0.5, "No mouse movement", transform=ax.transAxes, ha='center')
|
||||
ax.set_title('Mouse')
|
||||
|
||||
ax = axes[1]
|
||||
channels = list(log.keys())
|
||||
counts = list(log.values())
|
||||
ax.bar(channels, counts, color=['green', 'blue', 'purple', 'orange', 'red'])
|
||||
ax.set_ylabel('Event count')
|
||||
ax.set_title('Channel Activity')
|
||||
|
||||
plt.tight_layout()
|
||||
outdir = Path('outputs')
|
||||
outdir.mkdir(exist_ok=True)
|
||||
fig_path = outdir / 'io_test.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"\n可视化: {fig_path}")
|
||||
|
||||
|
||||
def live_full_sensory(n_steps: int, device: str):
|
||||
"""实时全感官感知循环"""
|
||||
screen_w, screen_h = get_screen_size()
|
||||
print(f"平台: {PLATFORM}")
|
||||
print(f"屏幕: {screen_w}x{screen_h}")
|
||||
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(device)
|
||||
|
||||
model_path = Path('outputs/multimodal_model.pt')
|
||||
if model_path.exists():
|
||||
ckpt = torch.load(model_path, map_location=device, weights_only=False)
|
||||
try:
|
||||
model.load_state_dict(ckpt['model'], strict=False)
|
||||
print(f"已加载模型: {model_path}")
|
||||
except Exception as e:
|
||||
print(f"模型加载失败,用随机初始化: {e}")
|
||||
else:
|
||||
print("未找到训练模型,用随机初始化")
|
||||
model.eval()
|
||||
|
||||
print(f"\n专家分工: {model.expert_modality}")
|
||||
|
||||
print("\n启动全感官流(摄像头+麦克风+屏幕+键盘+鼠标)...")
|
||||
stream = FullSensoryStream(
|
||||
use_camera=True, use_mic=True, use_desktop=True, img_size=(32, 32),
|
||||
)
|
||||
stream.start()
|
||||
|
||||
state = model.init_state(1, torch.device(device))
|
||||
step_log = []
|
||||
|
||||
print(f"\n运行 {n_steps} 步感知循环...")
|
||||
print("(移动鼠标、按键、对摄像头说话/做动作,模型会持续感知)")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
for step in range(n_steps):
|
||||
data = stream.get_latest()
|
||||
|
||||
modality = None
|
||||
input_tensor = None
|
||||
|
||||
if data['keyboard']:
|
||||
key_events = data['keyboard']
|
||||
key_ids = []
|
||||
for ev in key_events[:8]:
|
||||
k = ev.data if isinstance(ev.data, str) else ' '
|
||||
key_ids.append(ord(k[0]) if k and len(k) == 1 else 0)
|
||||
if key_ids:
|
||||
modality = 'keyboard'
|
||||
input_tensor = torch.tensor([key_ids], dtype=torch.long).to(device)
|
||||
|
||||
elif data['mouse']:
|
||||
ev = data['mouse'][0]
|
||||
x, y = ev.data
|
||||
x_norm = x / screen_w
|
||||
y_norm = y / screen_h
|
||||
click_l = 1.0 if ev.modifiers.get('button') == 'left' else 0.0
|
||||
click_r = 1.0 if ev.modifiers.get('button') == 'right' else 0.0
|
||||
modality = 'mouse'
|
||||
input_tensor = torch.tensor([[x_norm, y_norm, click_l, click_r]],
|
||||
dtype=torch.float32).to(device)
|
||||
|
||||
elif data['audio']:
|
||||
audio = data['audio'].data
|
||||
modality = 'audio'
|
||||
input_tensor = torch.tensor([audio], dtype=torch.float32).to(device)
|
||||
|
||||
elif data['screen']:
|
||||
img = data['screen'].data
|
||||
img_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device) / 127.5 - 1.0
|
||||
modality = 'screen'
|
||||
input_tensor = img_tensor
|
||||
|
||||
elif data['camera']:
|
||||
img = data['camera'].data
|
||||
img_tensor = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device) / 127.5 - 1.0
|
||||
modality = 'image'
|
||||
input_tensor = img_tensor
|
||||
|
||||
if modality is None:
|
||||
time.sleep(0.2)
|
||||
continue
|
||||
|
||||
with torch.no_grad():
|
||||
x = model.encode_modality(modality, input_tensor)
|
||||
if x.dim() == 1:
|
||||
x = x.unsqueeze(0)
|
||||
if x.dim() == 3:
|
||||
x = x[:, -1, :]
|
||||
if x.shape[0] != 1:
|
||||
x = x[-1:]
|
||||
state, _ = model.step(state, x)
|
||||
w = state['w']
|
||||
w_norm = w.norm(dim=-1).mean().item()
|
||||
|
||||
step_log.append({
|
||||
'step': step, 'modality': modality, 'w_norm': w_norm,
|
||||
})
|
||||
|
||||
if step % 5 == 0:
|
||||
if step % 10 == 0 and step > 0:
|
||||
with torch.no_grad():
|
||||
audio_out = model.audio_decoder(w)[0].cpu().numpy()
|
||||
stream.play_audio(audio_out)
|
||||
|
||||
mouse_pos = data.get('mouse_pos', (0, 0))
|
||||
key_buf = data.get('keyboard_buffer', '')[:20]
|
||||
print(f" step {step:3d} | mod {modality:8s} | ||w|| {w_norm:.3f} | "
|
||||
f"mouse={mouse_pos} | keys={repr(key_buf)}")
|
||||
|
||||
time.sleep(0.15)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n用户中断")
|
||||
finally:
|
||||
stream.stop()
|
||||
|
||||
# 可视化
|
||||
if step_log:
|
||||
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
|
||||
|
||||
ax = axes[0]
|
||||
steps = [s['step'] for s in step_log]
|
||||
w_norms = [s['w_norm'] for s in step_log]
|
||||
modalities = [s['modality'] for s in step_log]
|
||||
color_map = {'image': 'green', 'screen': 'purple', 'audio': 'blue',
|
||||
'keyboard': 'orange', 'mouse': 'red'}
|
||||
colors = [color_map.get(m, 'gray') for m in modalities]
|
||||
ax.scatter(steps, w_norms, c=colors, alpha=0.7, s=30)
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('||w||')
|
||||
ax.set_title('Workspace Norm by Modality')
|
||||
ax.grid(True, alpha=0.3)
|
||||
from matplotlib.patches import Patch
|
||||
legend_elements = [Patch(facecolor=c, label=m) for m, c in color_map.items()]
|
||||
ax.legend(handles=legend_elements)
|
||||
|
||||
ax = axes[1]
|
||||
mod_counts = {}
|
||||
for m in modalities:
|
||||
mod_counts[m] = mod_counts.get(m, 0) + 1
|
||||
ax.bar(mod_counts.keys(), mod_counts.values(),
|
||||
color=[color_map.get(m, 'gray') for m in mod_counts.keys()])
|
||||
ax.set_ylabel('Count')
|
||||
ax.set_title('Modality Distribution')
|
||||
|
||||
plt.tight_layout()
|
||||
outdir = Path('outputs')
|
||||
outdir.mkdir(exist_ok=True)
|
||||
fig_path = outdir / 'full_sensory_live.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"\n可视化: {fig_path}")
|
||||
|
||||
print(f"\n完成,共 {len(step_log)} 步")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='JspaceAI 全感官交互')
|
||||
parser.add_argument('--mode', type=str, default='test',
|
||||
choices=['test', 'live'])
|
||||
parser.add_argument('--steps', type=int, default=50)
|
||||
parser.add_argument('--device', type=str, default='cpu')
|
||||
args = parser.parse_args()
|
||||
|
||||
device = args.device
|
||||
if device == 'auto':
|
||||
if torch.cuda.is_available(): device = 'cuda'
|
||||
elif torch.backends.mps.is_available(): device = 'mps'
|
||||
else: device = 'cpu'
|
||||
|
||||
if args.mode == 'test':
|
||||
test_io()
|
||||
elif args.mode == 'live':
|
||||
live_full_sensory(args.steps, device)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
222
main_language.py
Normal file
222
main_language.py
Normal file
@@ -0,0 +1,222 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
JspaceAI 语言版 —— 自主进化实验
|
||||
|
||||
模型在 Shakespeare 文本上持续学习,边推理边进化。
|
||||
观察:
|
||||
1. loss 持续下降(在学习)
|
||||
2. 生成文本从乱码逐渐变成类 Shakespeare 风格
|
||||
3. 专家分工涌现(不同专家处理不同字符模式)
|
||||
4. EWC + 经验回放防止灾难性遗忘
|
||||
|
||||
运行:
|
||||
python main_language.py
|
||||
python main_language.py --steps 300 --device mps
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
|
||||
from jspaceai import (
|
||||
LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
|
||||
CharTokenizer, load_shakespeare,
|
||||
)
|
||||
|
||||
|
||||
def run_language_experiment(n_steps: int, device: str, outdir: Path):
|
||||
outdir.mkdir(parents=True, exist_ok=True)
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
# 1. 数据
|
||||
text = load_shakespeare()
|
||||
tokenizer = CharTokenizer.from_text(text)
|
||||
print(f"文本长度: {len(text)} 字符")
|
||||
print(f"词汇表大小: {tokenizer.vocab_size}")
|
||||
print(f"词汇表: {''.join(tokenizer.chars[:50])}...")
|
||||
|
||||
# 切成多段,模拟持续到来的文本流
|
||||
chunk_size = 200
|
||||
text_stream = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
|
||||
print(f"文本流: {len(text_stream)} 段, 每段 {chunk_size} 字符")
|
||||
|
||||
# 2. 模型
|
||||
config = LanguageConfig(
|
||||
vocab_size=tokenizer.vocab_size,
|
||||
embed_dim=16,
|
||||
input_dim=8, # J-space 输入维度
|
||||
workspace_dim=32, # 工作空间维度
|
||||
expert_dim=16, # 每个专家内部状态
|
||||
num_experts=5, # 5 个并行专家
|
||||
num_wells=4, # 每个专家 4 个吸引子
|
||||
ode_steps=3, # ODE 积分子步
|
||||
dt=0.1,
|
||||
tau_w=0.3,
|
||||
jacobian_sparsity=8,
|
||||
noise_std=0.005,
|
||||
)
|
||||
model = JSpaceLanguageModel(config)
|
||||
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
print(f"\n模型参数量: {n_params:,}")
|
||||
|
||||
# 3. 自主进化训练器
|
||||
trainer = EvolutionTrainer(
|
||||
model, config,
|
||||
lr=5e-3,
|
||||
ewc_lambda=0.05, # 较小的 EWC 权重,让模型能学新东西
|
||||
device=device,
|
||||
)
|
||||
|
||||
# 4. 进化
|
||||
print("\n" + "=" * 70)
|
||||
print("自主进化开始")
|
||||
print("=" * 70)
|
||||
|
||||
# 限制步数
|
||||
history = trainer.evolve(
|
||||
text_stream,
|
||||
tokenizer,
|
||||
seq_len=48,
|
||||
batch_size=4,
|
||||
consolidate_every=30,
|
||||
generate_every=30,
|
||||
max_steps=n_steps,
|
||||
prompt_text="To be",
|
||||
)
|
||||
|
||||
# 5. 总结
|
||||
summary = trainer.get_evolution_summary()
|
||||
print("\n" + "=" * 70)
|
||||
print("进化总结")
|
||||
print("=" * 70)
|
||||
print(f"总步数: {summary['steps']}")
|
||||
print(f"初始 loss: {summary['initial_loss']:.4f}")
|
||||
print(f"最终 loss: {summary['final_loss']:.4f}")
|
||||
print(f"最低 loss: {summary['min_loss']:.4f}")
|
||||
print(f"loss 下降: {(1 - summary['final_loss']/summary['initial_loss'])*100:.1f}%")
|
||||
print(f"\n专家最终使用率: {[f'{u:.3f}' for u in summary['expert_usage']]}")
|
||||
|
||||
# 专家专业化(每个专家最常处理的 top-5 字符)
|
||||
print("\n专家专业化(top-5 字符):")
|
||||
for i, spec in enumerate(summary['expert_specialization']):
|
||||
chars = tokenizer.decode([s[0] for s in spec])
|
||||
weights = [f"{s[1]:.1f}" for s in spec]
|
||||
display = ' '.join(f"{repr(c)}({w})" for c, w in zip(chars, weights))
|
||||
print(f" 专家 {i}: {display}")
|
||||
|
||||
# 6. 最终生成对比
|
||||
print("\n" + "=" * 70)
|
||||
print("最终生成对比")
|
||||
print("=" * 70)
|
||||
for prompt in ["To be", "Romeo", "The "]:
|
||||
prompt_ids = tokenizer.encode(prompt)
|
||||
generated = model.generate(prompt_ids, n_new=100, temperature=0.7, top_k=5)
|
||||
sample = prompt + tokenizer.decode(generated)
|
||||
print(f"\n提示 '{prompt}':")
|
||||
print(f" {sample}")
|
||||
|
||||
# 7. 可视化
|
||||
print("\n生成可视化...")
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
|
||||
# Loss 曲线
|
||||
ax = axes[0, 0]
|
||||
losses = [h['loss'] for h in history]
|
||||
ax.plot(losses, alpha=0.3, linewidth=0.5, color='blue', label='raw')
|
||||
if len(losses) > 10:
|
||||
smoothed = np.convolve(losses, np.ones(10)/10, mode='valid')
|
||||
ax.plot(smoothed, linewidth=2, color='blue', label='smoothed')
|
||||
ax.set_xlabel('Evolution step')
|
||||
ax.set_ylabel('Cross-entropy loss')
|
||||
ax.set_title('Loss during self-evolution')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# 专家使用率演化
|
||||
ax = axes[0, 1]
|
||||
alpha_history = np.array([h['alpha_mean'] for h in history])
|
||||
for i in range(config.num_experts):
|
||||
ax.plot(alpha_history[:, i], label=f'Expert {i}', alpha=0.8)
|
||||
ax.set_xlabel('Evolution step')
|
||||
ax.set_ylabel('Attention weight (mean)')
|
||||
ax.set_title('Expert usage during evolution')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# ||w|| 演化
|
||||
ax = axes[1, 0]
|
||||
w_norms = [h['w_norm_mean'] for h in history]
|
||||
ax.plot(w_norms, linewidth=1.5, color='green')
|
||||
ax.set_xlabel('Evolution step')
|
||||
ax.set_ylabel('||w|| (mean)')
|
||||
ax.set_title('Workspace norm during evolution')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# 生成样本随时间演化
|
||||
ax = axes[1, 1]
|
||||
ax.axis('off')
|
||||
samples = [(h['step'], h.get('sample', '')) for h in history if 'sample' in h]
|
||||
text_lines = []
|
||||
for step, sample in samples[:8]:
|
||||
s = sample[:60].replace('\n', ' ')
|
||||
text_lines.append(f"step {step:3d}: {s}")
|
||||
ax.text(0.05, 0.95, '\n'.join(text_lines),
|
||||
transform=ax.transAxes, fontsize=9, verticalalignment='top',
|
||||
fontfamily='monospace',
|
||||
bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
|
||||
ax.set_title('Generated samples during evolution')
|
||||
|
||||
plt.tight_layout()
|
||||
fig_path = outdir / 'language_evolution.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"可视化已保存: {fig_path}")
|
||||
plt.close()
|
||||
|
||||
# 保存模型
|
||||
torch.save({
|
||||
'model': model.state_dict(),
|
||||
'config': config,
|
||||
'tokenizer_chars': tokenizer.chars,
|
||||
'history': history,
|
||||
'summary': summary,
|
||||
}, outdir / 'language_model.pt')
|
||||
print(f"模型已保存: {outdir / 'language_model.pt'}")
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("实验完成")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='JspaceAI 语言版自主进化实验')
|
||||
parser.add_argument('--steps', type=int, default=100,
|
||||
help='文本流段数 (default: 100)')
|
||||
parser.add_argument('--device', type=str, default='cpu',
|
||||
help='设备 (cpu / cuda / mps / auto)')
|
||||
parser.add_argument('--outdir', type=str, default='outputs',
|
||||
help='输出目录')
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.device == 'auto':
|
||||
if torch.cuda.is_available():
|
||||
device = 'cuda'
|
||||
elif torch.backends.mps.is_available():
|
||||
device = 'mps'
|
||||
else:
|
||||
device = 'cpu'
|
||||
else:
|
||||
device = args.device
|
||||
|
||||
run_language_experiment(
|
||||
n_steps=args.steps,
|
||||
device=device,
|
||||
outdir=Path(args.outdir),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
296
main_multimodal.py
Normal file
296
main_multimodal.py
Normal file
@@ -0,0 +1,296 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
JspaceAI 多模态实时交互 demo
|
||||
|
||||
接入摄像头 + 麦克风 + 扬声器,原生支持图像/音频/文本。
|
||||
|
||||
模式:
|
||||
1. --mode train: 在合成多模态数据上训练(无需摄像头权限)
|
||||
2. --mode live: 实时感知-行动循环(需要摄像头/麦克风权限)
|
||||
3. --mode eval: 离线评估多模态对齐能力
|
||||
|
||||
运行:
|
||||
python main_multimodal.py --mode train
|
||||
python main_multimodal.py --mode live --steps 100
|
||||
python main_multimodal.py --mode eval
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
import time
|
||||
|
||||
from jspaceai import (
|
||||
MultimodalConfig, MultimodalJSpaceModel,
|
||||
MultimodalStream, SensoryMotorLoop,
|
||||
)
|
||||
|
||||
|
||||
def get_config(vocab_size: int = 50) -> MultimodalConfig:
|
||||
return MultimodalConfig(
|
||||
vocab_size=vocab_size,
|
||||
embed_dim=16,
|
||||
input_dim=8,
|
||||
workspace_dim=64,
|
||||
expert_dim=24,
|
||||
num_experts=8,
|
||||
num_wells=4,
|
||||
ode_steps=3,
|
||||
dt=0.1,
|
||||
tau_w=0.3,
|
||||
jacobian_sparsity=16,
|
||||
noise_std=0.01,
|
||||
img_size=32,
|
||||
audio_frame_size=1024,
|
||||
)
|
||||
|
||||
|
||||
def generate_synthetic_multimodal(batch_size: int = 4, device: str = 'cpu'):
|
||||
"""生成合成多模态数据:3 个概念对应 3 种模态特征"""
|
||||
import cv2
|
||||
concepts = ['A', 'B', 'C']
|
||||
colors = [(1, 0, 0), (0, 0, 1), (0, 1, 0)]
|
||||
freqs = [200, 800, 400]
|
||||
|
||||
images, audios, tokens = [], [], []
|
||||
for _ in range(batch_size):
|
||||
idx = np.random.randint(3)
|
||||
img = np.zeros((32, 32, 3), dtype=np.float32)
|
||||
color = colors[idx]
|
||||
if idx == 0:
|
||||
cv2.circle(img, (16, 16), 8, color, -1)
|
||||
elif idx == 1:
|
||||
img[8:24, 8:24] = color
|
||||
else:
|
||||
for y in range(32):
|
||||
w = min(y, 31 - y)
|
||||
if 8 < y < 24:
|
||||
img[31-y, 16-w:16+w] = color
|
||||
images.append(img.transpose(2, 0, 1))
|
||||
|
||||
t = np.linspace(0, 1024/16000, 1024)
|
||||
audio = np.sin(2 * np.pi * freqs[idx] * t).astype(np.float32) * 0.5
|
||||
audios.append(audio)
|
||||
tokens.append(idx)
|
||||
|
||||
images = torch.tensor(np.stack(images)).to(device) / 255.0 * 2 - 1
|
||||
audios = torch.tensor(np.stack(audios)).to(device)
|
||||
tokens = torch.tensor(tokens, dtype=torch.long).to(device)
|
||||
return images, audios, tokens
|
||||
|
||||
|
||||
def train_multimodal(n_steps: int, device: str, outdir: Path):
|
||||
outdir.mkdir(parents=True, exist_ok=True)
|
||||
torch.manual_seed(42); np.random.seed(42)
|
||||
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
|
||||
|
||||
print(f"模型参数: {sum(p.numel() for p in model.parameters()):,}")
|
||||
print(f"专家模态: {model.expert_modality}")
|
||||
print(f"\n训练 {n_steps} 步...")
|
||||
history = []
|
||||
|
||||
for step in range(n_steps):
|
||||
images, audios, tokens = generate_synthetic_multimodal(8, device)
|
||||
modality_choice = np.random.randint(3)
|
||||
total_loss = 0
|
||||
|
||||
if modality_choice == 0:
|
||||
outputs, _ = model.forward_multimodal('image', images)
|
||||
text_loss = F.cross_entropy(outputs['text_logits'], tokens)
|
||||
audio_loss = F.mse_loss(outputs['audio'], audios)
|
||||
total_loss = text_loss + 0.1 * audio_loss
|
||||
elif modality_choice == 1:
|
||||
outputs, _ = model.forward_multimodal('audio', audios)
|
||||
text_loss = F.cross_entropy(outputs['text_logits'], tokens)
|
||||
img_loss = F.mse_loss(outputs['image'], images)
|
||||
total_loss = text_loss + 0.1 * img_loss
|
||||
else:
|
||||
outputs, _ = model.forward_multimodal('text', tokens)
|
||||
img_loss = F.mse_loss(outputs['image'], images)
|
||||
audio_loss = F.mse_loss(outputs['audio'], audios)
|
||||
total_loss = 0.1 * img_loss + 0.1 * audio_loss
|
||||
|
||||
optimizer.zero_grad()
|
||||
total_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
history.append(total_loss.item())
|
||||
|
||||
if (step + 1) % 20 == 0:
|
||||
print(f" step {step+1:4d} | loss {total_loss.item():.4f} | "
|
||||
f"modality {['img','aud','txt'][modality_choice]}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("跨模态对齐评估")
|
||||
print("=" * 60)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
for mod_name, mod_data in [('image', images), ('audio', audios), ('text', tokens)]:
|
||||
outputs, _ = model.forward_multimodal(mod_name, mod_data[:3])
|
||||
pred_tokens = outputs['text_logits'].argmax(dim=-1)
|
||||
print(f" 输入 {mod_name:6s} → 预测 token: {pred_tokens.cpu().tolist()} "
|
||||
f"(真实: {tokens[:3].cpu().tolist()})")
|
||||
|
||||
fig, axes = plt.subplots(2, 3, figsize=(12, 8))
|
||||
|
||||
ax = axes[0, 0]
|
||||
ax.plot(history, alpha=0.5, linewidth=0.5)
|
||||
if len(history) > 10:
|
||||
smoothed = np.convolve(history, np.ones(10)/10, mode='valid')
|
||||
ax.plot(smoothed, linewidth=2)
|
||||
ax.set_title('Training Loss')
|
||||
ax.set_xlabel('Step')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
with torch.no_grad():
|
||||
test_img = images[:1]
|
||||
outputs, _ = model.forward_multimodal('image', test_img)
|
||||
|
||||
ax = axes[0, 1]
|
||||
ax.imshow(((test_img[0].cpu().permute(1,2,0) + 1) / 2).numpy())
|
||||
ax.set_title('Input Image')
|
||||
ax.axis('off')
|
||||
|
||||
ax = axes[0, 2]
|
||||
ax.imshow(((outputs['image'][0].cpu().permute(1,2,0) + 1) / 2).clamp(0,1).numpy())
|
||||
ax.set_title('Reconstructed Image')
|
||||
ax.axis('off')
|
||||
|
||||
ax = axes[1, 0]
|
||||
input_audio = audios[0].cpu().numpy()
|
||||
ax.plot(input_audio[:200], alpha=0.7, label='input')
|
||||
out_audio = outputs['audio'][0].cpu().numpy()
|
||||
ax.plot(out_audio[:200], alpha=0.7, label='reconstructed')
|
||||
ax.set_title('Audio Waveform')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
ax = axes[1, 1]
|
||||
logits = outputs['text_logits'][0].cpu().numpy()
|
||||
probs = np.exp(logits) / np.exp(logits).sum()
|
||||
ax.bar(['A', 'B', 'C'], probs[:3])
|
||||
ax.set_title('Text Logits (softmax)')
|
||||
|
||||
ax = axes[1, 2]
|
||||
w = outputs['w'][0].cpu().numpy()
|
||||
ax.bar(range(len(w)), w)
|
||||
ax.set_title('Workspace w (64-dim)')
|
||||
|
||||
plt.tight_layout()
|
||||
fig_path = outdir / 'multimodal_train.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"\n可视化: {fig_path}")
|
||||
|
||||
torch.save({'model': model.state_dict(), 'config': config},
|
||||
outdir / 'multimodal_model.pt')
|
||||
print(f"模型: {outdir / 'multimodal_model.pt'}")
|
||||
|
||||
|
||||
def live_demo(n_steps: int, device: str, outdir: Path):
|
||||
outdir.mkdir(parents=True, exist_ok=True)
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(device)
|
||||
|
||||
model_path = outdir / 'multimodal_model.pt'
|
||||
if model_path.exists():
|
||||
ckpt = torch.load(model_path, map_location=device)
|
||||
model.load_state_dict(ckpt['model'])
|
||||
print(f"已加载模型: {model_path}")
|
||||
else:
|
||||
print("未找到训练好的模型,用随机初始化运行")
|
||||
model.eval()
|
||||
|
||||
print("\n启动实时多模态流...")
|
||||
print("(macOS 会请求摄像头和麦克风权限,请允许)")
|
||||
stream = MultimodalStream(
|
||||
use_camera=True, use_mic=True,
|
||||
img_size=(32, 32), sample_rate=16000, audio_frame_size=1024,
|
||||
)
|
||||
loop = SensoryMotorLoop(model, stream, device=device)
|
||||
step_log = []
|
||||
|
||||
def on_step(info):
|
||||
mod = info.get('modality', '?')
|
||||
step = info.get('step', 0)
|
||||
w_norm = info.get('w_norm', torch.tensor([0])).mean().item()
|
||||
step_log.append({'step': step, 'modality': mod, 'w_norm': w_norm})
|
||||
if step % 5 == 0:
|
||||
print(f" step {step:3d} | modality {mod:5s} | ||w|| {w_norm:.3f}")
|
||||
|
||||
loop.run(n_steps=n_steps, interval=0.2, on_step=on_step)
|
||||
|
||||
if step_log:
|
||||
fig, ax = plt.subplots(1, 1, figsize=(10, 4))
|
||||
steps = [s['step'] for s in step_log]
|
||||
w_norms = [s['w_norm'] for s in step_log]
|
||||
modalities = [s['modality'] for s in step_log]
|
||||
colors = ['blue' if m == 'audio' else 'green' for m in modalities]
|
||||
ax.scatter(steps, w_norms, c=colors, alpha=0.6, s=20)
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('||w||')
|
||||
ax.set_title('Workspace Norm During Live (blue=audio, green=image)')
|
||||
ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
fig_path = outdir / 'multimodal_live.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"可视化: {fig_path}")
|
||||
|
||||
|
||||
def eval_multimodal(device: str, outdir: Path):
|
||||
model_path = outdir / 'multimodal_model.pt'
|
||||
if not model_path.exists():
|
||||
print("请先运行 --mode train")
|
||||
return
|
||||
|
||||
config = get_config()
|
||||
model = MultimodalJSpaceModel(config).to(device)
|
||||
ckpt = torch.load(model_path, map_location=device)
|
||||
model.load_state_dict(ckpt['model'])
|
||||
model.eval()
|
||||
|
||||
print("多模态对齐评估")
|
||||
print("=" * 60)
|
||||
|
||||
images, audios, tokens = generate_synthetic_multimodal(6, device)
|
||||
|
||||
with torch.no_grad():
|
||||
for mod_name, mod_data in [('image', images), ('audio', audios), ('text', tokens)]:
|
||||
outputs, _ = model.forward_multimodal(mod_name, mod_data[:3])
|
||||
pred_tokens = outputs['text_logits'].argmax(dim=-1)
|
||||
correct = (pred_tokens == tokens[:3]).float().mean().item()
|
||||
print(f" 输入 {mod_name:6s} → token 准确率: {correct:.3f}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='JspaceAI 多模态')
|
||||
parser.add_argument('--mode', type=str, default='train',
|
||||
choices=['train', 'live', 'eval'])
|
||||
parser.add_argument('--steps', type=int, default=200)
|
||||
parser.add_argument('--device', type=str, default='cpu')
|
||||
parser.add_argument('--outdir', type=str, default='outputs')
|
||||
args = parser.parse_args()
|
||||
|
||||
device = args.device
|
||||
if device == 'auto':
|
||||
if torch.cuda.is_available(): device = 'cuda'
|
||||
elif torch.backends.mps.is_available(): device = 'mps'
|
||||
else: device = 'cpu'
|
||||
|
||||
if args.mode == 'train':
|
||||
train_multimodal(args.steps, device, Path(args.outdir))
|
||||
elif args.mode == 'live':
|
||||
live_demo(args.steps, device, Path(args.outdir))
|
||||
elif args.mode == 'eval':
|
||||
eval_multimodal(device, Path(args.outdir))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
438
main_v2.py
Normal file
438
main_v2.py
Normal file
@@ -0,0 +1,438 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
JspaceAI v2 —— 基于 Anthropic J-space 论文优化版
|
||||
|
||||
新增功能:
|
||||
1. J-lens:观测模型内部"想法"(每个 ODE 子步的 workspace 读出)
|
||||
2. Directed Modulation:指令模型"想某概念",验证 workspace 可被 top-down 调制
|
||||
3. Selectivity 验证:ablate workspace,看是否只影响灵活推理
|
||||
4. W 轨迹记录:forward 时记录每个子步的 w,用于 J-lens 训练和可视化
|
||||
|
||||
运行:
|
||||
python main_v2.py
|
||||
python main_v2.py --steps 200 --device mps
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from pathlib import Path
|
||||
|
||||
from jspaceai import (
|
||||
LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
|
||||
CharTokenizer, load_shakespeare,
|
||||
JLensConfig, JLensSuite,
|
||||
WorkspaceAblator, DirectedModulation,
|
||||
)
|
||||
|
||||
|
||||
def count_params(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
def run_experiment(n_steps: int, device: str, outdir: Path):
|
||||
outdir.mkdir(parents=True, exist_ok=True)
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
# 1. 数据
|
||||
text = load_shakespeare()
|
||||
tokenizer = CharTokenizer.from_text(text)
|
||||
print(f"文本: {len(text)} 字符, 词汇表: {tokenizer.vocab_size}")
|
||||
|
||||
# 2. 模型
|
||||
config = LanguageConfig(
|
||||
vocab_size=tokenizer.vocab_size,
|
||||
embed_dim=16, input_dim=8,
|
||||
workspace_dim=32, expert_dim=16,
|
||||
num_experts=5, num_wells=4,
|
||||
ode_steps=4, dt=0.1, tau_w=0.3,
|
||||
jacobian_sparsity=8, noise_std=0.005,
|
||||
)
|
||||
model = JSpaceLanguageModel(config)
|
||||
print(f"模型参数: {count_params(model):,}")
|
||||
|
||||
# 3. J-lens 套件
|
||||
jlens_config = JLensConfig(
|
||||
n_substeps=config.ode_steps,
|
||||
workspace_dim=config.workspace_dim,
|
||||
vocab_size=config.vocab_size,
|
||||
)
|
||||
jlens_suite = JLensSuite(jlens_config).to(device)
|
||||
|
||||
# 4. 基础进化训练
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 1:基础进化训练")
|
||||
print("=" * 70)
|
||||
|
||||
trainer = EvolutionTrainer(model, config, lr=5e-3, ewc_lambda=0.05, device=device)
|
||||
chunks = [text[i:i+200] for i in range(0, len(text), 200)]
|
||||
trainer.evolve(
|
||||
chunks, tokenizer,
|
||||
seq_len=48, batch_size=4,
|
||||
consolidate_every=30, generate_every=50,
|
||||
max_steps=n_steps, prompt_text="To be",
|
||||
)
|
||||
|
||||
# 5. 训练 J-lens
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 2:训练 J-lens 探针")
|
||||
print("=" * 70)
|
||||
|
||||
jlens_optimizer = torch.optim.Adam(jlens_suite.parameters(), lr=3e-3)
|
||||
model.eval()
|
||||
all_tokens = tokenizer.encode(text)
|
||||
|
||||
for jlens_epoch in range(60):
|
||||
batch_tokens = []
|
||||
for _ in range(8):
|
||||
start = np.random.randint(0, len(all_tokens) - 64)
|
||||
seq = all_tokens[start:start+48]
|
||||
batch_tokens.append(seq)
|
||||
token_seq = torch.tensor(batch_tokens, dtype=torch.long).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
logits, info = model(token_seq, record_trajectory=True)
|
||||
|
||||
if 'w_trajectory' in info:
|
||||
w_traj = info['w_trajectory'] # (batch, T, n_substeps, workspace_dim)
|
||||
# target: 每个位置的下一个 token
|
||||
targets = token_seq[:, 1:] # (batch, T-1)
|
||||
|
||||
jlens_optimizer.zero_grad()
|
||||
total_loss = 0
|
||||
for substep in range(jlens_config.n_substeps):
|
||||
# (batch, T-1, workspace_dim) → 预测 targets
|
||||
w_sub = w_traj[:, :-1, substep, :] # (batch, T-1, workspace_dim)
|
||||
pred = jlens_suite.probes[substep](w_sub) # (batch, T-1, vocab)
|
||||
loss = F.cross_entropy(
|
||||
pred.reshape(-1, config.vocab_size),
|
||||
targets.reshape(-1),
|
||||
)
|
||||
total_loss += loss
|
||||
|
||||
total_loss.backward()
|
||||
jlens_optimizer.step()
|
||||
|
||||
if (jlens_epoch + 1) % 15 == 0:
|
||||
print(f" J-lens epoch {jlens_epoch+1}/60, loss={total_loss.item():.4f}")
|
||||
|
||||
model.train()
|
||||
|
||||
# 6. J-lens 观测
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 3:J-lens 观测——模型在想什么")
|
||||
print("=" * 70)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
prompt = "To be"
|
||||
prompt_ids = tokenizer.encode(prompt)
|
||||
token_tensor = torch.tensor([prompt_ids], dtype=torch.long).to(device)
|
||||
_, info = model(token_tensor, record_trajectory=True)
|
||||
|
||||
if 'w_trajectory' in info:
|
||||
w_traj = info['w_trajectory'][0]
|
||||
T, n_sub, _ = w_traj.shape
|
||||
sub_labels = ['sensory', 'workspace', 'workspace', 'motor']
|
||||
|
||||
print(f"\n提示: '{prompt}'")
|
||||
for t in range(min(T, 6)):
|
||||
char = prompt[t] if t < len(prompt) else '?'
|
||||
print(f" pos {t} ('{char}'):")
|
||||
for s in range(n_sub):
|
||||
w = w_traj[t, s].unsqueeze(0)
|
||||
probe = jlens_suite.probes[s]
|
||||
logits_s = probe(w)
|
||||
probs = F.softmax(logits_s, dim=-1)
|
||||
topk_probs, topk_idx = probs[0].topk(5)
|
||||
tokens_list = [tokenizer.idx_to_char.get(i.item(), '?') for i in topk_idx]
|
||||
probs_str = [f"{p:.2f}" for p in topk_probs.tolist()]
|
||||
pairs = ' '.join(f"{repr(tc)}({pr})" for tc, pr in zip(tokens_list, probs_str))
|
||||
label = sub_labels[s] if s < len(sub_labels) else f's{s}'
|
||||
print(f" 子步{s} ({label}): {pairs}")
|
||||
|
||||
# 7. Directed Modulation
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 4:Directed Modulation")
|
||||
print("=" * 70)
|
||||
|
||||
modulation = DirectedModulation(model, jlens_suite)
|
||||
test_concepts = ['R', 'd', 'o', ' ']
|
||||
|
||||
mod_results = []
|
||||
model.eval()
|
||||
for concept_char in test_concepts:
|
||||
if concept_char not in tokenizer.char_to_idx:
|
||||
continue
|
||||
concept_id = tokenizer.char_to_idx[concept_char]
|
||||
|
||||
with torch.no_grad():
|
||||
# 正常生成
|
||||
prompt_ids = list(tokenizer.encode("To be"))
|
||||
state = model.init_state(1, device)
|
||||
for tok in prompt_ids:
|
||||
state, _, _, _ = model.step(state, torch.tensor([tok], device=device))
|
||||
|
||||
normal_gen = []
|
||||
s = state
|
||||
last_tok = prompt_ids[-1]
|
||||
for _ in range(25):
|
||||
s, logits, _, _ = model.step(s, torch.tensor([last_tok], device=device))
|
||||
next_tok = logits[0].argmax().item()
|
||||
normal_gen.append(next_tok)
|
||||
last_tok = next_tok
|
||||
|
||||
# 调制生成
|
||||
mod_gen = []
|
||||
prompt_ids2 = list(tokenizer.encode("To be"))
|
||||
s = model.init_state(1, device)
|
||||
for tok in prompt_ids2:
|
||||
s, _, _, _ = model.step(s, torch.tensor([tok], device=device))
|
||||
|
||||
last_tok = prompt_ids2[-1]
|
||||
for _ in range(25):
|
||||
s = modulation.modulate_state(s, concept_id, strength=3.0)
|
||||
s, logits, _, _ = model.step(s, torch.tensor([last_tok], device=device))
|
||||
next_tok = logits[0].argmax().item()
|
||||
mod_gen.append(next_tok)
|
||||
last_tok = next_tok
|
||||
|
||||
normal_str = tokenizer.decode(normal_gen)
|
||||
mod_str = tokenizer.decode(mod_gen)
|
||||
mod_results.append((concept_char, normal_str, mod_str))
|
||||
print(f"\n 注入 '{concept_char}':")
|
||||
print(f" 正常: {repr(normal_str[:30])}")
|
||||
print(f" 调制: {repr(mod_str[:30])}")
|
||||
|
||||
# 8. Selectivity 验证——对比"自动任务"vs"需要 workspace 的任务"
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 5:Selectivity 验证")
|
||||
print("=" * 70)
|
||||
|
||||
# 任务 A: 简单续写(自动任务,应该 ablate 不影响)
|
||||
test_seqs = []
|
||||
for _ in range(8):
|
||||
start = np.random.randint(0, len(all_tokens) - 64)
|
||||
test_seqs.append(all_tokens[start:start+48])
|
||||
test_tensor = torch.tensor(test_seqs, dtype=torch.long).to(device)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
# 正常 forward
|
||||
normal_logits, _ = model(test_tensor)
|
||||
normal_pred = normal_logits[:, :-1].argmax(dim=-1)
|
||||
targets = test_tensor[:, 1:]
|
||||
normal_acc = (normal_pred == targets).float().mean().item()
|
||||
|
||||
# Ablate workspace forward
|
||||
ablate_state = model.init_state(test_tensor.shape[0], device)
|
||||
ablate_preds = []
|
||||
for t in range(test_tensor.shape[1] - 1):
|
||||
ablate_state, logits, _, _ = model.step(ablate_state, test_tensor[:, t])
|
||||
w = ablate_state['w']
|
||||
k = 5 # ablate top-5 J-lens 方向
|
||||
for b in range(w.shape[0]):
|
||||
probe = jlens_suite.probes[2]
|
||||
w_logits = probe(w[b:b+1])
|
||||
topk_vals, topk_idx = w_logits[0].topk(k)
|
||||
for idx in topk_idx:
|
||||
d = probe.lens.weight[idx]
|
||||
d_norm = d / (d.norm() + 1e-8)
|
||||
w[b] = w[b] - (w[b] @ d_norm) * d_norm
|
||||
ablate_state['w'] = w
|
||||
ablate_preds.append(logits.argmax(dim=-1))
|
||||
|
||||
ablate_preds = torch.stack(ablate_preds, dim=1)
|
||||
ablate_acc = (ablate_preds == targets).float().mean().item()
|
||||
|
||||
# 任务 B: 长程记忆(需要 workspace 持续装载信息)
|
||||
# 构造序列:前半段是"key",后半段需要回忆 key 的特征
|
||||
# 简化版:序列 [A, B, C, ..., A, ?] —— 第二次出现 A 后预测下一个
|
||||
# 对字符级,我们看"重复字符"任务:序列里某个字符重复出现,模型要"记住"它
|
||||
memory_seqs = []
|
||||
for _ in range(8):
|
||||
# 构造:随机字符 X 出现在位置 0,然后在位置 40 重复
|
||||
x = np.random.choice(all_tokens)
|
||||
seq = [np.random.choice(all_tokens) for _ in range(48)]
|
||||
seq[0] = x
|
||||
seq[40] = x # 40 步后重复
|
||||
memory_seqs.append(seq)
|
||||
memory_tensor = torch.tensor(memory_seqs, dtype=torch.long).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
# 正常
|
||||
mem_logits, _ = model(memory_tensor)
|
||||
mem_pred = mem_logits[:, :-1].argmax(dim=-1)
|
||||
mem_targets = memory_tensor[:, 1:]
|
||||
# 只看位置 40 之后(需要记忆的位置)
|
||||
mem_mask = torch.zeros_like(mem_targets, dtype=torch.bool)
|
||||
mem_mask[:, 40:] = True # 位置 40+ 需要"回忆"
|
||||
normal_mem_acc = (mem_pred[mem_mask] == mem_targets[mem_mask]).float().mean().item()
|
||||
|
||||
# Ablate
|
||||
abl_state = model.init_state(memory_tensor.shape[0], device)
|
||||
abl_preds = []
|
||||
for t in range(memory_tensor.shape[1] - 1):
|
||||
abl_state, logits, _, _ = model.step(abl_state, memory_tensor[:, t])
|
||||
w = abl_state['w']
|
||||
for b in range(w.shape[0]):
|
||||
probe = jlens_suite.probes[2]
|
||||
w_logits = probe(w[b:b+1])
|
||||
topk_vals, topk_idx = w_logits[0].topk(5)
|
||||
for idx in topk_idx:
|
||||
d = probe.lens.weight[idx]
|
||||
d_norm = d / (d.norm() + 1e-8)
|
||||
w[b] = w[b] - (w[b] @ d_norm) * d_norm
|
||||
abl_state['w'] = w
|
||||
abl_preds.append(logits.argmax(dim=-1))
|
||||
|
||||
abl_preds = torch.stack(abl_preds, dim=1)
|
||||
ablate_mem_acc = (abl_preds[mem_mask] == mem_targets[mem_mask]).float().mean().item()
|
||||
|
||||
print(f"\n 任务 A (简单续写 - 自动认知):")
|
||||
print(f" 正常: {normal_acc:.3f} Ablate: {ablate_acc:.3f} 下降: {normal_acc - ablate_acc:.3f}")
|
||||
print(f"\n 任务 B (长程记忆 - 需要 workspace):")
|
||||
print(f" 正常: {normal_mem_acc:.3f} Ablate: {ablate_mem_acc:.3f} 下降: {normal_mem_acc - ablate_mem_acc:.3f}")
|
||||
print(f"\n 解读: 任务 B 下降应大于任务 A —— workspace 对长程记忆更关键")
|
||||
|
||||
# 9. 可视化
|
||||
print("\n" + "=" * 70)
|
||||
print("阶段 6:可视化")
|
||||
print("=" * 70)
|
||||
|
||||
fig, axes = plt.subplots(2, 3, figsize=(18, 10))
|
||||
|
||||
ax = axes[0, 0]
|
||||
losses = [h['loss'] for h in trainer.history]
|
||||
ax.plot(losses, alpha=0.3, linewidth=0.5, color='blue')
|
||||
if len(losses) > 10:
|
||||
smoothed = np.convolve(losses, np.ones(10)/10, mode='valid')
|
||||
ax.plot(smoothed, linewidth=2, color='blue')
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('Loss')
|
||||
ax.set_title('Evolution Loss')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
ax = axes[0, 1]
|
||||
alpha_history = np.array([h['alpha_mean'] for h in trainer.history])
|
||||
for i in range(config.num_experts):
|
||||
ax.plot(alpha_history[:, i], label=f'Expert {i}', alpha=0.8)
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('Usage')
|
||||
ax.set_title('Expert Usage')
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
ax = axes[0, 2]
|
||||
w_norms = [h['w_norm_mean'] for h in trainer.history]
|
||||
ax.plot(w_norms, linewidth=1.5, color='green')
|
||||
ax.set_xlabel('Step')
|
||||
ax.set_ylabel('||w||')
|
||||
ax.set_title('Workspace Norm')
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
# J-lens 热力图
|
||||
ax = axes[1, 0]
|
||||
if 'w_trajectory' in info:
|
||||
w_traj = info['w_trajectory'][0]
|
||||
T, n_sub, _ = w_traj.shape
|
||||
sub_labels = ['sensory', 'workspace', 'workspace', 'motor']
|
||||
for s in range(min(n_sub, 4)):
|
||||
chars_row = []
|
||||
for t in range(T):
|
||||
w = w_traj[t, s].unsqueeze(0)
|
||||
probe = jlens_suite.probes[s]
|
||||
logits_s = probe(w)
|
||||
top1 = logits_s.argmax().item()
|
||||
chars_row.append(tokenizer.idx_to_char.get(top1, '?'))
|
||||
chars = ''.join(chars_row)
|
||||
label = sub_labels[s] if s < len(sub_labels) else f's{s}'
|
||||
ax.text(0.05, 0.95 - s*0.2, f"{label}: {chars}",
|
||||
transform=ax.transAxes, fontsize=9, fontfamily='monospace',
|
||||
verticalalignment='top')
|
||||
ax.set_title('J-lens Top-1 (position x substep)')
|
||||
ax.axis('off')
|
||||
|
||||
# Modulation 对比
|
||||
ax = axes[1, 1]
|
||||
ax.axis('off')
|
||||
mod_text = "Directed Modulation:\n\n"
|
||||
for concept_char, normal_str, mod_str in mod_results[:3]:
|
||||
mod_text += f"Inject '{concept_char}':\n"
|
||||
mod_text += f" N: {repr(normal_str[:20])}\n"
|
||||
mod_text += f" M: {repr(mod_str[:20])}\n\n"
|
||||
ax.text(0.05, 0.95, mod_text, transform=ax.transAxes, fontsize=8,
|
||||
verticalalignment='top', fontfamily='monospace',
|
||||
bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8))
|
||||
ax.set_title('Directed Modulation')
|
||||
|
||||
# Selectivity
|
||||
ax = axes[1, 2]
|
||||
tasks = ['Auto\n(continuation)', 'Memory\n(long-range)']
|
||||
normal_vals = [normal_acc, normal_mem_acc]
|
||||
ablate_vals = [ablate_acc, ablate_mem_acc]
|
||||
x = np.arange(len(tasks))
|
||||
width = 0.35
|
||||
ax.bar(x - width/2, normal_vals, width, label='Normal', color='green', alpha=0.7)
|
||||
ax.bar(x + width/2, ablate_vals, width, label='Ablated', color='red', alpha=0.7)
|
||||
ax.set_ylabel('Accuracy')
|
||||
ax.set_title('Selectivity: Workspace Ablation')
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(tasks)
|
||||
ax.legend()
|
||||
for i, (n, a) in enumerate(zip(normal_vals, ablate_vals)):
|
||||
ax.text(i - width/2, n + 0.005, f'{n:.3f}', ha='center', fontsize=8)
|
||||
ax.text(i + width/2, a + 0.005, f'{a:.3f}', ha='center', fontsize=8)
|
||||
|
||||
plt.tight_layout()
|
||||
fig_path = outdir / 'experiment_v2.png'
|
||||
plt.savefig(fig_path, dpi=150, bbox_inches='tight')
|
||||
print(f"可视化: {fig_path}")
|
||||
plt.close()
|
||||
|
||||
torch.save({
|
||||
'model': model.state_dict(),
|
||||
'jlens': jlens_suite.state_dict(),
|
||||
'config': config,
|
||||
'tokenizer_chars': tokenizer.chars,
|
||||
'selectivity': {
|
||||
'auto': {'normal': normal_acc, 'ablate': ablate_acc},
|
||||
'memory': {'normal': normal_mem_acc, 'ablate': ablate_mem_acc},
|
||||
},
|
||||
}, outdir / 'model_v2.pt')
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("完成")
|
||||
print("=" * 70)
|
||||
print(f"基础训练: {len(trainer.history)} 步")
|
||||
print(f"J-lens: {jlens_config.n_substeps} 个探针")
|
||||
print(f"Selectivity:")
|
||||
print(f" 自动任务下降: {normal_acc - ablate_acc:.3f}")
|
||||
print(f" 记忆任务下降: {normal_mem_acc - ablate_mem_acc:.3f}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='JspaceAI v2')
|
||||
parser.add_argument('--steps', type=int, default=200)
|
||||
parser.add_argument('--device', type=str, default='cpu')
|
||||
parser.add_argument('--outdir', type=str, default='outputs')
|
||||
args = parser.parse_args()
|
||||
|
||||
device = args.device
|
||||
if device == 'auto':
|
||||
if torch.cuda.is_available():
|
||||
device = 'cuda'
|
||||
elif torch.backends.mps.is_available():
|
||||
device = 'mps'
|
||||
else:
|
||||
device = 'cpu'
|
||||
|
||||
run_experiment(n_steps=args.steps, device=device, outdir=Path(args.outdir))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
9
requirements.txt
Normal file
9
requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
torch>=2.0
|
||||
numpy>=1.24,<2 # torch 2.2 不兼容 numpy 2.x
|
||||
matplotlib>=3.7
|
||||
opencv-python>=4.5
|
||||
sounddevice>=0.4
|
||||
soundfile>=0.12
|
||||
mss>=9.0
|
||||
pynput>=1.7
|
||||
av>=10.0
|
||||
Reference in New Issue
Block a user