菜单
Train Vision-Language-Action models for robotic control through world model simulation without real-world interaction—using pixel-based world models aligned with VLA features to enable self-correction and robust policy optimization.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill wmpo-world-model-vla-training命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
正在显示 SKILL.md
| name | wmpo-world-model-vla-training |
| title | WMPO: World Model-based Policy Optimization for VLA Models |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2511.09515 |
| keywords | ["Reinforcement Learning","World Models","Vision-Language-Action","Embodied AI","Robot Learning"] |
| description | Train Vision-Language-Action models for robotic control through world model simulation without real-world interaction—using pixel-based world models aligned with VLA features to enable self-correction and robust policy optimization. |
Training robots with reinforcement learning typically requires expensive real-world trial-and-error. WMPO (World Model Policy Optimization) enables effective RL by learning a world model (pixel-based simulator) that accurately predicts VLA features trained on web-scale vision data. The robot then learns from simulated trajectories, avoiding real-world costs while gaining robust control behaviors.
The approach achieves improved sample efficiency, stronger overall performance, emergent self-correction abilities, and robust generalization—all critical for practical robotic manipulation.
WMPO addresses a fundamental challenge: vision-language-action models are strong visual reasoners but lack direct RL training signals. The system creates a closed-loop learning environment:
This approach replaces real-world risk with model-based simulation risk, dramatically reducing training cost.
Step 1: Train Pixel-Based World Model
Build a generative model predicting next frames from states and actions.
import torch
import torch.nn as nn
from torchvision.models import resnet50
class PixelWorldModel(nn.Module):
():
().__init__()
.action_dim = action_dim
.latent_dim = latent_dim
.encoder = nn.Sequential(
nn.Conv2d(input_channels, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.Conv2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.Conv2d(, latent_dim, kernel_size=, stride=, padding=),
nn.ReLU(),
)
.action_embed = nn.Sequential(
nn.Linear(action_dim, ),
nn.ReLU(),
nn.Linear(, latent_dim)
)
.transition = nn.ModuleList([
nn.ConvTranspose2d(latent_dim + latent_dim, latent_dim,
kernel_size=, padding=)
_ (num_blocks)
])
.decoder = nn.Sequential(
nn.ConvTranspose2d(latent_dim, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.ConvTranspose2d(, , kernel_size=, stride=, padding=),
nn.ReLU(),
nn.ConvTranspose2d(, input_channels, kernel_size=, stride=, padding=),
nn.Sigmoid()
)
() -> torch.Tensor:
latent = .encoder(current_frame)
action_emb = .action_embed(action)
action_emb = action_emb.unsqueeze(-).unsqueeze(-)
action_emb = action_emb.expand(-, -, latent.shape[-], latent.shape[-])
combined = torch.cat([latent, action_emb], dim=)
next_latent = combined
transition .transition:
next_latent = transition(next_latent)
next_latent = torch.relu(next_latent)
next_frame = .decoder(next_latent)
next_frame
():
optimizer = torch.optim.Adam(model.parameters(), lr=)
criterion = nn.MSELoss()
epoch (num_epochs):
total_loss =
batch_idx, (frame_t, action, frame_t1) (train_dataloader):
pred_frame = model(frame_t, action)
loss = criterion(pred_frame, frame_t1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / (train_dataloader)
()
Step 2: Align World Model Features with VLA Features
Ensure world model outputs match the feature space of the VLA model.
class FeatureAlignment(nn.Module):
"""
Aligns world model output features with VLA model features.
"""
def __init__(self, world_model_dim: int = 256,
vla_feature_dim: int = 512, hidden_dim: int = 512):
"""
Args:
world_model_dim: Dimension of world model outputs
vla_feature_dim: Dimension of VLA feature space
hidden_dim: Hidden dimension for projection
"""
super().__init__()
self.projection = nn.Sequential(
nn.Linear(world_model_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, vla_feature_dim)
)
def align(self, world_features: torch.Tensor) -> torch.Tensor:
"""
Project world model features to VLA feature space.
Args:
world_features: Features from world model [batch, world_model_dim]
Returns:
aligned_features: Features in VLA space [batch, vla_feature_dim]
"""
return self.projection(world_features)
def compute_alignment_loss(world_features: torch.Tensor,
vla_features: torch.Tensor,
alignment_module: FeatureAlignment,
temperature: float = 0.1) -> torch.Tensor:
"""
Contrastive loss ensuring alignment between world and VLA features.
Args:
world_features: Features from world model
vla_features: Features from VLA model (from real observations)
alignment_module: Feature alignment network
temperature: Contrastive temperature
Returns:
alignment_loss: Loss value
"""
# Project world features
aligned = alignment_module.align(world_features)
# Normalize
aligned_norm = F.normalize(aligned, dim=-1)
vla_norm = F.normalize(vla_features, dim=-1)
# Cosine similarity (contrastive learning)
similarity = torch.mm(aligned_norm, vla_norm.t()) / temperature
# Diagonal should be high (matching features), off-diagonal low
labels = torch.arange(aligned_norm.shape[0], device=aligned_norm.device)
loss = F.cross_entropy(similarity, labels) + \
F.cross_entropy(similarity.t(), labels)
return loss
Step 3: Trajectory Rollout in World Model
Simulate episodes using the world model for RL training.
class WorldModelSimulator:
"""
Simulates trajectories in learned world model.
"""
def __init__(self, world_model: PixelWorldModel, vla_model,
feature_alignment: FeatureAlignment, max_episode_length: int = 50):
"""
Args:
world_model: Trained PixelWorldModel
vla_model: Vision-Language-Action model
feature_alignment: Feature alignment network
max_episode_length: Rollout horizon
"""
self.world_model = world_model
self.vla = vla_model
self.alignment = feature_alignment
self.max_length = max_episode_length
def simulate_trajectory(self, initial_frame: torch.Tensor,
task_instruction: str) -> Dict[str, Any]:
"""
Rollout episode in world model.
Args:
initial_frame: Starting image [3, H, W]
task_instruction: Natural language task description
Returns:
trajectory: {frames, actions, rewards, features}
"""
frames = [initial_frame]
actions = []
features = []
current_frame = initial_frame.unsqueeze(0)
for step in range(self.max_length):
# Get VLA action
with torch.no_grad():
action, feature = self.vla.forward(
current_frame, task_instruction,
return_features=True
)
actions.append(action)
features.append(feature)
# Predict next frame in world model
with torch.no_grad():
next_frame = self.world_model(current_frame, action)
frames.append(next_frame.squeeze(0))
current_frame = next_frame
# Early stopping if task complete (optional: use classifier)
if self._is_task_complete(next_frame, task_instruction):
break
return {
'frames': torch.stack(frames),
'actions': torch.stack(actions),
'features': torch.stack(features),
'trajectory_length': len(actions)
}
def _is_task_complete(self, frame: torch.Tensor, instruction: str) -> bool:
"""Heuristic: detect task completion."""
# In practice: train a task completion classifier
return False
Step 4: On-Policy RL with GRPO
Optimize VLA policy using simulated trajectories.
def compute_trajectory_reward(trajectory: Dict, task_instruction: str,
reward_model=None) -> torch.Tensor:
"""
Compute reward for simulated trajectory.
Args:
trajectory: Simulated episode
task_instruction: Task description
reward_model: Optional learned reward model
Returns:
rewards: Per-step reward estimates
"""
if reward_model is None:
# Simple reward: larger movements toward end = progress
positions = trajectory['frames']
movement = torch.norm(positions[1:] - positions[:-1], dim=(2, 3, 4))
# Normalize and scale
rewards = (movement - movement.mean()) / (movement.std() + 1e-6)
else:
# Learned reward model
features = trajectory['features']
rewards = reward_model.predict(features)
return rewards
def grpo_update(vla_model, trajectories: List[Dict],
task_instruction: str, learning_rate: float = 1e-6):
"""
Group Relative Policy Optimization update.
Args:
vla_model: VLA to optimize
trajectories: List of simulated trajectories
task_instruction: Task instruction
learning_rate: Optimizer learning rate
"""
optimizer = torch.optim.Adam(vla_model.parameters(), lr=learning_rate)
# Compute rewards for each trajectory
rewards = []
for traj in trajectories:
traj_reward = compute_trajectory_reward(traj, task_instruction)
rewards.append(traj_reward.mean())
rewards = torch.tensor(rewards)
# Compute advantages (group relative)
mean_reward = rewards.mean()
advantages = rewards - mean_reward
# Policy gradient update
for trajectory, advantage in zip(trajectories, advantages):
# Get log probabilities of actions in trajectory
log_probs = vla_model.compute_log_prob(
trajectory['frames'],
trajectory['actions'],
task_instruction
)
# Policy loss: gradient ascent on log_prob * advantage
loss = -(log_probs.mean() * advantage)
optimizer.zero_grad()
loss.backward()
optimizer.step()
When to Use WMPO:
When NOT to Use:
Hyperparameters and Configuration:
Pitfalls to Avoid:
Reference: https://arxiv.org/abs/2511.09515
基于 SOC 职业分类
