adu2021 Train Vision-Language-Action Robots Through Simulated World Models
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.
Core Concept
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:
- Pixel-Based World Model - Predicts next video frames given actions, ensuring alignment with visual domain
- Feature Alignment - World model outputs are regularized to match VLA features, bridging vision and action
- On-Policy RL - Uses GRPO on simulated trajectories to optimize robot behaviors
- Self-Correction - Robot learns to recover from failures through iterative refinement
This approach replaces real-world risk with model-based simulation risk, dramatically reducing training cost.
Architecture Overview
- Vision-Language-Action Model (VLA): Perceives images, reasons about tasks, outputs actions
- Pixel-Based World Model: Predicts next frames [s_t, a_t] → s_{t+1}
- Feature Alignment Layer: Ensures world model outputs match VLA feature space
- Trajectory Simulator: Rolls out imagined episodes using world model
- Policy Optimizer (GRPO): Updates VLA parameters based on simulated rewards
- Real-World Validator: Periodic real validation to detect model divergence
Implementation Steps
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):
"""
World model predicting next pixel frames given current state and action.
"""
def __init__(self, input_channels: int = 3, action_dim: int = 7,
latent_dim: int = 256, num_blocks: int = 4):
"""
Args:
input_channels: RGB channels (3)
action_dim: Dimension of action vector
latent_dim: Latent representation dimension
num_blocks: Number of residual blocks for prediction
"""
super().__init__()
self.action_dim = action_dim
self.latent_dim = latent_dim
# Encoder: compress frames to latent
self.encoder = nn.Sequential(
nn.Conv2d(input_channels, 32, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, latent_dim, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
)
# Action embedding
self.action_embed = nn.Sequential(
nn.Linear(action_dim, 128),
nn.ReLU(),
nn.Linear(128, latent_dim)
)
# Transition predictor: [latent + action_embed] → next_latent
self.transition = nn.ModuleList([
nn.ConvTranspose2d(latent_dim + latent_dim, latent_dim,
kernel_size=3, padding=1)
for _ in range(num_blocks)
])
# Decoder: decompress latent to frames
self.decoder = nn.Sequential(
nn.ConvTranspose2d(latent_dim, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, input_channels, kernel_size=4, stride=2, padding=1),
nn.Sigmoid() # Output in [0, 1]
)
def forward(self, current_frame: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
"""
Predict next frame given current frame and action.
Args:
current_frame: Current image [batch, 3, H, W]
action: Action vector [batch, action_dim]
Returns:
next_frame: Predicted next frame [batch, 3, H, W]
"""
# Encode current frame
latent = self.encoder(current_frame)
# Embed and concatenate action
action_emb = self.action_embed(action)
action_emb = action_emb.unsqueeze(-1).unsqueeze(-1)
action_emb = action_emb.expand(-1, -1, latent.shape[-2], latent.shape[-1])
combined = torch.cat([latent, action_emb], dim=1)
# Predict next latent through transitions
next_latent = combined
for transition in self.transition:
next_latent = transition(next_latent)
next_latent = torch.relu(next_latent)
# Decode to frame
next_frame = self.decoder(next_latent)
return next_frame
def train_world_model(model, train_dataloader, num_epochs: int = 10):
"""
Train world model on video/trajectory data.
Args:
model: PixelWorldModel instance
train_dataloader: Dataloader with (frame_t, action, frame_t+1) tuples
num_epochs: Training epochs
"""
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()
for epoch in range(num_epochs):
total_loss = 0
for batch_idx, (frame_t, action, frame_t1) in enumerate(train_dataloader):
# Predict next frame
pred_frame = model(frame_t, action)
# L2 loss on pixel prediction
loss = criterion(pred_frame, frame_t1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_dataloader)
print(f"Epoch {epoch}: World model loss {avg_loss:.4f}")
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()
Practical Guidance
When to Use WMPO:
- Robotic manipulation tasks with expensive real-world time
- Vision-language models requiring RL refinement
- Tasks where world model errors are acceptable (simulation bias tolerable)
When NOT to Use:
- Tasks requiring perfect real-world dynamics (world model too inaccurate)
- Scenarios with limited simulation-to-reality transfer
- Real-time systems where planning latency is critical
Hyperparameters and Configuration:
- World model update frequency: Every 50 real episodes, retrain on collected rollouts
- Feature alignment weight: 0.1-0.5 (balance world model accuracy with VLA alignment)
- Rollout horizon: 20-50 steps (longer = more RL signal, higher model error risk)
- GRPO update frequency: After every 10 simulated episodes
Pitfalls to Avoid:
- Compounding error - World model errors accumulate over long rollouts; use short horizons initially
- Distribution shift - World model trained on one distribution; regularly validate on real data
- Reward sparsity - Simulated rewards may be weak; use shaped rewards or learned models
- Alignment failure - If world model features drift from VLA, training fails; monitor alignment loss
Reference: https://arxiv.org/abs/2511.09515