VLA 架构深度剖析：视觉语言动作模型与机器人操控

1
2017  - Attention Is All You Need (Transformer)
2

3
2021  - CLIP (Radford et al.) — 视觉语言预训练
4
2022  - RT-1 (Brohan et al.) — 机器人 Transformer (350K 演示数据)
5

6
2023  - RT-2 (Brohan et al.) — VLA 开山之作 ★
7
         PaLM-E (Driess et al.) — 具身多模态 LLM
8
         RoboGPT (Kumar et al.) — LLM 作为机器人控制器
9
         HiveFormer — 多机器人协作
10

11
2024  - OpenVLA (Zhao et al.) — 开源 VLA (7B, 97K 演示数据) ★
12
         π₀ (Blackwell et al.) — 动作先验 + 完整机器人策略
13
         Act (Zhang et al.) — 扩散策略 + VLM 编码器
14
         RDT (Peng et al.) — 扩散 Transformer + 物理引擎数据
15
         VoxPoser (Huang et al.) — LLM 编程机器人
16
         UniARM — 通用机械臂策略
17

18
2025  - π₀.5 (Physical Intelligence) — 双手多指操控
19
         OpenHands-VLA (Ji et al.) — 开放式手动物体操控
20
         Helix — 完整 7 轴全身 VLA (Physical Intelligence)
21
         GR00T (NVIDIA) — 机器人基础模型

1
class ActionSpace:
2
    """
3
    机器人动作空间的三种主要类型。
4
    """
5
    DISCRETE = "discrete"
6
    CONTINUOUS = "continuous"
7
    HYBRID = "hybrid"
8

9

10
@dataclass
11
class ActionSpec:
12
    """
13
    动作规范。
14
    """
15
    space: str
16
    action_dim: int
17
    frequency_hz: float
18
    chunk_size: int   # 每次预测的动作步数

1
class DiscreteActionTokenizer:
2
    """
3
    RT-2 的离散动作 tokenizer。
4
    把连续动作量化成 token ID。
5
    """
6
    def __init__(self, num_bins=2**15):
7
        self.num_bins = num_bins   # 32768 个离散 bin
8
        self.action_dim = 14       # 7 joint + 6 ee + 1 gripper
9

10
    def tokenize_action(self, action):
11
        """
12
        连续动作 → 离散 token ID。
13

14
        参数:
15
            action: (action_dim,) 连续动作向量
16
        返回:
17
            token_ids: (action_dim,) 离散 token ID 列表
18
        """
19
        # 每个维度独立量化
20
        token_ids = []
21
        for a in action:
22
            # 归一化到 [-1, 1] (数据集内)
23
            norm_a = (a - self.action_mean) / (self.action_std + 1e-8)
24
            # 映射到 [0, num_bins)
25
            bin_id = int((norm_a + 1) / 2 * (self.num_bins - 1))
26
            bin_id = max(0, min(self.num_bins - 1, bin_id))
27
            token_ids.append(bin_id)
28
        return token_ids
29

30
    def detokenize_action(self, token_ids):
31
        """
32
        离散 token ID → 连续动作。
33
        """
34
        actions = []
35
        for tid in token_ids:
36
            norm_a = (tid / (self.num_bins - 1)) * 2 - 1
37
            a = norm_a * self.action_std + self.action_mean
38
            actions.append(a)
39
        return torch.tensor(actions)
40

41
    def forward(self, action_sequence):
42
        """
43
        动作序列 → token 序列。
44
        用于训练时的标签构建。
45
        """
46
        tokens = []
47
        for step in action_sequence:
48
            tokens.extend(self.tokenize_action(step))
49
        return tokens  # 14 × T 个 token IDs

1
class ContinuousActionHead(nn.Module):
2
    """
3
    OpenVLA 的连续动作头。
4
    输出: 每个动作维度的均值和方差（用于采样）。
5
    """
6
    def __init__(self, llm_dim, action_dim, chunk_size=8):
7
        super().__init__()
8
        self.chunk_size = chunk_size
9
        self.action_dim = action_dim
10

11
        self.action_head = nn.Sequential(
12
            nn.Linear(llm_dim, llm_dim // 2),
13
            nn.ReLU(),
14
            nn.Linear(llm_dim // 2, chunk_size * action_dim * 2),  # mean + log_std
15
        )
16

17
    def forward(self, llm_output):
18
        """
19
        llm_output: (B, llm_dim)  LLM 的最后一层 hidden state
20
        返回:
21
            action_chunk: (B, chunk_size, action_dim) 预测的动作块
22
        """
23
        raw = self.action_head(llm_output)  # (B, chunk_size*action_dim*2)
24
        mean, log_std = raw.chunk(2, dim=-1)  # 各 (B, chunk_size*action_dim)
25

26
        mean   = mean.view(-1, self.chunk_size, self.action_dim)
27
        log_std = log_std.view(-1, self.chunk_size, self.action_dim)
28
        std    = (log_std / 2).exp() + 1e-6   # 确保标准差为正
29

30
        # 训练时：直接预测均值（简化版）
31
        action_chunk = torch.tanh(mean)  # 动作通常在 [-1, 1] 范围
32
        return action_chunk

1
@dataclass
2
class ActionChunkConfig:
3
    """动作分块配置。"""
4
    chunk_size: int = 8        # 每次预测 8 步
5
    execution_freq_hz: float = 50.0  # 机器人控制频率
6
    lookahead_seconds: float = chunk_size / execution_freq_hz  # 0.16 秒
7

8
    # 置信度衰减
9
    use_ema_decay: bool = True   # 越远的动作权重越低
10
    ema_gamma: float = 0.99
11

12
    # 安全约束
13
    clip_actions: bool = True
14
    action_min: float = -1.0
15
    action_max: float = 1.0
16

17

18
def predict_with_chunking(model, observation, config):
19
    """
20
    用动作分块进行推理。
21
    """
22
    action_chunk = model(observation)  # (chunk_size, action_dim)
23

24
    if config.use_ema_decay:
25
        # 对远处的动作施加置信度衰减
26
        decays = config.ema_gamma ** torch.arange(config.chunk_size)
27
        action_chunk = action_chunk * decays.view(-1, 1)
28

29
    # 执行第一步
30
    action = action_chunk[0].cpu().numpy()
31

32
    # 安全约束
33
    if config.clip_actions:
34
        action = np.clip(action, config.action_min, config.action_max)
35

36
    return action, action_chunk[1:]  # 返回剩余 chunk 用于下次执行

1
class ActionNormalizer:
2
    """
3
    动作归一化器。
4
    每个数据集的动作有不同的均值和方差。
5
    """
6
    def __init__(self, action_dim):
7
        self.action_dim = action_dim
8
        self.register_buffer("mean", torch.zeros(action_dim))
9
        self.register_buffer("std",  torch.ones(action_dim))
10

11
    def normalize(self, action):
12
        """连续动作 → [-1, 1] 范围。"""
13
        return (action - self.mean) / (self.std + 1e-8)
14

15
    def denormalize(self, action_normed):
16
        """[-1, 1] → 原始物理单位。"""
17
        return action_normed * self.std + self.mean
18

19
    def compute_stats(self, dataset):
20
        """从数据集计算归一化统计量。"""
21
        all_actions = []
22
        for traj in dataset.trajectories:
23
            all_actions.append(traj.actions)  # (T, action_dim)
24

25
        all_actions = torch.cat(all_actions, dim=0)  # (N_total, action_dim)
26
        self.mean = all_actions.mean(dim=0)
27
        self.std  = all_actions.std(dim=0)
28

29
    @classmethod
30
    def from_dataset_config(cls, dataset_name):
31
        """预定义数据集的统计量。"""
32
        stats = {
33
            "rt-1": {
34
                "action_dim": 7,   # 7 自由度臂
35
                "mean": torch.zeros(7),
36
                "std":  torch.ones(7),
37
            },
38
            "bridge": {
39
                "action_dim": 8,   # 7DOF + gripper
40
                "mean": torch.zeros(8),
41
                "std":  torch.ones(8),
42
            },
43
            "open-x-embodiment": {
44
                "action_dim": 14,  # 7 joint + 6 ee + 1 gripper
45
                "mean": torch.zeros(14),
46
                "std":  torch.ones(14),
47
            },
48
        }
49
        return stats.get(dataset_name)

1
┌─────────────────────────────────────────────────────────────┐
2
│  VLA = 视觉编码器 + LLM 骨干 + 动作头                       │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                             │
5
│  [图像/视频] → 视觉编码器 → 视觉 tokens                      │
6
│       ↓                                                       │
7
│  [语言指令] → LLM Tokenizer → 语言 tokens                    │
8
│       ↓                                                       │
9
│  [视觉 tokens + 语言 tokens] → LLM Transformer ★            │
10
│       ↓                                                       │
11
│  LLM 输出 hidden state → 动作头 → [动作块]                  │
12
│                                                             │
13
└─────────────────────────────────────────────────────────────┘
14

15
三种设计模式:
16
  A. VLA-as-VLM (RT-2):   把动作当作"超级 token"，用语言头预测
17
  B. VLM + Action Head (OpenVLA): 冻住 VLM，加专用动作头
18
  C. Diffusion Action Head (π₀):   动作头用扩散模型

1
class RT2LikeVLA(nn.Module):
2
    """
3
    RT-2 风格的 VLA（离散动作版）。
4
    """
5
    def __init__(self, llm_name="google/palm2-e"):
6
        super().__init__()
7
        # 1) 视觉编码器 (冻结，用于预训练)
8
        self.vision_encoder = AutoModel.from_pretrained("google/pali-x")
9
        # 2) LLM 骨干
10
        self.llm = AutoModelForCausalLM.from_pretrained(llm_name)
11
        # 3) 动作 tokenizer（与语言 tokenizer 共用）
12
        self.action_tokenizer = DiscreteActionTokenizer(num_bins=2**15)
13
        # 4) 动作嵌入
14
        self.action_embed = nn.Embedding(2**15, self.llm.config.hidden_size)
15

16
    def forward(self, image, instruction, action_sequence=None):
17
        """
18
        训练时的前向传播。
19

20
        参数:
21
            image:          (B, C, H, W) 视觉输入
22
            instruction:    str 或 token IDs 语言指令
23
            action_sequence: (B, T, action_dim) 动作序列（标签）
24
        """
25
        # 视觉编码
26
        vision_tokens = self.vision_encoder(image)  # (B, N_v, D_v)
27

28
        # 语言编码
29
        if isinstance(instruction, str):
30
            instruction_ids = self.tokenizer(instruction, return_tensors="pt").to(image.device)
31
        else:
32
            instruction_ids = instruction
33

34
        lang_embeds = self.llm.get_input_embeddings()(instruction_ids["input_ids"])
35

36
        # 融合
37
        multimodal_embeds = torch.cat([vision_tokens, lang_embeds], dim=1)
38

39
        if action_sequence is not None:
40
            # 训练: 预测动作 tokens
41
            action_tokens = self.action_tokenizer.forward(action_sequence)
42
            action_embeds = self.action_embed(torch.tensor(action_tokens).to(image.device))
43
            multimodal_embeds = torch.cat([multimodal_embeds, action_embeds], dim=1)
44

45
            labels = instruction_ids["input_ids"].clone()
46
            labels[labels == self.tokenizer.pad_token_id] = -100  # 忽略 padding
47

48
            outputs = self.llm(inputs_embeds=multimodal_embeds, labels=labels)
49
            return outputs
50

51
        else:
52
            # 推理: 自回归生成动作
53
            outputs = self.llm.generate(
54
                inputs_embeds=multimodal_embeds,
55
                max_new_tokens=max_action_tokens,
56
            )
57
            action_tokens = outputs[0, -num_action_tokens:]
58
            actions = self.action_tokenizer.detokenize(action_tokens)
59
            return actions

1
class OpenVLAConfig:
2
    """OpenVLA 配置。"""
3
    vision_encoder: str = "蒲公英/siglip-so400m-224"  # SigLIP-So400m
4
    llm_model:      str = "meta-llama/Meta-Llama-3-8B"
5
    action_dim:     int = 7           # 7 自由度
6
    chunk_size:     int = 8           # 每次预测 8 步
7
    norm_stats:     str = "openx_embod"  # 动作归一化统计
8

9

10
class OpenVLA(nn.Module):
11
    """
12
    OpenVLA 架构。
13
    论文: "OpenVLA: An Open-Source Vision-Language-Action Model"
14
    """
15

16
    def __init__(self, config: OpenVLAConfig):
17
        super().__init__()
18
        # 1) 视觉编码器 (冻结)
19
        self.vision_encoder = AutoModel.from_pretrained(config.vision_encoder)
20
        self.vision_encoder.requires_grad_(False)
21

22
        # 2) 模态投影器 (训练)
23
        vision_dim = self.vision_encoder.config.hidden_size
24
        lang_dim  = 4096  # Llama-3 8B
25
        self.projector = nn.Sequential(
26
            nn.Linear(vision_dim, lang_dim),
27
            nn.GELU(),
28
            nn.Linear(lang_dim, lang_dim),
29
        )
30

31
        # 3) LLM (冻结)
32
        self.llm = AutoModelForCausalLM.from_pretrained(
33
            config.llm_model,
34
            torch_dtype=torch.float16,
35
            device_map="auto",
36
        )
37
        self.llm.requires_grad_(False)
38

39
        # 4) 动作头 (训练)
40
        self.action_head = nn.Sequential(
41
            nn.Linear(lang_dim, lang_dim // 2),
42
            nn.ReLU(),
43
            nn.Linear(lang_dim // 2, config.chunk_size * config.action_dim),
44
        )
45

46
        # 5) 动作归一化
47
        self.normalizer = ActionNormalizer(config.action_dim)
48
        self.normalizer.load_stats(config.norm_stats)
49

50
        # 6) 特殊 token
51
        self.tokenizer = AutoTokenizer.from_pretrained(config.llm_model)
52
        obs_tokens = ["[ OBS ]", "[ / OBS ]"]
53
        self.tokenizer.add_tokens(obs_tokens, special_tokens=True)
54
        self.obs_start = self.tokenizer("[ OBS ]", add_special_tokens=False).input_ids[0]
55
        self.obs_end   = self.tokenizer("[ / OBS ]", add_special_tokens=False).input_ids[0]
56

57
    def encode_observation(self, image):
58
        """编码图像观察。"""
59
        with torch.no_grad():
60
            vision_out = self.vision_encoder(image)
61
            vision_embeds = vision_out.last_hidden_state  # (B, N+1, D_v)
62
            vision_embeds = vision_embeds[:, 1:, :]       # 去掉 [CLS]
63

64
        return self.projector(vision_embeds)  # (B, N, D_l)
65

66
    def forward(self, image, instruction, action=None):
67
        """
68
        完整前向传播。
69
        """
70
        # 编码观察
71
        vision_embeds = self.encode_observation(image)  # (B, N, D_l)
72

73
        # 编码指令
74
        instr_ids = self.tokenizer(instruction, return_tensors="pt", padding=True,
75
                                   truncation=True, max_length=512).to(image.device)
76
        lang_embeds = self.llm.get_input_embeddings()(instr_ids["input_ids"])
77

78
        # 拼接: [OBS] + 视觉 tokens + [OBS] + 语言指令
79
        B = image.shape[0]
80
        obs_start_emb = self.llm.get_input_embeddings()(
81
            torch.tensor([self.obs_start]).to(image.device)
82
        ).expand(B, 1, -1)
83
        obs_end_emb = self.llm.get_input_embeddings()(
84
            torch.tensor([self.obs_end]).to(image.device)
85
        ).expand(B, 1, -1)
86

87
        multimodal_embeds = torch.cat([obs_start_emb, vision_embeds,
88
                                        obs_end_emb, lang_embeds], dim=1)
89

90
        # LLM 前向（冻结权重，但用于特征提取）
91
        with torch.no_grad():
92
            llm_out = self.llm(inputs_embeds=multimodal_embeds,
93
                               attention_mask=create_attention_mask(multimodal_embeds))
94

95
        # 动作预测（训练）
96
        last_hidden = llm_out.last_hidden_state[:, -1, :]  # (B, D_l)
97
        action_flat = self.action_head(last_hidden)       # (B, chunk*action_dim)
98
        action_chunk = action_flat.view(-1, self.chunk_size, self.action_dim)
99
        action_chunk = torch.tanh(action_chunk)            # 限制在 [-1, 1]
100

101
        if action is not None:
102
            # 训练模式: 计算 MSE 损失
103
            action_normed = self.normalizer.normalize(action)  # (B, T, action_dim)
104
            # 取前 chunk_size 步作为标签
105
            loss = F.mse_loss(action_chunk, action_normed[:, :self.chunk_size, :])
106
            return {"loss": loss, "action_pred": action_chunk}
107

108
        return action_chunk

1
class DiffusionActionPrior(nn.Module):
2
    """
3
    π₀ 的扩散动作先验。
4
    条件扩散模型: noise → action_chunk
5
    """
6
    def __init__(self, lang_dim, action_dim, chunk_size, hidden_dim=512):
7
        super().__init__()
8
        self.action_dim = action_dim
9
        self.chunk_size = chunk_size
10
        self.noise_steps = 100
11

12
        # 时间步嵌入
13
        self.time_embed = nn.Sequential(
14
            nn.Linear(1, hidden_dim),
15
            nn.SiLU(),
16
            nn.Linear(hidden_dim, hidden_dim),
17
        )
18

19
        # 条件: 语言 + 视觉
20
        self.cond_proj = nn.Linear(lang_dim, hidden_dim)
21

22
        # U-Net 风格的 denoiser
23
        self.net = nn.Sequential(
24
            nn.Linear(chunk_size * action_dim + hidden_dim, hidden_dim * 4),
25
            nn.SiLU(),
26
            nn.Linear(hidden_dim * 4, hidden_dim * 4),
27
            nn.SiLU(),
28
            nn.Linear(hidden_dim * 4, chunk_size * action_dim),
29
        )
30

31
        # 噪声调度
32
        self.register_buffer("betas", self.cosine_beta_schedule())
33

34
    def cosine_beta_schedule(self):
35
        """余弦噪声调度。"""
36
        steps = self.noise_steps + 1
37
        s = 0.008
38
        x = torch.linspace(0, self.noise_steps, steps)
39
        alphas_cumprod = torch.cos((x / self.noise_steps + s) / (1 + s) * torch.pi * 0.5) ** 2
40
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
41
        betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1]
42
        return torch.clamp(betas, 0.0001, 0.9999)
43

44
    def forward(self, x_t, t, cond):
45
        """
46
        Denoiser: 从噪声 x_t 中预测 noise (DDPM 目标)。
47

48
        参数:
49
            x_t:    (B, chunk_size, action_dim) 带噪声的动作
50
            t:      (B,)  时间步
51
            cond:   (B, lang_dim) 条件（语言+视觉）
52
        返回:
53
            noise_pred: (B, chunk_size, action_dim) 预测的噪声
54
        """
55
        t_embed = self.time_embed(t.float().unsqueeze(-1) / self.noise_steps)
56
        cond_embed = self.cond_proj(cond)
57

58
        x_flat = x_t.flatten(1)  # (B, chunk*action_dim)
59
        x_cat = torch.cat([x_flat, t_embed + cond_embed], dim=-1)  # (B, chunk*action_dim + hidden_dim)
60

61
        noise_pred = self.net(x_cat)
62
        noise_pred = noise_pred.view(-1, self.chunk_size, self.action_dim)
63
        return noise_pred
64

65
    @torch.no_grad()
66
    def sample(self, cond, num_inference_steps=50):
67
        """
68
        DDIM 采样：从条件生成动作块。
69

70
        参数:
71
            cond: (B, lang_dim) 条件
72
            num_inference_steps: 推理步数
73
        返回:
74
            action_chunk: (B, chunk_size, action_dim)
75
        """
76
        B = cond.shape[0]
77
        # 从纯噪声开始
78
        x = torch.randn(B, self.chunk_size, self.action_dim, device=cond.device)
79

80
        # DDIM 时间步
81
        step_ratio = self.noise_steps // num_inference_steps
82
        timesteps = (torch.arange(0, num_inference_steps) * step_ratio).flip(0)
83

84
        for i, t in enumerate(timesteps):
85
            t_tensor = torch.full((B,), t, device=cond.device, dtype=torch.long)
86

87
            # 预测噪声
88
            noise_pred = self.forward(x, t_tensor, cond)
89

90
            # DDIM 采样公式
91
            alpha_prod_t     = self.alphas_cumprod[t]
92
            alpha_prod_t_prev = self.alphas_cumprod[t - step_ratio] if t > 0 else torch.tensor(1.0)
93

94
            # 去噪
95
            x0_pred = (x - (1 - alpha_prod_t).sqrt() * noise_pred) / alpha_prod_t.sqrt()
96

97
            # 估计 x_{t-1}
98
            x = alpha_prod_t_prev.sqrt() * x0_pred + (1 - alpha_prod_t_prev).sqrt() * noise_pred
99

100
        return x  # (B, chunk_size, action_dim)
101

102

103
class PiZero(nn.Module):
104
    """
105
    π₀ 完整架构。
106
    """
107
    def __init__(self):
108
        super().__init__()
109
        # 1) VLM 骨干 (冻结)
110
        self.vlm = load_pretrained_vlm()  # e.g., CLIP ViT-L + Llama-3
111

112
        # 2) π₀ 扩散动作先验
113
        lang_dim = 4096
114
        action_dim = 7
115
        chunk_size = 16
116
        self.diffusion_prior = DiffusionActionPrior(lang_dim, action_dim, chunk_size)
117

118
    def forward(self, image, instruction, action=None):
119
        """
120
        完整前向。
121
        """
122
        # VLM 提取特征
123
        vision_embeds, lang_embeds = self.vlm(image, instruction)
124
        cond = vision_embeds + lang_embeds  # (B, lang_dim)
125

126
        if action is not None:
127
            # 训练: 加噪，计算噪声预测损失
128
            B, T, A = action.shape
129
            if T > self.diffusion_prior.chunk_size:
130
                action = action[:, :self.diffusion_prior.chunk_size]
131
            elif T < self.diffusion_prior.chunk_size:
132
                pad = torch.zeros(B, self.diffusion_prior.chunk_size - T, A, device=action.device)
133
                action = torch.cat([action, pad], dim=1)
134

135
            t = torch.randint(0, self.diffusion_prior.noise_steps, (B,), device=action.device)
136
            noise = torch.randn_like(action)
137
            alpha_prod = self.diffusion_prior.alphas_cumprod[t].view(B, 1, 1)
138
            noisy_action = alpha_prod.sqrt() * action + (1 - alpha_prod).sqrt() * noise
139

140
            noise_pred = self.diffusion_prior(noisy_action, t, cond)
141
            loss = F.mse_loss(noise_pred, noise)
142
            return {"loss": loss}
143

144
        else:
145
            # 推理: 扩散采样
146
            action_chunk = self.diffusion_prior.sample(cond)
147
            return action_chunk

1
┌─────────────────────────────────────────────────────────────┐
2
│  阶段 1: 视觉语言预训练 (与 VLM 相同)                        │
3
│  目标: 让视觉编码器和 LLM 学到强大的语义理解                  │
4
│  数据: 图文对、网页文本、数万亿 token                        │
5
│  LLM: 训练 🔥 | 视觉编码器: 训练 🔥                         │
6
├─────────────────────────────────────────────────────────────┤
7
│  阶段 2: 机器人动作预训练 (Imitation Learning)             │
8
│  目标: 让 VLM 的语义能力迁移到动作预测                       │
9
│  数据: 示教数据 (demonstration data)                        │
10
│        - Open X-Embodiment: 100+ 机器人, 100K+ episodes    │
11
│        - Bridge Dataset: 家庭场景 7DoF 臂                   │
12
│        - RT-1/RT-2 数据: Google 机器人                      │
13
│  LLM: 解冻 🔥 | 视觉编码器: 解冻 🔥 | 动作头: 训练 🔥       │
14
├─────────────────────────────────────────────────────────────┤
15
│  阶段 3: 任务特定微调 (可选)                                 │
16
│  目标: 在目标任务上进一步优化                                │
17
│  数据: 目标场景的少量示教数据                                │
18
│  方法: 行为克隆 / DAPO / PPO                               │
19
│  特点: 可以用 LoRA 高效微调                                 │
20
└─────────────────────────────────────────────────────────────┘

1
class RobotDataset(Dataset):
2
    """
3
    机器人示教数据集。
4
    """
5
    def __init__(self, dataset_paths, chunk_size=8):
6
        self.trajectories = []
7
        for path in dataset_paths:
8
            self.trajectories.extend(self.load(path))
9

10
        self.chunk_size = chunk_size
11
        self.normalizer = ActionNormalizer(action_dim)
12

13
    def load(self, path):
14
        """加载轨迹数据。"""
15
        if path.endswith(".hdf5"):
16
            return self.load_hdf5(path)
17
        elif path.endswith(".zarr"):
18
            return self.load_zarr(path)
19

20
    def __getitem__(self, idx):
21
        traj = self.trajectories[idx]
22
        # 随机采样一个起始位置
23
        T = len(traj["observations"])
24
        start = np.random.randint(0, max(1, T - self.chunk_size))
25
        end   = min(start + self.chunk_size, T)
26

27
        obs   = traj["observations"][start:end]      # (T, obs_dim)
28
        actions = traj["actions"][start:end]          # (T, action_dim)
29
        instruction = traj["language_instruction"]
30

31
        # 图像处理（可能多视角）
32
        images = [obs[i]["image"] for i in range(start, end)]
33
        # 多视角 concat 或选择主视角
34
        image = np.concatenate(images, axis=-1) if len(images) > 1 else images[0]
35

36
        return {
37
            "image":       self.preprocess_image(image),
38
            "instruction": instruction,
39
            "action":      torch.FloatTensor(actions),
40
        }
41

42
    def preprocess_image(self, image):
43
        # 调整大小到 [224, 224]
44
        # 归一化到 [0, 1]
45
        return torch.FloatTensor(image).permute(2, 0, 1) / 255.0
46

47

48
def train_openx_vla(model, dataset_paths, num_epochs=10, lr=1e-4):
49
    """
50
    在 Open X-Embodiment 数据集上训练 VLA。
51
    """
52
    dataset = RobotDataset(dataset_paths)
53
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=8)
54

55
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
56
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
57

58
    model.train()
59
    for epoch in range(num_epochs):
60
        epoch_loss = 0.0
61
        for batch in tqdm(dataloader):
62
            images       = batch["image"].to(device)  # (B, C, H, W)
63
            instructions = batch["instruction"]
64
            actions      = batch["action"].to(device)  # (B, T, A)
65

66
            optimizer.zero_grad()
67

68
            output = model(images, instructions, actions)
69
            loss   = output["loss"]
70

71
            loss.backward()
72
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
73
            optimizer.step()
74

75
            epoch_loss += loss.item()
76

77
        scheduler.step()
78
        print(f"Epoch {epoch}: loss={epoch_loss / len(dataloader):.4f}")
79

80
    # 保存
81
    torch.save(model.state_dict(), "openx_vla.pt")

1
def dapo_finetune(model, env, num_episodes=1000):
2
    """
3
    DAPO (Diffusion Augmented Policy Optimization) 微调。
4
    在模仿学习的基础上，用 RL 进一步优化策略。
5
    """
6
    from torch.distributions import Normal
7

8
    optimizer = torch.optim.AdamW(model.action_head.parameters(), lr=1e-5)
9

10
    for episode in range(num_episodes):
11
        # 收集轨迹
12
        obs     = env.reset()
13
        done    = False
14
        rewards = []
15
        states  = []
16

17
        while not done:
18
            with torch.no_grad():
19
                action = model.predict(obs)  # (action_dim,)
20

21
            next_obs, reward, done, info = env.step(action)
22

23
            rewards.append(reward)
24
            states.append((obs, action, reward, next_obs, done))
25

26
            obs = next_obs
27

28
        # 计算 GAE
29
        returns = compute_gae(rewards, gamma=0.99, lam=0.95)
30

31
        # 策略梯度更新
32
        optimizer.zero_grad()
33
        total_loss = 0.0
34

35
        for (obs, action, reward, next_obs, done), ret in zip(states, returns):
36
            output = model(obs["image"], obs["instruction"])
37
            action_pred = output["action_pred"][0]  # (action_dim,)
38

39
            # 负对数似然（假设高斯策略）
40
            log_prob = Normal(action_pred, std=0.1).log_prob(torch.FloatTensor(action)).mean()
41
            loss = -log_prob * ret  # 策略梯度
42

43
            total_loss += loss
44

45
        total_loss.backward()
46
        optimizer.step()
47

48
        if episode % 100 == 0:
49
            print(f"Episode {episode}: avg_reward={np.mean(rewards):.2f}")

1
class DomainRandomization:
2
    """
3
    领域随机化：增强 VLA 的跨机器人泛化能力。
4
    """
5
    def __init__(self):
6
        self.robot_configs = {
7
            "panda":       {"action_dim": 7, "limits": [(-2.9,2.9)]*7},
8
            " WidowX":     {"action_dim": 5, "limits": [(-1.6,1.6)]*5},
9
            "UR5e":        {"action_dim": 6, "limits": [(-np.pi,np.pi)]*6},
10
            "xArm7":       {"action_dim": 7, "limits": [(-2.0,2.0)]*7},
11
        }
12

13
    def sample_robot_config(self):
14
        """随机选择一个机器人配置。"""
15
        import random
16
        name = random.choice(list(self.robot_configs.keys()))
17
        return name, self.robot_configs[name]
18

19
    def adapt_action(self, action, source_robot, target_robot):
20
        """
21
        把源机器人的动作映射到目标机器人。
22
        """
23
        src_cfg = self.robot_configs[source_robot]
24
        tgt_cfg = self.robot_configs[target_robot]
25

26
        # 简单的维度映射（实际中更复杂）
27
        min_dim = min(src_cfg["action_dim"], tgt_cfg["action_dim"])
28
        adapted = action[:min_dim]
29

30
        # 重新缩放到目标机器人的范围
31
        for i in range(min_dim):
32
            src_lo, src_hi = src_cfg["limits"][i]
33
            tgt_lo, tgt_hi = tgt_cfg["limits"][i]
34
            normalized = (adapted[i] - src_lo) / (src_hi - src_lo)
35
            adapted[i] = normalized * (tgt_hi - tgt_lo) + tgt_lo
36

37
        return adapted

1
@dataclass
2
class ActionSpaceConfig:
3
    """
4
    VLA 动作空间的完整配置。
5
    """
6
    # 关节控制 vs 末端执行器控制
7
    control_mode: str = "joint"  # "joint" | "ee_pose" | "ee_delta"
8

9
    # 动作分块
10
    chunk_size: int = 8         # 每次预测的动作步数
11
    execution_freq_hz: float = 50.0
12

13
    # 动作表示
14
    representation: str = "continuous"  # "continuous" | "discrete" | "diffusion"
15

16
    # 连续动作的约束
17
    use_tanh_squash: bool = True    # tanh 把输出压缩到 [-1, 1]
18
    use_clipping: bool = True
19
    action_bound: float = 1.0
20

21
    # 速度/加速度/位置
22
    action_type: str = "delta"  # "position" | "delta" | "velocity"
23

24
    # 多臂协调 (π₀.5)
25
    num_arms: int = 1
26
    shared_base: bool = False   # 双臂是否共享底座
27

28

29
def apply_action_constraints(action, config: ActionSpaceConfig):
30
    """
31
    应用动作约束。
32
    """
33
    if config.use_tanh_squash:
34
        # tanh 压缩：输出在 (-1, 1)，然后缩放到实际范围
35
        action = config.action_bound * torch.tanh(action)
36

37
    if config.use_clipping:
38
        action = torch.clamp(action, -config.action_bound, config.action_bound)
39

40
    return action