具身智能（二）：机器人大脑与行为控制

# 多模态感知融合示例
import torch
import torch.nn as nn
from typing import Dict, Tuple

class MultimodalPerception(nn.Module):
    """具身智能的多模态感知融合模块"""
    
    def __init__(self, rgb_dim: int = 512, depth_dim: int = 256,
                 force_dim: int = 128, proprio_dim: int = 64,
                 fusion_dim: int = 512):
        super().__init__()
        
        # 各模态的编码器
        self.rgb_encoder = nn.Sequential(
            nn.Linear(rgb_dim, fusion_dim),
            nn.LayerNorm(fusion_dim),
            nn.GELU(),
        )
        self.depth_encoder = nn.Sequential(
            nn.Linear(depth_dim, fusion_dim),
            nn.LayerNorm(fusion_dim),
            nn.GELU(),
        )
        self.force_encoder = nn.Sequential(
            nn.Linear(force_dim, fusion_dim),
            nn.LayerNorm(fusion_dim),
            nn.GELU(),
        )
        self.proprio_encoder = nn.Sequential(
            nn.Linear(proprio_dim, fusion_dim),
            nn.LayerNorm(fusion_dim),
            nn.GELU(),
        )
        
        # 交叉注意力融合
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=fusion_dim,
            num_heads=8,
            batch_first=True,
        )
        self.fusion_mlp = nn.Sequential(
            nn.Linear(fusion_dim, fusion_dim * 2),
            nn.GELU(),
            nn.Linear(fusion_dim * 2, fusion_dim),
        )
    
    def forward(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        observations: {
            'rgb': (B, rgb_dim),
            'depth': (B, depth_dim),
            'force': (B, force_dim),
            'proprio': (B, proprio_dim),
        }
        """
        # 各模态编码
        rgb_feat = self.rgb_encoder(observations['rgb'])      # (B, D)
        depth_feat = self.depth_encoder(observations['depth']) # (B, D)
        force_feat = self.force_encoder(observations['force']) # (B, D)
        proprio_feat = self.proprio_encoder(observations['proprio'])  # (B, D)
        
        # 拼接为序列: (B, 4, D)
        features = torch.stack([rgb_feat, depth_feat, force_feat, proprio_feat], dim=1)
        
        # 交叉注意力融合（以 RGB 为 query）
        query = rgb_feat.unsqueeze(1)  # (B, 1, D)
        attended, _ = self.cross_attention(query, features, features)
        
        # 融合输出
        fused = self.fusion_mlp(attended.squeeze(1))  # (B, D)
        return fused

# 简化的世界模型实现
import torch
import torch.nn as nn
import torch.distributions as dist

class RSSM(nn.Module):
    """Recurrent State Space Model - 经典世界模型架构
    出自: Hafner et al., "Learning Latent Dynamics for Planning from Pixels"
    """
    
    def __init__(self, obs_dim: int = 64, action_dim: int = 10,
                 state_dim: int = 30, deterministic_dim: int = 200,
                 stochastic_dim: int = 30):
        super().__init__()
        self.state_dim = state_dim
        
        # GRU 核心（确定性状态转移）
        self.gru = nn.GRUCell(deterministic_dim, deterministic_dim)
        
        # 随机状态的后验编码器（观测 + 确定性状态 → 随机状态）
        self.posterior_encoder = nn.Sequential(
            nn.Linear(obs_dim + deterministic_dim, 128),
            nn.ELU(),
            nn.Linear(128, stochastic_dim * 2),  # mean, std
        )
        
        # 随机状态的先验（上一状态 + 动作 → 随机状态）
        self.prior_encoder = nn.Sequential(
            nn.Linear(state_dim + action_dim, 128),
            nn.ELU(),
            nn.Linear(128, stochastic_dim * 2),
        )
        
        # 观测解码器（状态 → 重建观测）
        self.obs_decoder = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ELU(),
            nn.Linear(128, obs_dim),
        )
        
        # 奖励预测
        self.reward_predictor = nn.Sequential(
            nn.Linear(state_dim + action_dim, 64),
            nn.ELU(),
            nn.Linear(64, 1),
        )
    
    def forward(self, observations: torch.Tensor, 
                actions: torch.Tensor) -> dict:
        """
        observations: (T, B, obs_dim)
        actions: (T, B, action_dim)
        """
        T, B, _ = observations.shape
        h = torch.zeros(B, 200)  # GRU 隐状态
        states = []
        posteriors = []
        priors = []
        
        for t in range(T):
            # 后验编码（有观测时）
            z_post = self._posterior(observations[t], h)
            
            # 先验预测（无观测时也可用）
            if t > 0:
                z_prior = self._prior(states[-1], actions[t-1])
                priors.append(z_prior)
            
            states.append(z_post)
            posteriors.append(z_post)
            
            # 更新 GRU
            concat = torch.cat([z_post, actions[t]], dim=-1)
            h = self.gru(concat, h)
        
        return {
            "states": torch.stack(states),
            "posteriors": posteriors,
            "priors": priors if priors else None,
        }
    
    def _posterior(self, obs: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
        inp = torch.cat([obs, h], dim=-1)
        mean, logstd = self.posterior_encoder(inp).chunk(2, dim=-1)
        return dist.Normal(mean, torch.exp(logstd)).rsample()
    
    def _prior(self, prev_state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        inp = torch.cat([prev_state, action], dim=-1)
        mean, logstd = self.prior_encoder(inp).chunk(2, dim=-1)
        return dist.Normal(mean, torch.exp(logstd)).rsample()

# VLA 模型推理示例（基于 OpenVLA 风格）
import torch
from transformers import AutoModelForVision2Seq, AutoProcessor
from PIL import Image
import numpy as np

class VLAPolicy:
    """简化的 VLA 策略模型"""
    
    def __init__(self, model_name: str = "openvla/openvla-7b"):
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.model = AutoModelForVision2Seq.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
    
    def predict_action(self, image: Image.Image, 
                       instruction: str,
                       state: np.ndarray = None) -> np.ndarray:
        """
        给定图像和语言指令，预测机器人动作
        
        Args:
            image: 当前相机观测
            instruction: 自然语言任务指令，如"拿起红色方块"
            state: 可选的机器人本体状态（关节角度等）
        
        Returns:
            action: 7 维动作向量 [x, y, z, roll, pitch, yaw, gripper]
        """
        # 构建提示词
        prompt = f"In: What action should the robot take to execute: '{instruction}'?
Out:"
        
        # 处理输入
        inputs = self.processor(
            text=prompt,
            images=image,
            return_tensors="pt",
        ).to(self.model.device)
        
        # 自回归生成动作 token
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=32,
            do_sample=False,
        )
        
        # 解码动作 token 为连续动作向量
        action_tokens = outputs[:, -7:]  # 取最后 7 个 token
        action = self._decode_tokens(action_tokens)
        return action
    
    def _decode_tokens(self, tokens: torch.Tensor) -> np.ndarray:
        """将离散的 action token 解码为连续动作值"""
        # OpenVLA 将每个动作维度离散化为 256 个 bin
        # 需要反归一化到实际动作范围
        action = (tokens.float() / 255.0) * 2.0 - 1.0  # [-1, 1]
        return action.squeeze().cpu().numpy()

# 使用示例
# policy = VLAPolicy()
# image = Image.open("robot_camera_view.jpg")
# action = policy.predict_action(image, "把杯子放到桌子上")
# robot.execute_action(action)

"""
具身智能入门：MuJoCo 机械臂抓取仿真
训练一个 RL 策略让机械臂抓取目标物体
"""
import mujoco
import mujoco_viewer
import numpy as np
from collections import deque

class RobotArmEnv:
    """简化的机械臂抓取环境"""
    
    def __init__(self, xml_path: str = "robot_arm.xml"):
        self.model = mujoco.MjModel.from_xml_path(xml_path)
        self.data = mujoco.MjData(self.model)
        self.viewer = None
        
        # 环境配置
        self.max_steps = 200
        self.action_scale = 0.05  # 动作幅度（米）
        
        # 状态历史（用于时间信息）
        self.state_history = deque(maxlen=4)
    
    def reset(self) -> np.ndarray:
        """重置环境"""
        mujoco.mj_resetData(self.model, self.data)
        
        # 随机化目标物体位置
        target_pos = self.data.geom("target").pos
        target_pos[:2] += np.random.uniform(-0.1, 0.1, 2)
        self.data.geom("target").pos = target_pos
        
        # 随机化初始关节角度
        self.data.qpos[:7] = np.random.uniform(
            self.model.jnt_range[:7, 0],
            self.model.jnt_range[:7, 1],
        )
        mujoco.mj_forward(self.model, self.data)
        
        self.state_history.clear()
        self.step_count = 0
        
        return self._get_observation()
    
    def step(self, action: np.ndarray) -> tuple:
        """执行一步动作"""
        # 动作空间：7 维 [dx, dy, dz, droll, dpitch, dyaw, gripper]
        clipped_action = np.clip(action, -1, 1) * self.action_scale
        
        # 设置关节速度目标
        self.data.ctrl[:6] = self.data.qpos[3:9] + clipped_action[:6]
        self.data.ctrl[6] = clipped_action[6]  # 夹爪开合
        
        # 物理步进
        mujoco.mj_step(self.model, self.data)
        self.step_count += 1
        
        # 计算奖励
        reward = self._compute_reward()
        done = self.step_count >= self.max_steps
        truncated = False
        
        return self._get_observation(), reward, done, truncated, {}
    
    def _get_observation(self) -> np.ndarray:
        """获取观测：关节状态 + 末端位置 + 目标位置"""
        qpos = self.data.qpos[:7]  # 关节角度
        qvel = self.data.qvel[:6]  # 关节速度
        end_effector = self.data.body("end_effector").xpos[:3]
        target_pos = self.data.geom("target").pos[:3]
        
        obs = np.concatenate([qpos, qvel, end_effector, target_pos])
        self.state_history.append(obs)
        
        # 拼接历史观测（提供时间信息）
        while len(self.state_history) < 4:
            self.state_history.append(np.zeros_like(obs))
        
        return np.concatenate(list(self.state_history))
    
    def _compute_reward(self) -> float:
        """奖励函数设计"""
        end_effector = self.data.body("end_effector").xpos[:3]
        target_pos = self.data.geom("target").pos[:3]
        
        # 距离奖励（负距离）
        distance = np.linalg.norm(end_effector - target_pos)
        dist_reward = -distance
        
        # 成功抓取奖励
        gripper_pos = self.data.ctrl[6]
        if distance < 0.02 and gripper_pos < 0:  # 靠近且夹爪闭合
            return 100.0
        
        return dist_reward

# RL 训练（简化的 SAC 算法）
from stable_baselines3 import SAC
from stable_baselines3.common.env_util import make_vec_env

def train_robot_arm():
    env = make_vec_env(lambda: RobotArmEnv(), n_envs=4)
    
    model = SAC(
        "MultiInputPolicy",
        env,
        learning_rate=3e-4,
        buffer_size=100000,
        batch_size=256,
        gamma=0.99,
        verbose=1,
    )
    
    model.learn(total_timesteps=500000)
    model.save("robot_arm_sac")
    return model

# 训练完成后部署
# model = SAC.load("robot_arm_sac")
# env = RobotArmEnv()
# obs = env.reset()
# for _ in range(200):
#     action, _ = model.predict(obs, deterministic=True)
#     obs, reward, done, _, _ = env.step(action)
#     if done:
#         break