【AI】RMSS：循环状态空间模型

RSSM （Recurrent State Space Model，循环状态空间模型）是 Dreamer/DreamerV3 等世界模型的核心架构，由 Google DeepMind 的 Hafner 等人提出。它的核心思想是：不在高维像素空间直接预测未来，而是在压缩后的潜在空间（Latent Space）里学习环境的动态规律，让 AI 能够"闭眼想象"未来。

---

一、核心原理：双状态分离设计

RSSM 将世界的状态拆分为两个互补部分，这是它区别于普通 RNN/Transformer 的关键：

状态类型	作用	类比
确定性状态 (h)	用 GRU/RNN 记住历史轨迹，捕捉"过去发生了什么"	人的工作记忆：我知道方块在左上角，机械臂正在靠近
随机状态 (z)	用概率分布（通常是对角高斯）捕捉环境中的不确定性	人的直觉：推这个方块，摩擦力不确定，可能滑 3cm 也可能滑 5cm

完整状态表示 ：s_t = [h_t, z_t]（将两部分拼接）

数学流程（简化版）

1. 编码观测：
   embed_t = Encoder(o_t)           # CNN 把图像/传感器数据压缩为向量

2. 确定性路径（必须记住历史）：
   h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1})  
   # 上一时刻的确定性状态 + 随机状态 + 动作 → 当前确定性状态

3. 随机路径（捕捉不确定性）：
   # 训练时（有真实观测）：
   z_t ~ q(z_t | h_t, embed_t)      # 后验分布：结合历史+当前观测
   
   # 想象时（无真实观测）：
   ẑ_t ~ p(z_t | h_t)               # 先验分布：只凭历史预测未来

4. 解码/预测：
   ô_t, r̂_t, donê_t = Decoder(h_t, z_t)   # 重构画面、预测奖励、预测是否终止

---

二、关键机制：Latent Imagination（潜在空间想象）

这是 RSSM 最具革命性的能力。传统强化学习必须在真实环境中试错（费时间、费硬件、有危险），而 RSSM 可以在潜在空间中"做梦"：

想象流程

当前状态 (h_t, z_t)
    ↓
选择动作 a_t（来自 Actor 策略网络）
    ↓
RSSM 动态模型预测 → (h_{t+1}, z_{t+1})
    ↓
预测奖励 r̂_{t+1} 和终止标志 donê_{t+1}
    ↓
重复 N 步，生成一条完整的"想象轨迹"
    ↓
在想象轨迹上训练 Actor（策略）和 Critic（价值函数）

关键 ：想象时完全脱离真实环境，只依赖 p(z_t | h_t) 先验网络，因此可以并行生成数千条未来分支，像 AlphaGo 的蒙特卡洛树搜索一样在脑中推演。

---

三、举例说明：机器人推方块

场景设定

• 观测：64×64 像素的摄像头画面（桌子、红色方块、机械臂）
• 动作 ：二维向量 [dx, dy]（末端执行器的水平位移）
• 目标：把方块推到目标区域

RSSM 如何工作

Step 1: 编码现实

机器人看到第一帧画面 → CNN 编码为 embed_1 → RSSM 初始化 h_1（GRU 隐藏态）和 z_1（高斯采样）

Step 2: 学习动态（训练阶段）

• 机器人执行动作 a_1 = [+2cm, 0]，看到第二帧
• RSSM 比较：先验预测 （仅凭 h_1 猜 z_1） vs 后验 （结合真实画面 embed_2）
• 损失函数迫使两者接近，同时重构画面要准确、奖励预测要准
• 此时 h_2 记住了"我刚向右推了 2cm"

Step 3: 想象推演（规划阶段）

策略网络想测试"如果向左推会怎样"：

• 在潜在空间 rollout 15 步想象：
  • a = [-1cm, 0] → RSSM 预测 h'、z' → 解码器预测画面（方块左移）→ 预测奖励（-0.5，靠近桌边危险）
• a = [0, +1cm] → 想象 rollout → 预测奖励（+1.0，靠近目标区）
• 无需移动真实机械臂，0.1 秒内完成 1000 条想象轨迹的评估

Step 4: 执行最优动作

选择在想象中奖励最高的动作，发给真实机器人执行。

随机状态 (z) 的具体作用

• 如果桌面有微小倾斜（机器人不知道），每次推方块滑动距离有随机性
• z_t 的高斯分布方差会捕捉这种不确定性
• 想象时采样不同的 z，机器人能预见到"推这个动作可能产生多种结果"，从而选择更鲁棒的策略

---

四、应用场景

1. 机器人操作（Robot Manipulation）

• 代表：DreamerV3 在机械臂抓取、人形机器人平衡控制
• 价值：在仿真中想象数百万次抓取，迁移到真实机械臂时样本效率提升 10-100 倍

2. 自动驾驶（Autonomous Driving）

• 代表：CarDreamer（基于 DreamerV2/V3 的自动驾驶平台）
• 价值：车辆可以在脑中"预演"前方车辆急刹、行人横穿等危险场景，提前规划避险路径，而非仅靠反应式规则

3. 游戏 AI 与 NPC

• 代表：MineWorld（Minecraft 中的 RSSM 应用）
• 价值：NPC 不是按脚本行动，而是"想象"玩家可能的走位，提前埋伏或逃跑

4. 人形机器人全身控制

• 代表：Humanoid World Models (HWM)
• 价值：人形机器人从第一视角视频预测未来画面，学习走路时"如果左腿迈太大，身体会向右倾"，在潜在空间预演平衡调整

5. 样本高效强化学习（Sample-Efficient RL）

• 价值：Atari 游戏等环境中，DreamerV3 仅用 100 万帧（传统方法需数亿帧）即可达到人类水平，因为大部分训练在"想象"中完成

---

五、RSSM vs 其他架构

对比项	RSSM	纯 Transformer（如 GPT）	纯 Diffusion（如 Sora）
预测空间	压缩潜在空间（低维）	Token 空间	像素空间（高维）
长程规划	专为 RL 设计，支持闭环控制	擅长开环生成	擅长视频生成，不直接输出动作
不确定性建模	显式概率分布（z）	隐式在注意力中	隐式在去噪过程中
计算效率	高，可实时 rollout	注意力复杂度 O(N²)	低，需多步去噪
动作耦合	原生支持动作条件预测	需额外适配	需额外适配

---

六、一句话总结

RSSM 让 AI 拥有"想象力"：用 GRU 记住历史（确定性），用概率分布捕捉不确定性（随机性），在压缩的潜在空间中低成本推演未来，从而在不触碰真实世界的情况下学会复杂控制策略。