🧠 VLA大模型深度解析¶

# RT-2 伪代码结构
class RT2(nn.Module):  # 继承nn.Module定义网络层
    def __init__(self, vlm_backbone, action_bins=256, action_dim=7):
        super().__init__()  # super()调用父类方法
        self.vlm = vlm_backbone  # PaLI-X / PaLM-E
        self.action_bins = action_bins
        self.action_dim = action_dim

        # 将连续动作空间均匀量化
        # 每个维度 [-1, 1] → 256个bin
        self.bin_edges = torch.linspace(-1, 1, action_bins + 1)

    def tokenize_action(self, continuous_action):
        """连续动作 → 离散token"""
        tokens = []
        for d in range(self.action_dim):
            # 找到最近的bin
            val = continuous_action[..., d]
            bin_idx = torch.bucketize(val, self.bin_edges[1:])
            bin_idx = bin_idx.clamp(0, self.action_bins - 1)
            tokens.append(bin_idx)
        return torch.stack(tokens, dim=-1)  # (B, 7)  # torch.stack沿新维度拼接张量

    def detokenize_action(self, tokens):
        """离散token → 连续动作"""
        bin_centers = (self.bin_edges[:-1] + self.bin_edges[1:]) / 2
        actions = []
        for d in range(self.action_dim):
            actions.append(bin_centers[tokens[..., d]])
        return torch.stack(actions, dim=-1)

    def forward(self, images, text_instruction):
        """
        images: (B, T, C, H, W) 历史图像序列
        text_instruction: 语言指令
        """
        # 拼接图像token + 文本token → VLM
        output_tokens = self.vlm(images, text_instruction)

        # 最后7个token即为动作
        action_tokens = output_tokens[:, -self.action_dim:]
        action = self.detokenize_action(action_tokens)

        return action  # (B, 7)

# OpenVLA 推理示例
from transformers import AutoModelForVision2Seq, AutoProcessor
from PIL import Image
import torch

class OpenVLAInference:
    """OpenVLA模型推理封装"""

    def __init__(self, model_path="openvla/openvla-7b"):
        self.processor = AutoProcessor.from_pretrained(
            model_path, trust_remote_code=True
        )
        self.model = AutoModelForVision2Seq.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        ).cuda()

    def predict_action(self, image: Image.Image, instruction: str):
        """
        输入一帧图像和语言指令，输出7DoF动作
        """
        # 构造prompt
        prompt = f"In: What action should the robot take to {instruction}?\nOut:"

        # 编码
        inputs = self.processor(prompt, image).to("cuda", dtype=torch.bfloat16)  # 移至GPU/CPU

        # 生成动作token
        action = self.model.predict_action(**inputs, unnorm_key="bridge_orig")

        # action: [Δx, Δy, Δz, Δroll, Δpitch, Δyaw, gripper_open]
        return action

    def run_episode(self, env, instruction, max_steps=200):
        """运行一个完整episode"""
        obs = env.reset()

        for step in range(max_steps):
            image = Image.fromarray(obs['image'])
            action = self.predict_action(image, instruction)

            obs, reward, done, info = env.step(action)

            if done:
                return info.get('success', False)

        return False

import torch
import torch.nn as nn

class FlowMatchingActionHead(nn.Module):
    """
    π0的核心：流匹配动作生成头

    基本思想:
    - 训练时: 学习从噪声分布到动作分布的向量场
    - 推理时: 从噪声出发，沿向量场积分得到动作

    对比DDPM: 流匹配只需要~10步去噪（DDPM需要50-100步）
    """

    def __init__(self, context_dim, action_dim, action_horizon=16):
        super().__init__()
        self.action_dim = action_dim        # 每步动作维度 (7)
        self.action_horizon = action_horizon # 预测未来16步动作
        flat_dim = action_dim * action_horizon

        # 去噪网络: 预测向量场 v(x_t, t, context)
        self.denoiser = nn.Sequential(
            nn.Linear(flat_dim + context_dim + 1, 512),
            nn.SiLU(),
            nn.Linear(512, 512),
            nn.SiLU(),
            nn.Linear(512, 512),
            nn.SiLU(),
            nn.Linear(512, flat_dim)
        )

    def forward(self, noisy_actions, timestep, context):
        """预测向量场 v"""
        t_embed = timestep.unsqueeze(-1)  # unsqueeze增加一个维度
        x = torch.cat([noisy_actions.flatten(-2), context, t_embed], dim=-1)  # torch.cat沿已有维度拼接张量
        return self.denoiser(x).view_as(noisy_actions)

    def training_loss(self, actions, context):
        """
        流匹配训练损失
        x_0: 噪声 (从标准正态分布采样)
        x_1: 真实动作
        x_t = (1-t)*x_0 + t*x_1  (线性插值)
        v_target = x_1 - x_0     (最优向量场)
        """
        B = actions.shape[0]
        x_1 = actions  # 真实动作
        x_0 = torch.randn_like(x_1)  # 噪声

        # 随机时间步
        t = torch.rand(B, device=actions.device)

        # 线性插值
        t_expand = t.view(B, 1, 1)  # 重塑张量形状
        x_t = (1 - t_expand) * x_0 + t_expand * x_1

        # 目标向量场
        v_target = x_1 - x_0

        # 预测向量场
        v_pred = self.forward(x_t, t, context)

        # MSE损失
        loss = ((v_pred - v_target) ** 2).mean()
        return loss

    @torch.no_grad()  # 禁用梯度计算，节省内存
    def sample(self, context, num_steps=10):
        """
        推理: ODE积分生成动作
        x_0 → x_1 (10步欧拉法)
        """
        B = context.shape[0]
        x = torch.randn(B, self.action_horizon, self.action_dim,
                         device=context.device)

        dt = 1.0 / num_steps
        for i in range(num_steps):
            t = torch.full((B,), i * dt, device=context.device)
            v = self.forward(x, t, context)
            x = x + v * dt

        return x  # (B, horizon, action_dim)

class DiffusionPolicy(nn.Module):
    """
    Diffusion Policy (Chi et al., RSS 2023)
    条件扩散模型用于行为克隆
    """

    def __init__(self, obs_encoder, noise_pred_net,
                 obs_horizon=2, pred_horizon=16, action_dim=7,
                 num_diffusion_steps=100):
        super().__init__()
        self.obs_encoder = obs_encoder
        self.noise_net = noise_pred_net
        self.obs_horizon = obs_horizon
        self.pred_horizon = pred_horizon
        self.action_dim = action_dim
        self.num_steps = num_diffusion_steps

        # DDPM噪声调度
        betas = torch.linspace(1e-4, 0.02, num_diffusion_steps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        self.register_buffer('alphas_cumprod', alphas_cumprod)
        self.register_buffer('sqrt_alphas_cumprod', alphas_cumprod.sqrt())
        self.register_buffer('sqrt_one_minus_alphas_cumprod',
                            (1 - alphas_cumprod).sqrt())
        self.register_buffer('betas', betas)

    def forward(self, obs_seq, actions=None):
        """
        obs_seq: (B, obs_horizon, obs_dim) 观察历史
        actions: (B, pred_horizon, action_dim) 训练用
        """
        # 编码观察
        obs_features = self.obs_encoder(obs_seq)  # (B, feat_dim)

        if actions is not None:
            return self._compute_loss(obs_features, actions)
        else:
            return self._sample(obs_features)

    def _compute_loss(self, obs_features, actions):
        """训练: 预测噪声"""
        B = actions.shape[0]
        noise = torch.randn_like(actions)

        # 随机时间步
        t = torch.randint(0, self.num_steps, (B,), device=actions.device)

        # 加噪
        # view(B,1,1)将每个样本的标量α值扩展为(B,1,1)，以广播方式与(B,预测步长,动作维度)的张量逐元素相乘
        sqrt_alpha = self.sqrt_alphas_cumprod[t].view(B, 1, 1)
        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t].view(B, 1, 1)
        noisy_actions = sqrt_alpha * actions + sqrt_one_minus_alpha * noise

        # 预测噪声
        noise_pred = self.noise_net(noisy_actions, t, obs_features)

        return ((noise_pred - noise) ** 2).mean()

    @torch.no_grad()
    def _sample(self, obs_features, use_ddim=True, ddim_steps=10):
        """推理: 从噪声去噪采样"""
        B = obs_features.shape[0]
        x = torch.randn(B, self.pred_horizon, self.action_dim,
                         device=obs_features.device)

        if use_ddim:
            # DDIM加速采样 (100步 → 10步)
            step_indices = torch.linspace(0, self.num_steps - 1, ddim_steps).long()
            for i in reversed(range(len(step_indices))):
                t = step_indices[i].expand(B).to(obs_features.device)
                noise_pred = self.noise_net(x, t, obs_features)

                alpha = self.alphas_cumprod[t].view(B, 1, 1)
                alpha_prev = self.alphas_cumprod[step_indices[i-1]].view(B, 1, 1) \
                             if i > 0 else torch.ones(B, 1, 1, device=x.device)

                x0_pred = (x - (1 - alpha).sqrt() * noise_pred) / alpha.sqrt()
                x = alpha_prev.sqrt() * x0_pred + \
                    (1 - alpha_prev).sqrt() * noise_pred

        return x  # (B, pred_horizon, action_dim)

from transformers import AutoModelForVision2Seq, AutoProcessor
from torch.utils.data import DataLoader
import torch

class VLAFineTuner:
    """VLA模型微调（以OpenVLA为例）"""

    def __init__(self, model_path, learning_rate=2e-5, use_lora=True):
        self.processor = AutoProcessor.from_pretrained(
            model_path, trust_remote_code=True
        )
        self.model = AutoModelForVision2Seq.from_pretrained(
            model_path, torch_dtype=torch.bfloat16, trust_remote_code=True
        ).cuda()

        if use_lora:
            self._apply_lora()

        self.optimizer = torch.optim.AdamW(
            # filter+lambda过滤出requires_grad=True的参数（排除冻结的视觉编码器），只优化LoRA等可训练部分
            filter(lambda p: p.requires_grad, self.model.parameters()),  # lambda匿名函数
            lr=learning_rate, weight_decay=0.01
        )

    def _apply_lora(self, rank=32, alpha=64):
        """对Llama部分应用LoRA（冻结视觉编码器）"""
        from peft import LoraConfig, get_peft_model

        lora_config = LoraConfig(
            r=rank,
            lora_alpha=alpha,
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
            lora_dropout=0.05,
            bias="none"
        )
        self.model = get_peft_model(self.model, lora_config)

        # 冻结视觉编码器
        for name, param in self.model.named_parameters():
            if "vision" in name.lower():
                param.requires_grad = False

        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.model.parameters())
        print(f"可训练参数: {trainable/1e6:.1f}M / {total/1e6:.1f}M "
              f"({100*trainable/total:.1f}%)")

    def train_epoch(self, dataloader):
        """训练一个epoch"""
        self.model.train()  # train()训练模式
        total_loss = 0

        for batch in dataloader:
            images = batch['images'].cuda()
            instructions = batch['instructions']
            target_actions = batch['actions'].cuda()

            # 前向传播
            outputs = self.model(
                images=images,
                text=instructions,
                labels=target_actions  # 动作token作为标签
            )

            loss = outputs.loss

            # 反向传播
            self.optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()  # 更新参数

            total_loss += loss.item()  # 将单元素张量转为Python数值

        return total_loss / len(dataloader)

# 微调超参数建议
# - 学习率: 2e-5 (全参) / 5e-5 (LoRA)
# - 批大小: 16-64 (取决于显存)
# - 训练轮数: 50-200 (取决于数据量)
# - 数据量: 至少50条有效轨迹（小规模）, 500+条（鲁棒）