TGM-VLA: Task-Guided Mixup for Sampling-Efficient and Robust Robotic Manipulation

TL;DR

A novel holistic framework that significantly improves both model performance and training efficiency is proposed and a task-guided mixup technique that dynamically fuses point clouds and action heatmaps according to task instructions is proposed, greatly improving robustness to distractors and performance in multi-goal scenarios.

Abstract

The performance of robotic imitation learning is fundamentally limited by data quality and training strategies. Prevalent sampling strategies on RLBench suffer from severe keyframe redundancy and imbalanced temporal distribution, leading to inefficient memory usage and unstable optimization. Moreover, reprojecting point clouds onto multi-view images with a black background--while more efficient than voxel-based methods--often causes dark objects to be indistinguishable and hard to manipulate. In this work, we propose a novel holistic framework that significantly improves both model performance and training efficiency. First, we redesign and optimize the keyframe sampling strategy, reducing memory consumption by 80% and accelerating training speed by 5x. Second, we augment the model with a color inversion projection branch--a simple yet effective module that resolves the ambiguity of dark objects. Finally, we propose a task-guided mixup technique that dynamically fuses point clouds and action heatmaps according to task instructions, greatly improving robustness to distractors and performance in multi-goal scenarios. Extensive experiments demonstrate that our method achieves state-of-the-art performance with a 90.5% success rate on RLBench and 68.8% on the COLOSSEUM benchmark under challenging interference conditions. Our code and checkpoints are available at https://github.com/PuFanqi23/TGM-VLA.

Figures (7)

• Figure 1: Data Distribution and Sampling. Red dots represent next target poses. Orange dots represent observed states sampled every 10 frames. Yellow arrows depict keyframe samples, while white arrows depict augmented samples. The original sampling strategy causes data redundancy and temporal imbalance, as each augmented sample forces repeated sampling of all subsequent keyframe samples.
• Figure 2: Overview of TGM-VLA. Given the current point cloud and a task instruction, TGM-VLA predicts the next key-frame pose. This model employs data augmentations including 3D transformations, intra-task, and cross-task mixup. At inference, it first renders point cloud into orthogonal views, enhanced with inverted-color images for improved contrast, to coarsely localize the area of interest. The second stage uses zoomed-in views to precisely predict the gripper pose.
• Figure 3: Remedies for Failure Scenarios. Each subfigure pairs a specific failure mode with our corrective measures.
• Figure 4: Low Contrast Problem and Color Inversion.
• Figure 5: Alignment of Different Visual Modalities with Text Instructions.