feat: 更新模块结构，添加训练相关功能和可恢复训练会话，优化模型生成逻辑

jspaceai/__init__.py

@@ -3,11 +3,12 @@ JspaceAI —— 全局工作空间 + J-space 广播的智慧系统
 
 模块：
     1. core.py: 核心架构（Expert + JSpaceWorkspace + JSpaceModel，含 RK4/LayerNorm/异构专家）
-    2. language_model.py: 语言版 + 自主进化
+    2. language_model.py: 语言版 JSpace 模型
     3. jlens.py: J-lens 可解释性工具
     4. multimodal.py: 多模态（图像/音频/视频/文本）
-    5. realtime.py: 实时 I/O（摄像头/麦克风/扬声器）
-    6. evolution.py: 自主进化训练器
+    5. policy.py / runtime.py: 动作策略与 workspace 主循环
+    6. events.py / memory.py: 可序列化事件与记忆接口
+    7. training.py / continual.py: 可恢复训练与在线学习
 """
 from .core import (
     Expert,
@@ -59,6 +60,13 @@ from .multimodal import (
 from .continual import (
     OnlineLanguageLearner,
 )
+from .training import (
+    LanguageTrainingConfig,
+    LanguageTrainingSession,
+    TokenBatchSampler,
+    expert_integration_mode,
+    save_language_checkpoint,
+)
 from .policy import (
     ACTION_LABELS,
     compose_action_params,
@@ -156,4 +164,7 @@ __all__ = [
     # 自主进化
     "EvolutionTrainer",
     "OnlineLanguageLearner",
+    "LanguageTrainingConfig", "LanguageTrainingSession",
+    "TokenBatchSampler", "expert_integration_mode",
+    "save_language_checkpoint",
 ]

jspaceai/embodied.py

@@ -1,21 +1,9 @@
 """
-输出执行器层 + 神经系统
+Embodied runtime adapter.
 
-对应人类神经系统的各部分：
-    - 大脑皮层: workspace w + 专家池（已在 multimodal.py）
-    - 小脑: 运动控制器（前向模型+逆模型，精细动作）
-    - 中枢神经: 动作调度器（反射弧+决策门控）
-    - 海马体: 外部情景记忆库
-    - 基底神经节: 动作价值学习（习惯化）
-    - 执行器: 鼠标控制 + 键盘输出 + 音频输出 + 屏幕绘制
-
-核心思想：输出和输入对称。
-    输入：摄像头/麦克风/屏幕/键盘/鼠标 → 编码 → workspace
-    输出：workspace → 解码 → 鼠标移动/键盘按键/音频播放/屏幕绘制
-
-workspace 是模态无关的"意图空间"。
-"想点击左上角"这个意图，在 workspace 里是一个向量，
-解码到鼠标控制器就是移动+点击，解码到键盘就是 Tab+Enter。
+This module owns local sensors/effectors and delegates decision making to
+ActionPolicy. The shared state remains the workspace vector plus serializable
+events, so memory and runtime orchestration can evolve independently.
 """
 from __future__ import annotations
 

jspaceai/language_model.py

@@ -199,6 +199,7 @@ class JSpaceLanguageModel(nn.Module):
             temperature: 采样温度
             top_k: top-k 采样
         """
+        was_training = self.training
         self.eval()
         device = next(self.parameters()).device
         state = self.init_state(1, device)
@@ -226,7 +227,10 @@ class JSpaceLanguageModel(nn.Module):
             generated.append(next_tok)
             tokens.append(next_tok)
 
-        self.train()
+        if was_training:
+            self.train()
+        else:
+            self.eval()
         return generated
 
 
@@ -285,11 +289,18 @@ class EWCOptimizer:
     """
 
     def __init__(self, model: nn.Module, lr: float = 1e-3,
-                 ewc_lambda: float = 1.0, max_grad_norm: float = 1.0):
+                 ewc_lambda: float = 1.0, max_grad_norm: float = 1.0,
+                 weight_decay: float = 0.0):
         self.model = model
-        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+        if weight_decay > 0:
+            self.optimizer = torch.optim.AdamW(
+                model.parameters(), lr=lr, weight_decay=weight_decay,
+            )
+        else:
+            self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
         self.ewc_lambda = ewc_lambda
         self.max_grad_norm = max_grad_norm
+        self.weight_decay = weight_decay
 
         # Fisher 信息和锚定参数
         self.fisher: dict[str, torch.Tensor] = {}
@@ -361,6 +372,25 @@ class EWCOptimizer:
         self.optimizer.step()
         return total_loss.item()
 
+    def state_dict(self) -> dict:
+        return {
+            "optimizer": self.optimizer.state_dict(),
+            "ewc_lambda": self.ewc_lambda,
+            "max_grad_norm": self.max_grad_norm,
+            "weight_decay": self.weight_decay,
+            "fisher": self.fisher,
+            "anchored_params": self.anchored_params,
+        }
+
+    def load_state_dict(self, state: dict):
+        if "optimizer" in state:
+            self.optimizer.load_state_dict(state["optimizer"])
+        self.ewc_lambda = state.get("ewc_lambda", self.ewc_lambda)
+        self.max_grad_norm = state.get("max_grad_norm", self.max_grad_norm)
+        self.weight_decay = state.get("weight_decay", self.weight_decay)
+        self.fisher = state.get("fisher", {})
+        self.anchored_params = state.get("anchored_params", {})
+
 
 class ExpertPlasticity:
     """

jspaceai/training.py

@@ -0,0 +1,311 @@
+"""
+Reusable training utilities for the JSpace language model.
+
+The goal is to keep training scalable without hard-wiring it to one script:
+sampling, validation, checkpointing, replay, and EWC all live behind a single
+session object that can be reused by CLI tools, tests, and future workers.
+"""
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from .language_data import CharTokenizer
+from .language_model import (
+    EWCOptimizer,
+    ExperienceReplay,
+    ExpertPlasticity,
+    JSpaceLanguageModel,
+    LanguageConfig,
+)
+
+
+@dataclass
+class LanguageTrainingConfig:
+    seq_len: int = 64
+    batch_size: int = 8
+    lr: float = 1e-3
+    weight_decay: float = 0.0
+    ewc_lambda: float = 0.05
+    max_grad_norm: float = 0.5
+    replay_capacity: int = 500
+    replay_batch_size: int = 4
+    replay_weight: float = 0.5
+    consolidate_every: int = 50
+    consolidate_samples: int = 10
+    validate_every: int = 50
+    validate_batches: int = 4
+    save_every: int = 100
+    train_fraction: float = 0.98
+    use_euler_during_train: bool = True
+
+
+@contextmanager
+def expert_integration_mode(model, use_rk4: bool):
+    """Temporarily set expert integration mode."""
+    original = [expert.use_rk4 for expert in model.experts]
+    for expert in model.experts:
+        expert.use_rk4 = use_rk4
+    try:
+        yield
+    finally:
+        for expert, enabled in zip(model.experts, original):
+            expert.use_rk4 = enabled
+
+
+class TokenBatchSampler:
+    """Random contiguous sampler over a token stream with a train/val split."""
+
+    def __init__(
+        self,
+        token_ids: list[int],
+        seq_len: int = 64,
+        train_fraction: float = 0.98,
+    ):
+        if not token_ids:
+            raise ValueError("token_ids must not be empty")
+        self.seq_len = seq_len
+        min_len = seq_len + 2
+        if len(token_ids) < min_len:
+            repeats = min_len // len(token_ids) + 1
+            token_ids = (token_ids * repeats)[:min_len]
+
+        split = int(len(token_ids) * train_fraction)
+        split = min(max(split, min_len), len(token_ids))
+        self.train_tokens = token_ids[:split]
+        self.val_tokens = token_ids[split:] if len(token_ids) - split >= min_len else token_ids[:split]
+
+    def sample(self, batch_size: int, split: str = "train", device: str = "cpu") -> torch.Tensor:
+        tokens = self.train_tokens if split == "train" else self.val_tokens
+        max_start = max(1, len(tokens) - self.seq_len - 1)
+        rows = []
+        for _ in range(batch_size):
+            start = torch.randint(0, max_start, (1,)).item()
+            rows.append(tokens[start:start + self.seq_len])
+        return torch.tensor(rows, dtype=torch.long, device=device)
+
+
+def save_language_checkpoint(
+    path: str | Path,
+    model: JSpaceLanguageModel,
+    config: LanguageConfig,
+    tokenizer: CharTokenizer,
+    trainer_state: dict | None = None,
+    metadata: dict | None = None,
+):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({
+        "model": model.state_dict(),
+        "config": config,
+        "tokenizer_chars": tokenizer.chars,
+        "trainer": trainer_state or {},
+        "metadata": metadata or {},
+    }, path)
+
+
+class LanguageTrainingSession:
+    """Stateful trainer for scalable language-model training."""
+
+    def __init__(
+        self,
+        model: JSpaceLanguageModel,
+        model_config: LanguageConfig,
+        tokenizer: CharTokenizer,
+        train_config: LanguageTrainingConfig | None = None,
+        device: str = "cpu",
+    ):
+        self.model = model.to(device)
+        self.model_config = model_config
+        self.tokenizer = tokenizer
+        self.train_config = train_config or LanguageTrainingConfig()
+        self.device = device
+        self.optimizer = EWCOptimizer(
+            self.model,
+            lr=self.train_config.lr,
+            ewc_lambda=self.train_config.ewc_lambda,
+            max_grad_norm=self.train_config.max_grad_norm,
+            weight_decay=self.train_config.weight_decay,
+        )
+        self.replay_buffer = ExperienceReplay(
+            capacity=self.train_config.replay_capacity,
+            seq_len=self.train_config.seq_len,
+        )
+        self.plasticity = ExpertPlasticity(num_experts=model_config.num_experts)
+        self.global_step = 0
+        self.history: list[dict] = []
+
+    def learn_batch(self, token_seq: torch.Tensor) -> dict:
+        token_seq = token_seq.to(self.device)
+        self.model.train()
+        logits, info = self.model(token_seq)
+        pred = logits[:, :-1]
+        target = token_seq[:, 1:]
+        loss = F.cross_entropy(
+            pred.reshape(-1, self.model_config.vocab_size),
+            target.reshape(-1),
+        )
+
+        replay_loss = torch.tensor(0.0, device=self.device)
+        replay_seq = self.replay_buffer.sample(self.train_config.replay_batch_size)
+        if replay_seq is not None:
+            replay_seq = replay_seq.to(self.device)
+            replay_logits, _ = self.model(replay_seq)
+            replay_pred = replay_logits[:, :-1]
+            replay_target = replay_seq[:, 1:]
+            replay_loss = F.cross_entropy(
+                replay_pred.reshape(-1, self.model_config.vocab_size),
+                replay_target.reshape(-1),
+            )
+
+        total_task_loss = loss + self.train_config.replay_weight * replay_loss
+        total_loss = self.optimizer.step(total_task_loss)
+
+        self.plasticity.update(info["alpha"].detach(), token_seq.detach())
+        self.replay_buffer.push(token_seq.detach().cpu())
+
+        return {
+            "loss": float(loss.item()),
+            "replay_loss": float(replay_loss.item()),
+            "total_loss": float(total_loss),
+            "w_norm_mean": float(info["w_norm"].mean().item()),
+            "expert_usage": self.plasticity.usage.tolist(),
+        }
+
+    @torch.no_grad()
+    def evaluate(self, sampler: TokenBatchSampler) -> float:
+        was_training = self.model.training
+        self.model.eval()
+        losses = []
+        for _ in range(max(1, self.train_config.validate_batches)):
+            token_seq = sampler.sample(
+                self.train_config.batch_size,
+                split="val",
+                device=self.device,
+            )
+            logits, _ = self.model(token_seq)
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, self.model_config.vocab_size),
+                token_seq[:, 1:].reshape(-1),
+            )
+            losses.append(loss.item())
+        if was_training:
+            self.model.train()
+        return float(sum(losses) / len(losses))
+
+    def fit_text(
+        self,
+        text: str,
+        max_steps: int,
+        checkpoint_path: str | Path | None = None,
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        return self.fit_tokens(
+            self.tokenizer.encode(text),
+            max_steps=max_steps,
+            checkpoint_path=checkpoint_path,
+            on_progress=on_progress,
+        )
+
+    def fit_tokens(
+        self,
+        token_ids: list[int],
+        max_steps: int,
+        checkpoint_path: str | Path | None = None,
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        sampler = TokenBatchSampler(
+            token_ids,
+            seq_len=self.train_config.seq_len,
+            train_fraction=self.train_config.train_fraction,
+        )
+        use_rk4 = not self.train_config.use_euler_during_train
+        with expert_integration_mode(self.model, use_rk4=use_rk4):
+            for _ in range(max_steps):
+                batch = sampler.sample(
+                    self.train_config.batch_size,
+                    split="train",
+                    device=self.device,
+                )
+                stats = self.learn_batch(batch)
+                self.global_step += 1
+                stats["step"] = self.global_step
+
+                if (
+                    self.train_config.validate_every > 0
+                    and self.global_step % self.train_config.validate_every == 0
+                ):
+                    stats["val_loss"] = self.evaluate(sampler)
+
+                if (
+                    self.train_config.consolidate_every > 0
+                    and self.global_step % self.train_config.consolidate_every == 0
+                ):
+                    self.optimizer.consolidate(
+                        batch.detach(),
+                        n_samples=self.train_config.consolidate_samples,
+                    )
+
+                self.history.append(stats)
+                if on_progress:
+                    on_progress(stats)
+
+                if (
+                    checkpoint_path is not None
+                    and self.train_config.save_every > 0
+                    and self.global_step % self.train_config.save_every == 0
+                ):
+                    self.save_checkpoint(checkpoint_path)
+
+        if checkpoint_path is not None:
+            self.save_checkpoint(checkpoint_path)
+        return self.history
+
+    def state_dict(self) -> dict:
+        return {
+            "global_step": self.global_step,
+            "optimizer": self.optimizer.state_dict(),
+            "replay_buffer": list(self.replay_buffer.buffer),
+            "plasticity": {
+                "usage": self.plasticity.usage,
+                "expert_specialization": self.plasticity.expert_specialization,
+            },
+            "train_config": asdict(self.train_config),
+            "history": self.history[-200:],
+        }
+
+    def load_state_dict(self, state: dict):
+        self.global_step = int(state.get("global_step", 0))
+        if "optimizer" in state:
+            self.optimizer.load_state_dict(state["optimizer"])
+        replay = state.get("replay_buffer", [])
+        self.replay_buffer.buffer.clear()
+        for seq in replay:
+            self.replay_buffer.push(seq)
+        plasticity = state.get("plasticity", {})
+        if "usage" in plasticity:
+            self.plasticity.usage = plasticity["usage"].detach().cpu()
+        if "expert_specialization" in plasticity:
+            self.plasticity.expert_specialization = plasticity["expert_specialization"]
+        self.history = list(state.get("history", []))
+
+    def save_checkpoint(self, path: str | Path, metadata: dict | None = None):
+        save_language_checkpoint(
+            path,
+            self.model,
+            self.model_config,
+            self.tokenizer,
+            trainer_state=self.state_dict(),
+            metadata=metadata,
+        )
+
+    def load_checkpoint_state(self, path: str | Path):
+        ckpt = torch.load(path, map_location=self.device, weights_only=False)
+        trainer_state = ckpt.get("trainer")
+        if trainer_state:
+            self.load_state_dict(trainer_state)