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

Fanqi Pu; Lei Jiang; Wenming Yang

cs.ROFeb 28, 2026

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

Authors: Fanqi Pu, Lei Jiang, Wenming Yang

Organizations: Shenzhen International Graduate School, Tsinghua University, Shenzhen, China. · The National and Local Co-Build Humanoid Robotics Innovation Center, Shanghai, China.

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 & tables

Fig. 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.

Fig. 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.

Fig. 3: Remedies for Failure Scenarios. Each subfigure pairs a specific failure mode with our corrective measures.

Fig. 4: Low Contrast Problem and Color Inversion.

Fig. 5: Alignment of Different Visual Modalities with Text Instructions.

	Overall		Task Success Rate (%)
Models	Avg. SR (%) $\uparrow$	Reported Hardware	Close Jar	Drag Stick	Insert Peg	Meat off Grill	Open Drawer	Place Cups	Place Wine	Push Buttons
PerAct [ 12 ]	49.4	8V100 $\times$ 16d	55.2 $\pm$ 4.7	89.6 $\pm$ 4.1	5.6 $\pm$ 4.1	70.4 $\pm$ 2.0	88.0 $\pm$ 5.7	2.4 $\pm$ 3.2	44.8 $\pm$ 7.8	92.8 $\pm$ 3.0
RVT [ 13 ]	62.9	8V100 $\times$ 1d	52.0 $\pm$ 2.5	99.2 $\pm$ 1.6	11.2 $\pm$ 3.0	88.0 $\pm$ 2.5	71.2 $\pm$ 6.9	4.0 $\pm$ 2.5	91.0 $\pm$ 5.2	100.0 $\pm$ 0.0
RVT-2 [ 14 ]	81.4	8V100 $\times$ 20h	100.0 $\pm$ 0.0	99.0 $\pm$ 1.7	40.0 $\pm$ 0.0	99.0 $\pm$ 1.7	74.0 $\pm$ 11.8	38.0 $\pm$ 4.5	95.0 $\pm$ 3.3	100.0 $\pm$ 0.0
SAM2ACT [ 21 ]	86.8	8H100 $\times$ 12h or 2A6000 $\times$ 5d	99.0 $\pm$ 2.0	99.0 $\pm$ 2.0	84.0 $\pm$ 5.7	98.0 $\pm$ 2.3	83.0 $\pm$ 6.0	47.0 $\pm$ 6.0	93.0 $\pm$ 3.8	100.0 $\pm$ 0.0
BridgeVLA [ 22 ]	88.2	48H100 $\times$ 20h	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	88.0 $\pm$ 2.8	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	58.4 $\pm$ 10.0	88.0 $\pm$ 2.8	98.4 $\pm$ 2.2

TABLE I: Results on RLBench. This table compares our TGM-VLA with prior works and ablations. “Avg. SR” is the average success rate over 18 tasks (higher is better). “Reported Hardware” lists the training resources (GPUs

\times

Type

\times

Time). “w/o CTM” and “w/o ITM” denote ablations without Cross-Task and Intra-Task Mixup, respectively.

Fig. 6: Illustration of Intra-Task Mixup (top row) and Cross-Task Mixup (bottom row).

	Overall		Success Rate (%)
Models	Avg. SR (%) $\uparrow$	Avg. Rank $\downarrow$	All Perturbations	MO-COLOR	RO-COLOR	MO-TEXTURE	RO-TEXTURE	MO-SIZE
PerAct [ 12 ]	27.9	5.71	7.2	24.0	29.2	28.8	17.71	35.6
RVT [ 13 ]	35.4	5.28	6.4	26.0	31.3	44.8	41.1	35.3
RVT-2 [ 14 ]	56.7	3.92	15.6 $\pm$ 0.8	53.0 $\pm$ 0.9	54.6 $\pm$ 0.6	59.7 $\pm$ 0.7	56.7 $\pm$ 1.4	60.9 $\pm$ 0.9
SAM2ACT [ 21 ]	61.9	2.93	18.5 $\pm$ 3.5	62.1 $\pm$ 0.9	60.9 $\pm$ 0.3	58.6 $\pm$ 0.6	63.4 $\pm$ 2.1	61.8 $\pm$ 0.5
BridgeVLA [ 22 ]	64.0	2.14	18.7 $\pm$ 2.2	60.5 $\pm$ 1.1	63.8 $\pm$ 0.1	63.5 $\pm$ 1.5	68.4 $\pm$ 3.3	69.3 $\pm$ 1.0

TABLE II: Results on the COLOSSEUM Benchmark. This table shows the success rates across 14 generalization settings (12 unseen environmental perturbations + all perturbations + RLBench). The “Avg. Rank” column reports the average rank of each method across all perturbations, where lower values indicate better overall performance.

Method	Training		Testing
Method	Stack Cups	Setup Chess	Stack Cups	Setup Chess
w/o CIB	80.0 $\pm$ 3.0	16.0 $\pm$ 3.8	74.0 $\pm$ 4.6	12.0 $\pm$ 2.3
Ours	94.0 $\pm$ 1.0	46.0 $\pm$ 7.3	87.0 $\pm$ 2.0	30.0 $\pm$ 4.0
Improvement	+14.0	+30.0	+13.0	+18.0

TABLE III: Ablation Study on Color Inversion. ”w/o CIB” denotes ablations without a color inversion branch.

Method	Basic Tasks			Generalization
Method	Stack Blocks	Put in Drawer	Put on Plate	Distractor	Background
w/o TGM	60%	60%	80%	40%	60%
TGM-VLA	80%	70%	90%	90%	80%
Improvement	+20%	+10%	+10%	+50%	+20%

TABLE IV: Results in the real world. ”w/o TGM” denotes ablations without task-guided mixup technique.

Explore similar work

Data-Asymmetric Latent Imagination and Reranking for 3D Robotic Imitation Learning
May 11, 2026Lianghao Luo, Xizhou Bu, Ruyan Liu +5Robot Skill Learning World Models
Morphi
MoRI: Mixture of RL and IL Experts for Long-Horizon Manipulation Tasks
Apr 11, 2026Yaohang Xu, Lianjie Ma, Gewei Zuo +3Robotic RL Long-Horizon Robotic Manipulation
School of Artificial Intelligence and Automation, Huazhong University of Science and Technology, Wuhan, China. · DRAGON Lab, The University of Tokyo, Tokyo, Japan. · School of Mechanical Science and Engineering, Huazhong University of Science and Technology, Wuhan, China.
$M^2$ -VLA: Boosting Vision-Language Models for Generalizable Manipulation via Layer Mixture and Meta-Skills
Apr 27, 2026Siyao Xiao, Yuhong Zhang, Zhifang Liu +9Language-Conditioned Robot Manipulation VLM Adaptation
Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China. · PengCheng Laboratory, Shenzhen, China. · College of Intelligent Robotics and Advanced Manufacturing, Fudan University, Shanghai, China. +1

May 11, 2026J/K move · Enter open · S save

Data-Asymmetric Latent Imagination and Reranking for 3D Robotic Imitation Learning

Lianghao Luo, Xizhou Bu, Ruyan Liu, Qingqiu Huang+4

Morphi

Robot Skill Learning World Models Robotic Manipulation Data-Efficient Robot Learning Learning from Demonstration Robot Imitation Learning

Open paper AskPDF arXiv

Robotic imitation learning typically assumes access to optimal demonstrations, yet real-world data collection often yields suboptimal, exploratory, or even failed trajectories. Discarding such data wastes valuable information about environment dynamics and failure modes, which can instead be leveraged to improve decision-making. While 3D policies reduce reliance on high-quality demonstrations through strong spatial generalization, they still require large-scale data to achieve high task success. To address this, we propose DALI-R, a Data-Asymmetric Latent Imagination and Reranking framework for 3D robotic imitation learning from mixed-quality trajectories. It learns a Latent World Model over 3D point clouds for imagined rollouts and a Task Completion Scorer that reranks candidate action chunks, improving decision-making without additional high-quality demonstrations. We instantiate DALI-R with both diffusion and efficient flow-matching policies and evaluate it on Adroit and MetaWorld benchmarks. Across the two evaluated 3D base policies, DALI-R achieves an average

6.8

% improvement in success rate while incurring less than

0.7\times

additional inference overhead.

Abstract

Explore similar work

Data-Asymmetric Latent Imagination and Reranking for 3D Robotic Imitation Learning

MoRI: Mixture of RL and IL Experts for Long-Horizon Manipulation Tasks

M2M^2M2-VLA: Boosting Vision-Language Models for Generalizable Manipulation via Layer Mixture and Meta-Skills

$M^2$ -VLA: Boosting Vision-Language Models for Generalizable Manipulation via Layer Mixture and Meta-Skills