Simplevla Rl Scaling Vla Training

Scaling Vision-Language-Action Models Through Reinforcement Learning

Outcome: Train VLA models with minimal demonstration data (one trajectory per task) to achieve near-supervised performance with improved generalization and emergent manipulation strategies.

Problem Context

Vision-Language-Action (VLA) models trained via supervised fine-tuning (SFT) face two critical bottlenecks:

1. Data Scarcity: Collecting diverse, high-quality robot demonstrations is expensive and time-consuming. Current methods require substantial trajectory datasets per task.

2. Generalization Failure: SFT models suffer poor generalization to distribution shifts—different object appearances, spatial layouts, or task variants cause performance collapse. On long-horizon tasks, baseline performance drops from 91% to 17.1%.

3. Lack of Exploration: Demonstration data constrains the learned policy to observed behaviors. Suboptimal demonstration patterns become ingrained, preventing discovery of more efficient manipulation strategies.

SimpleVLA-RL addresses these limitations by applying reinforcement learning to VLA models, replacing demonstration-heavy SFT with online RL that explores diverse solution spaces through closed-loop robot interaction simulation.

Core Concept

Instead of treating VLA training as supervised behavior cloning, SimpleVLA-RL reframes it as a sequential decision problem. A pre-trained VLA model generates action tokens in closed-loop with simulated robot environments. Successful task completion provides outcome rewards that propagate to individual actions via policy gradient optimization.

Three core insights drive the approach:

1. Outcome Rewards Suffice: Binary task-completion rewards eliminate the need for hand-crafted, per-task reward functions. A single reward signal works across diverse manipulation domains.

2. Interactive Sampling Matters: Unlike language models that generate independent tokens, robotic policies must observe environment responses. Temperature-based sampling on VLA action logits produces diverse but coherent rollouts.

3. Exploration Enhancements Are Critical: Standard PPO exploration is insufficient for discovery. Targeted modifications—dynamic batch filtering, raised clipping thresholds, elevated sampling temperatures—unlock policy improvement and emergent strategies.

Architecture Overview

System Components:

• Base VLA Model: OpenVLA-OFT (pre-trained vision-language encoder mapping observations to action token distributions). Frozen during RL to preserve learned representations.

• Trajectory Sampling Module: Generates rollouts by sampling action tokens from the VLA logits, executing sampled actions in the simulator, and observing state transitions.

• Environment Simulator: Parallelized multi-environment rendering (PyBullet-based) supporting LIBERO and RoboTwin task variations. Enables efficient batch trajectory collection.

• Reward Evaluator: Applies binary task-completion checking at trajectory level. Success distributes reward uniformly across trajectory steps.

• Policy Optimizer: Proximal Policy Optimization (PPO) engine with VLA-specific modifications. Updates action token logits via policy gradients while keeping frozen backbone.

• veRL Runtime: Distributed RL framework managing parallelization, gradient synchronization, and efficient GPU utilization across multi-environment rollouts.

Implementation

Step 1: Environment Setup and Trajectory Sampling

The first challenge is generating diverse rollouts from a frozen VLA model. Unlike text generation, robot trajectories depend on environment state transitions. Each sampled action affects subsequent observations, creating a closed-loop dependency.

# VLA-Specific Trajectory Sampling
import torch
import numpy as np

class VLATrajectoryCollector:
    def __init__(self, vla_model, simulator, temperature=1.6):
        self.vla = vla_model
        self.simulator = simulator
        self.temperature = temperature

    def rollout_trajectory(self, task_id, max_steps=10):
        """
        Generate single trajectory with temperature-based action sampling.
        Returns: (trajectory_states, trajectory_actions, success)
        """
        obs = self.simulator.reset(task_id)
        trajectory = []
        success = False

        for step in range(max_steps):
            # Get VLA logits for current observation
            with torch.no_grad():
                vla_logits = self.vla.forward(
                    image=obs['rgb'],
                    instruction=obs['task_instruction']
                )  # Shape: (vocab_size,)

            # Temperature-scaled sampling for action diversity
            action_probs = torch.softmax(
                vla_logits / self.temperature,
                dim=-1
            )
            action_token = torch.multinomial(
                action_probs,
                num_samples=1
            ).item()

            # Execute action in simulator, observe next state
            next_obs, done = self.simulator.step(action_token)

            # Store transition for RL training
            trajectory.append({
                'obs': obs,
                'action_token': action_token,
                'action_logits': vla_logits.cpu(),
                'next_obs': next_obs,
            })

            obs = next_obs
            if done:
                success = self.simulator.check_task_success()
                break

        return trajectory, success

    def collect_batch(self, task_ids, batch_size=32):
        """Parallelized trajectory collection across multiple tasks."""
        trajectories = []
        successes = []

        for task_id in task_ids:
            for _ in range(batch_size):
                traj, success = self.rollout_trajectory(task_id)
                trajectories.append(traj)
                successes.append(success)

        return trajectories, np.array(successes)

Key Design: Temperature=1.6 (higher than standard 1.0) increases exploration diversity. The VLA model's frozen representations provide a strong prior, preventing collapse to random exploration.

Step 2: Outcome-Based Reward Computation

Standard robotic RL often requires per-task reward shaping (distance-to-goal, intermediate milestones). SimpleVLA-RL uses purely outcome-based rewards: 1 if task succeeds, 0 otherwise. These rewards distribute uniformly across trajectory steps.

# Outcome Reward Propagation
class OutcomeRewardAssigner:
    def __init__(self):
        pass

    def assign_rewards(self, trajectories, success_labels):
        """
        Convert trajectory-level success to per-step rewards.
        Args:
            trajectories: List of trajectory dicts
            success_labels: Binary array (1=success, 0=failure)
        Returns:
            rewards_per_step: List matching trajectory structure
        """
        trajectory_rewards = []

        for traj, success in zip(trajectories, success_labels):
            # Binary outcome reward
            outcome_reward = 1.0 if success else 0.0
            traj_length = len(traj)

            # Distribute uniformly across steps
            step_rewards = [outcome_reward / traj_length] * traj_length

            trajectory_rewards.append({
                'steps': traj,
                'rewards': step_rewards,
                'return': outcome_reward  # Discounted return (gamma=1)
            })

        return trajectory_rewards

    def compute_advantages(self, trajectory_rewards, value_estimates):
        """
        Compute advantage estimates for policy gradient.
        Args:
            trajectory_rewards: Outcomes with per-step rewards
            value_estimates: Baseline value predictions
        Returns:
            advantages: Normalized advantage estimates
        """
        all_advantages = []

        for traj_data in trajectory_rewards:
            step_rewards = np.array(traj_data['rewards'])
            baselines = value_estimates[traj_data['traj_id']]

            # Advantage = Return - Baseline
            advantages = step_rewards - baselines
            all_advantages.extend(advantages)

        # Normalize advantages for training stability
        advantages_array = np.array(all_advantages)
        normalized = (advantages_array - advantages_array.mean()) / (
            advantages_array.std() + 1e-8
        )

        return normalized

Why This Works: Task-specific reward crafting introduces bias and requires domain expertise. Binary outcomes provide a clean supervision signal that generalizes across task families (LIBERO, RoboTwin, real-world).

Step 3: Exploration-Enhanced PPO

Standard PPO can under-explore in VLA settings where the action space is large (discrete action tokens). SimpleVLA-RL implements three targeted modifications.

# Exploration-Enhanced PPO for VLAs
import torch
import torch.nn as nn
from torch.distributions import Categorical

class VLA_PPO_Optimizer:
    def __init__(
        self,
        vla_model,
        lr=1e-4,
        ppo_eps_clip=0.2,
        entropy_coeff=0.01,
        upper_clip_ratio=1.28,  # Key: Raised from 1.2
        temperature=1.6
    ):
        self.vla = vla_model
        self.lr = lr
        self.ppo_eps_clip = ppo_eps_clip
        self.entropy_coeff = entropy_coeff
        self.upper_clip_ratio = upper_clip_ratio
        self.temperature = temperature
        self.optimizer = torch.optim.Adam(
            [p for p in vla_model.parameters() if p.requires_grad],
            lr=lr
        )

    def ppo_loss_with_exploration(
        self,
        new_logits,
        old_logits,
        actions,
        advantages,
        dynamic_batch_filter=True
    ):
        """
        Modified PPO loss with exploration enhancements.

        Modifications:
        1. Dynamic Batch Filtering: Skip batches with 0% or 100% success
        2. Raised Upper Clipping: Allow larger probability increases for
           rare actions (lower old_prob → larger ratio)
        3. Temperature Scaling: Higher temperature increases exploration
        """

        # Compute old and new action probabilities
        old_dist = torch.softmax(old_logits / self.temperature, dim=-1)
        new_dist = torch.softmax(new_logits / self.temperature, dim=-1)

        # Gather probabilities for sampled actions
        old_probs = old_dist.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
        new_probs = new_dist.gather(-1, actions.unsqueeze(-1)).squeeze(-1)

        # Probability ratio
        prob_ratio = new_probs / (old_probs + 1e-10)

        # Clipped surrogate objective with exploration modifications
        # Standard PPO: min(ratio * adv, clip(ratio, 1-eps, 1+eps) * adv)
        # Modified PPO: Asymmetric clipping to encourage rare actions

        clipped_ratio = torch.clamp(
            prob_ratio,
            min=max(1.0 - self.ppo_eps_clip, 0.0),
            max=1.0 + self.upper_clip_ratio  # Increased upper bound
        )

        policy_loss = -torch.min(
            prob_ratio * advantages,
            clipped_ratio * advantages
        ).mean()

        # Entropy regularization encourages exploration
        entropy = -(new_dist * torch.log(new_dist + 1e-10)).sum(dim=-1).mean()

        total_loss = policy_loss - self.entropy_coeff * entropy

        return total_loss, policy_loss, entropy

    def dynamic_batch_filtering(self, trajectories, success_labels):
        """
        Exclude batches where all trajectories succeed or all fail.
        Maintains meaningful advantage signals for gradient flow.
        """
        success_rate = success_labels.mean()

        # Skip batches with extreme success rates
        if success_rate < 0.05 or success_rate > 0.95:
            return None, None  # Skip this batch

        return trajectories, success_labels

    def update_step(self, batch_trajectories, batch_success_labels):
        """Single PPO update with exploration enhancements."""

        # Dynamic Batch Filtering (Modification 1)
        filtered_trajs, filtered_success = self.dynamic_batch_filtering(
            batch_trajectories,
            batch_success_labels
        )

        if filtered_trajs is None:
            return {"policy_loss": 0.0, "skipped": True}

        # Collect old logits (before update)
        old_logits = torch.cat([
            torch.stack([s['action_logits'] for s in traj])
            for traj in filtered_trajs
        ])

        # Forward pass to get new logits
        new_logits = []
        actions = []

        for traj in filtered_trajs:
            for step in traj:
                obs_tensor = self._prepare_observation(step['obs'])
                with torch.no_grad():
                    logits = self.vla.forward(obs_tensor)
                new_logits.append(logits)
                actions.append(step['action_token'])

        new_logits = torch.stack(new_logits)
        actions = torch.tensor(actions)

        # Compute advantages from outcome rewards
        advantages = self._compute_advantages(filtered_success)

        # Compute loss with exploration modifications
        loss, policy_loss, entropy = self.ppo_loss_with_exploration(
            new_logits,
            old_logits,
            actions,
            torch.tensor(advantages),
            dynamic_batch_filter=True
        )

        # Gradient update
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            [p for p in self.vla.parameters() if p.requires_grad],
            max_norm=1.0
        )
        self.optimizer.step()

        return {
            "policy_loss": policy_loss.item(),
            "entropy": entropy.item(),
            "total_loss": loss.item()
        }

    def _prepare_observation(self, obs):
        """Convert observation dict to model input tensor."""
        # Implementation depends on VLA model's input format
        rgb_tensor = torch.from_numpy(obs['rgb']).float() / 255.0
        return rgb_tensor

    def _compute_advantages(self, success_labels):
        """Placeholder for advantage computation from outcomes."""
        return success_labels.astype(np.float32)

Exploration Rationale:

• Dynamic Filtering (Modification 1): Batches with 0% or 100% success provide no gradient signal. Skipping preserves meaningful advantages.

• Raised Clipping (Modification 2): Standard PPO uses upper clip of 1.2. Raising to 1.28 allows low-probability actions larger probability increases. Enables policy to adopt rare strategies not in demonstration data.

• Temperature=1.6 (Modification 3): Higher temperature flattens action distribution, increasing sample diversity. Drives exploration toward novel strategies like "pushcut."

Step 4: Multi-Environment Parallelization

Efficient scaling requires collecting rollouts across many environments simultaneously. The framework leverages PyBullet's batched physics and parallelized rendering.

# Parallelized Multi-Environment Rollout
from multiprocessing import Pool
import concurrent.futures

class ParallelEnvironmentSimulator:
    def __init__(self, task_configs, num_processes=8):
        self.task_configs = task_configs
        self.num_processes = num_processes
        self.simulators = [
            self._create_simulator(config)
            for config in task_configs
        ]

    def _create_simulator(self, task_config):
        """Initialize PyBullet simulator for task."""
        import pybullet as p

        client = p.connect(p.DIRECT)
        p.setAdditionalSearchPath(task_config['asset_path'])
        p.loadURDF("plane.urdf")

        robot = p.loadURDF(
            task_config['robot_urdf'],
            basePosition=task_config['robot_base']
        )

        return {
            'client': client,
            'robot': robot,
            'task_config': task_config
        }

    def rollout_single_trajectory(self, env_idx, vla_model, task_id):
        """Execute trajectory in single environment."""
        import pybullet as p

        env = self.simulators[env_idx]
        obs = self._get_observation(env, task_id)
        trajectory = []

        for step in range(10):
            # VLA inference
            logits = vla_model.forward(obs['rgb'], task_id)
            action = torch.multinomial(
                torch.softmax(logits / 1.6, dim=-1),
                num_samples=1
            ).item()

            # Environment step
            joint_targets = self._token_to_joint_targets(action)
            for _ in range(10):  # Substeps for stability
                p.setJointMotorControlArray(
                    env['robot'],
                    jointIndices=range(7),
                    controlMode=p.POSITION_CONTROL,
                    targetPositions=joint_targets,
                    forces=[50] * 7
                )
                p.stepSimulation()

            obs = self._get_observation(env, task_id)
            trajectory.append((action, obs))

        success = self._check_task_success(env, task_id)
        return trajectory, success

    def collect_batch_parallel(
        self,
        vla_model,
        task_ids,
        trajectories_per_task=4,
        num_workers=8
    ):
        """
        Collect trajectories across all environments in parallel.
        Returns: (list of trajectories, success labels)
        """
        all_trajectories = []
        all_successes = []

        with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = []

            for task_id in task_ids:
                for _ in range(trajectories_per_task):
                    env_idx = len(futures) % len(self.simulators)
                    future = executor.submit(
                        self.rollout_single_trajectory,
                        env_idx,
                        vla_model,
                        task_id
                    )
                    futures.append(future)

            for future in concurrent.futures.as_completed(futures):
                trajectory, success = future.result()
                all_trajectories.append(trajectory)
                all_successes.append(success)

        return all_trajectories, np.array(all_successes)

    def _get_observation(self, env, task_id):
        """Capture RGB + task instruction."""
        import pybullet as p

        width, height = 512, 512
        view_matrix = p.computeViewMatrix(
            cameraEyePosition=[0.5, 0.5, 0.8],
            cameraTargetPosition=[0, 0, 0],
            cameraUpVector=[0, 0, 1]
        )
        proj_matrix = p.computeProjectionMatrixFOV(
            fov=60,
            aspect=1.0,
            nearVal=0.01,
            farVal=10
        )

        _, _, rgb, _, _ = p.getCameraImage(
            width, height,
            viewMatrix=view_matrix,
            projectionMatrix=proj_matrix,
            renderer=p.ER_BULLET_HARDWARE_OPENGL
        )

        return {
            'rgb': np.array(rgb),
            'task_instruction': self.task_configs[task_id]['instruction']
        }

    def _check_task_success(self, env, task_id):
        """Task-specific success evaluation."""
        # Implementation varies per task suite
        pass

    def _token_to_joint_targets(self, action_token):
        """Decode action token to 7-DOF arm targets."""
        # Map discrete tokens to continuous control
        pass

Performance Gain: Parallelizing across 8+ environments increases data collection throughput while maintaining GPU utilization for inference. Critical for scaling to longer training runs.

Practical Guidance

When to Use SimpleVLA-RL

Ideal Scenarios:

1. Limited Demonstration Data: You have 1-10 trajectories per task but task simulator is available.

2. Diverse Task Families: Training across 10+ related tasks (e.g., LIBERO's 50 manipulation tasks). Outcome rewards apply universally.

3. Generalization Matters: Distribution shifts are expected (new objects, layouts, real-world deployment). RL improves generalization by 20-36% over SFT.

4. Discovery of Novel Strategies: Current demonstrations use suboptimal approaches. Exploration can uncover more efficient behaviors (e.g., pushcut strategy).

5. Sim-to-Real Transfer: Task simulators exist (PyBullet, MuJoCo) enabling cheap online learning before real deployment.

When NOT to Use SimpleVLA-RL

1. Minimal Base Competency: RL completely fails if the pre-trained VLA model achieves 0% success. A frozen backbone provides essential exploration prior. Consider supervised pre-training first.

2. No Simulator Available: RL requires rollout environments. If only demonstration data exists, stick with SFT or offline RL methods.

3. Extremely Complex Tasks: Very long-horizon tasks (100+ steps) with sparse rewards face exploration collapse. Consider reward shaping or hierarchical approaches.

4. Real-Robot-Only Scenarios: Direct real-world RL is unsafe and data-inefficient. Sim-to-real transfer via simulation is core to this method.

5. High Latency Requirements: Parallelized trajectory collection requires computational resources. Edge-deployment or resource-constrained scenarios favor lighter SFT.

Hyperparameters & Configuration

Common Pitfalls & Solutions

Pitfall 1: Training Collapse on Zero-Success Tasks

Problem: RL fails if initial success rate is <1%.

Solution: Validate pre-trained VLA achieves >10% success before starting RL. If not, apply supervised fine-tuning on 50+ expert trajectories first.

Pitfall 2: Exploration Too High, No Convergence

Problem: Temperature=2.0+ causes erratic trajectories; success rate stays near 50%.

Solution: Reduce temperature to 1.4–1.6. Monitor variance of action probabilities; they should be meaningfully concentrated.

Pitfall 3: Advantage Signal Noise

Problem: Binary outcome rewards create sparse signal; most steps have zero advantage.

Solution: Use value-function baselines to estimate per-step returns. Or apply reward-weighted trajectory sampling (favor successful trajectories).

Pitfall 4: Sim-to-Real Gap

Problem: RL achieves 95% in simulation but fails on real robot.

Solution: Add environment randomization during RL (randomize object textures, lighting, physics parameters). Train on diverse object categories; test transfer to unseen objects.

Pitfall 5: Memory Overflow During Parallelization

Problem: Storing 1000s of trajectories with RGB observations exhausts VRAM.

Solution: Stream trajectories directly to optimizer without buffering. Use on-policy RL (PPO) not off-policy (SAC), which requires replay buffers.

Debugging & Monitoring

Key Metrics to Track:

• Success Rate: Primary outcome metric. Should increase from baseline 17% to 90%+ within 5-10K environment steps.

• Advantage Variance: Low variance (<0.1) indicates collapsed exploration. Increase temperature or entropy coefficient.

• KL Divergence (New vs Old Policy): Should stay <0.1. Higher indicates unsafe policy updates; lower temperature-sampling.

• Entropy: Monitor action distribution entropy. Collapsed entropy means exploration failure.

• Gradient Norm: Should remain ~0.1–1.0. Very high (>10) suggests instability; use gradient clipping.

Reference

Paper: SimpleVLA-RL: Scaling VLA Training via Reinforcement Learning

arXiv: https://arxiv.org/abs/2509.09674

Key Results:

• LIBERO-90: 99.1% average success (+8% vs baseline)
• LIBERO-Long: 91.7% vs 17.1% baseline (+74.6%)
• RoboTwin 1.0: 30.6% relative improvement
• RoboTwin 2.0: 80% relative improvement
• Real-world dual-arm tasks: 21 percentage point improvement

Citation:

@article{simplevlrl2024,
  title={SimpleVLA-RL: Scaling VLA Training via Reinforcement Learning},
  author={...},
  journal={arXiv preprint arXiv:2509.09674},
  year={2024}
}