Roboscape Physics World Model

RoboScape: Physics-Informed World Models for Robotic Manipulation

Existing video diffusion models excel at generating visually realistic content, but they lack physical awareness needed for robotic manipulation tasks. A model might generate a video where a robot arm smoothly glides through an object, violating physics constraints, or where keypoints shift inconsistently across frames. This visual plausibility without physical correctness makes generated videos unsuitable for training robotic policies—policies trained on physically implausible demonstrations learn implausible behaviors.

RoboScape solves this by integrating physics into the video generation process through multi-task learning. Instead of generating RGB frames in isolation, the model simultaneously predicts depth, tracks physical keypoints, and generates RGB, with these tasks constraining each other. The result is a world model that generates physically consistent robotic videos suitable for both visual understanding and policy training.

Core Concept

RoboScape's innovation is treating video generation as a multi-task problem where physical constraints emerge from joint learning. The key insight is that:

1. Depth prediction enforces 3D geometric consistency across frames
2. Keypoint tracking ensures object deformations follow physical laws
3. RGB generation becomes implicitly constrained by the geometry and keypoints it must align with

By fusing intermediate features between the RGB and depth branches during decoding, the model learns that realistic videos have consistent geometry. Keypoint trajectories with high temporal continuity tell the model that abrupt position changes violate physics. This turns physics enforcement from an explicit loss term into an implicit emergent property of multi-task learning.

Architecture Overview

The RoboScape architecture uses an autoregressive Transformer with dual processing branches:

• RGB Branch: Decodes color information with causal temporal attention (current frame depends on previous frames) and bidirectional spatial attention within each frame
• Depth Branch: Predicts per-pixel depth maps with the same spatial-temporal attention structure, enabling 3D reconstruction
• Feature Fusion: Intermediate representations from the depth branch are additively fused with RGB features at each decoding layer, enforcing geometric consistency
• Keypoint Head: A lightweight predictor identifies and tracks high-motion points across frames, learning which visual elements have physical significance
• Temporal Consistency: Cross-frame attention mechanisms ensure depth and keypoints maintain smooth trajectories across time

Implementation

Step 1: Prepare training data with physics annotations

Extract depth and keypoint information from robot videos using pretrained models. This creates multi-modal training signals without manual annotation.

def prepare_physics_annotated_video(video_path, frame_height=256, frame_width=256):
    """
    Process raw video to extract RGB, depth, and keypoint trajectories.
    Uses pretrained depth and pose estimation models to generate physics annotations.
    """
    frames = load_video(video_path)
    depth_estimator = load_pretrained_depth_model('dpt-hybrid')
    keypoint_detector = load_pretrained_keypoint_model('mediapipe')

    rgb_frames = []
    depth_frames = []
    keypoint_sequences = []

    for frame in frames:
        # Resize and normalize
        frame = preprocess_frame(frame, height=frame_height, width=frame_width)
        rgb_frames.append(frame)

        # Estimate depth with high confidence in object regions
        depth = depth_estimator(frame)
        depth_frames.append(depth)

        # Detect keypoints (joints, object corners, contact points)
        keypoints = keypoint_detector(frame)
        keypoint_sequences.append(keypoints)

    # Ensure temporal consistency in keypoint assignments
    keypoint_tracks = temporal_keypoint_matching(keypoint_sequences)

    return {
        'rgb': torch.stack(rgb_frames),
        'depth': torch.stack(depth_frames),
        'keypoints': keypoint_tracks
    }

Step 2: Build the dual-branch autoregressive transformer

Create an architecture with separate RGB and depth branches that share intermediate representations through fusion.

import torch
import torch.nn as nn

class RoboScapeDualBranch(nn.Module):
    """
    Autoregressive transformer with RGB and depth branches.
    Features are fused between branches to enforce geometric consistency.
    """
    def __init__(self, hidden_dim=768, num_layers=12, vocab_size=8192):
        super().__init__()

        # Shared token embedding and position encoding
        self.embed_dim = hidden_dim
        self.token_embed = nn.Embedding(vocab_size, hidden_dim)
        self.pos_embed = nn.Parameter(torch.randn(1, 2048, hidden_dim))

        # RGB branch: causal temporal attention, bidirectional spatial
        self.rgb_temporal_attn = nn.ModuleList([
            CausalTemporalAttention(hidden_dim) for _ in range(num_layers)
        ])
        self.rgb_spatial_attn = nn.ModuleList([
            BidirectionalSpatialAttention(hidden_dim) for _ in range(num_layers)
        ])
        self.rgb_mlp = nn.ModuleList([
            FeedForward(hidden_dim) for _ in range(num_layers)
        ])

        # Depth branch: same attention pattern but separate parameters
        self.depth_temporal_attn = nn.ModuleList([
            CausalTemporalAttention(hidden_dim) for _ in range(num_layers)
        ])
        self.depth_spatial_attn = nn.ModuleList([
            BidirectionalSpatialAttention(hidden_dim) for _ in range(num_layers)
        ])
        self.depth_mlp = nn.ModuleList([
            FeedForward(hidden_dim) for _ in range(num_layers)
        ])

        # Cross-branch fusion: additively combine features for consistency
        self.fusion_weights = nn.ParameterList([
            nn.Parameter(torch.ones(1) * 0.5) for _ in range(num_layers)
        ])

        # Output heads
        self.rgb_head = nn.Linear(hidden_dim, vocab_size)
        self.depth_head = nn.Linear(hidden_dim, vocab_size)

    def forward(self, rgb_tokens, depth_tokens, keypoint_mask=None):
        """
        Process tokens through dual branches with cross-branch fusion.
        Keypoint mask ensures tracked points receive consistent gradients.
        """
        batch_size = rgb_tokens.shape[0]

        # Embed tokens
        rgb_x = self.token_embed(rgb_tokens) + self.pos_embed[:, :rgb_tokens.shape[1]]
        depth_x = self.token_embed(depth_tokens) + self.pos_embed[:, :depth_tokens.shape[1]]

        # Process through layers with fusion
        for layer_idx in range(len(self.rgb_temporal_attn)):
            # RGB path
            rgb_temporal = self.rgb_temporal_attn[layer_idx](rgb_x)
            rgb_spatial = self.rgb_spatial_attn[layer_idx](rgb_temporal)
            rgb_x = rgb_x + self.rgb_mlp[layer_idx](rgb_spatial)

            # Depth path
            depth_temporal = self.depth_temporal_attn[layer_idx](depth_x)
            depth_spatial = self.depth_spatial_attn[layer_idx](depth_temporal)
            depth_x = depth_x + self.depth_mlp[layer_idx](depth_spatial)

            # Fusion: additively combine features with learnable weight
            # Higher weight on depth features in regions with keypoints
            fusion_weight = self.fusion_weights[layer_idx]
            if keypoint_mask is not None:
                fusion_weight = fusion_weight * (1 + keypoint_mask.unsqueeze(-1))

            rgb_x = rgb_x + fusion_weight * depth_x
            depth_x = depth_x + fusion_weight * rgb_x

        # Generate outputs
        rgb_logits = self.rgb_head(rgb_x)
        depth_logits = self.depth_head(depth_x)

        return rgb_logits, depth_logits

Step 3: Train with multi-task objectives combining RGB, depth, and keypoint consistency

The training loss combines three objectives to ensure the model learns physically consistent generation.

def compute_multitask_loss(rgb_logits, depth_logits, keypoint_preds,
                           target_rgb_tokens, target_depth_tokens,
                           target_keypoints, weights={'rgb': 1.0, 'depth': 0.5, 'keypoint': 0.3}):
    """
    Multi-task training loss combining:
    - RGB cross-entropy for visual quality
    - Depth prediction loss for geometric consistency
    - Keypoint continuity loss for physics plausibility
    """

    # RGB generation loss: standard language modeling
    rgb_loss = torch.nn.functional.cross_entropy(
        rgb_logits.view(-1, rgb_logits.shape[-1]),
        target_rgb_tokens.view(-1),
        reduction='mean'
    )

    # Depth prediction loss: L2 error on depth maps
    depth_loss = torch.nn.functional.mse_loss(
        depth_logits,
        target_depth_tokens,
        reduction='mean'
    )

    # Keypoint continuity loss: minimize jumps in keypoint positions
    keypoint_velocity = torch.diff(keypoint_preds, dim=1)
    target_velocity = torch.diff(target_keypoints, dim=1)
    keypoint_loss = torch.nn.functional.huber_loss(
        keypoint_velocity,
        target_velocity,
        delta=5.0,  # Robust to outliers
        reduction='mean'
    )

    # Attention-weighted keypoint loss: high motion keypoints receive more weight
    motion_magnitude = torch.norm(target_velocity, dim=-1, keepdim=True)
    keypoint_weight = 1.0 + motion_magnitude / motion_magnitude.max()
    weighted_keypoint_loss = (keypoint_loss * keypoint_weight).mean()

    # Combine losses
    total_loss = (weights['rgb'] * rgb_loss +
                  weights['depth'] * depth_loss +
                  weights['keypoint'] * weighted_keypoint_loss)

    return total_loss, {'rgb': rgb_loss, 'depth': depth_loss, 'keypoint': weighted_keypoint_loss}

Step 4: Evaluate on downstream manipulation policy training

The ultimate validation is whether generated videos effectively train robotic policies.

def evaluate_world_model_for_policy_training(world_model, real_videos, num_synthetic=1000):
    """
    Generate synthetic training data and compare policy performance.
    A good world model produces synthetic data that transfers to real policies.
    """

    # Train policy on real videos (ground truth)
    real_policy = train_diffusion_policy(real_videos)
    real_eval = evaluate_on_simulator(real_policy)

    # Generate synthetic videos from world model
    prompts = extract_task_prompts(real_videos)
    synthetic_videos = world_model.generate(prompts, num_samples=num_synthetic)

    # Train policy on synthetic videos
    synthetic_policy = train_diffusion_policy(synthetic_videos)
    synthetic_eval = evaluate_on_simulator(synthetic_policy)

    # Compute correlation: should be > 0.85 for effective transfer
    correlation = compute_pearson_correlation(
        real_eval['success_rate'],
        synthetic_eval['success_rate']
    )

    return {
        'correlation': correlation,
        'real_success': real_eval['success_rate'],
        'synthetic_success': synthetic_eval['success_rate'],
        'is_effective': correlation > 0.85
    }

Practical Guidance

When to use RoboScape:

• You need synthetic video data for training robotic manipulation policies
• You have real robotic video datasets for pre-training
• Physical consistency matters more than maximum visual quality
• You want to evaluate whether generated videos correlate with simulator performance

When NOT to use RoboScape:

• Your primary goal is photorealistic video generation without physics constraints
• You have unlimited real robotic data (generating synthetic data adds no value)
• You're working with non-manipulation domains (RoboScape is specialized for robotics)
• You need real-time generation (autoregressive models have high latency)

Common pitfalls:

• Depth branch collapse: If depth predictions become uniform, increase the depth weight or add depth regularization. Ensure depth ground truth is accurate.
• Keypoint drift: Temporal consistency loss might not be strong enough. Increase the keypoint weight or use a stricter continuity constraint.
• Mode collapse to simple motions: The model might avoid complex physics. Add diverse task demonstrations and increase the number of training steps.
• Low policy correlation: If synthetic data doesn't transfer to policy training, the world model isn't capturing task-relevant physics. Debug by visualizing depth predictions and keypoint trajectories.

Reference

RoboScape: Physics-informed Embodied World Model https://arxiv.org/abs/2506.23135