XiuchengWu
300e86956b
- Add ConsensusSnapshot and ConsensusSlot classes for structured consensus state representation. - Introduce OnlineLanguageLearner for stateful online learning with experience replay and EWC regularization. - Create WorkspaceEvent and ActionEvent classes for serializable event types. - Develop InMemoryVectorMemoryStore for storing and querying workspace events. - Implement ActionPolicy to manage action selection and motor parameterization. - Establish a unified WorkspaceRuntime for managing the agent's lifecycle and interactions. - Enhance tests in test_smoke.py to cover new functionalities and ensure correctness.
470 lines
16 KiB
Python
470 lines
16 KiB
Python
"""
|
||
多模态感知-行动系统
|
||
|
||
原生支持图片、音频、视频、文本四种模态。
|
||
所有模态编码到统一的 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, torch.Tensor, list]:
|
||
"""
|
||
单步前向
|
||
|
||
Args:
|
||
state: {'w': ..., 'm': [...]}
|
||
x: (batch, input_dim) 已编码的输入（任意模态）
|
||
record_trajectory: 是否记录 w 轨迹
|
||
|
||
Returns:
|
||
new_state, alpha, 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, alpha, 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)
|
||
|
||
# 处理序列或单步
|
||
w_traj_all = []
|
||
alpha_last = None
|
||
if x.dim() == 2: # 单步 (batch, input_dim)
|
||
state, alpha_last, w_traj = self.step(state, x, record_trajectory)
|
||
if record_trajectory:
|
||
w_traj_all = w_traj
|
||
w = state['w']
|
||
else: # 序列 (batch, T, input_dim)
|
||
for t in range(x.shape[1]):
|
||
state, alpha_last, 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 alpha_last is not None:
|
||
info['alpha'] = alpha_last
|
||
if record_trajectory and w_traj_all:
|
||
info['w_trajectory'] = w_traj_all
|
||
|
||
return outputs, info
|