R2RDreamer: 3D-aware Data Augmentation for Spatially-generalized 2D Manipulation Policies

The evolution of robotic visual manipulation has transitioned from classical feature-based methods to deep learning-driven strategies, enabling robots to learn complex behaviors from demonstrations. Early imitation learning approaches relied on handcrafted features and simple policies, which limited their adaptability. With the advent of deep neural networks, end-to-end visuomotor policies emerged, exemplified by works like GQN (Generative Query Network) and Behavior Cloning. These methods demonstrated impressive capabilities in controlled environments but struggled with generalization across different object poses, viewpoints, and scene configurations.

Recent efforts have focused on leveraging large-scale datasets and simulation environments, such as MuJoCo and PyBullet, to generate diverse training samples. Simulation-based augmentation methods like MimicGen and its variants attempted to synthesize multiple execution trajectories from limited demonstrations, but they faced challenges related to the sim-to-real gap—differences between simulated and real-world physics, appearance, and sensor noise. To mitigate this, real-to-real augmentation methods like DemoGen and R2RGen emerged, directly editing real RGB-D data to produce new samples. These approaches rely heavily on accurate scene parsing, geometry completion, and precise camera pose estimation, which are difficult in cluttered or low-quality data.

Meanwhile, the rise of video models and generative architectures, such as diffusion models and GANs, opened new avenues for visual data synthesis. However, their application in robotic data augmentation has been limited by issues like temporal coherence, scene consistency, and the need for extensive training data. Overall, the field has made significant strides but still faces the fundamental challenge of balancing geometric fidelity, visual realism, and scalability in data augmentation for robotic manipulation.

The core problem addressed by this work is how to achieve effective spatial generalization in robotic manipulation policies with minimal demonstration data. Existing methods either require large, diverse datasets, which are costly and impractical, or depend on complex scene reconstruction in simulation or real-world data, which is computationally expensive and sensitive to perception errors. These limitations hinder the deployment of robots in dynamic, unstructured environments where scene configurations change frequently.

Specifically, the challenge is to generate augmented demonstration data that reflects a wide range of spatial variations—such as object positions, orientations, and camera viewpoints—without sacrificing geometric consistency or visual quality. Achieving this requires a method that can perform scene editing in 3D, handle occlusions realistically, and produce high-fidelity RGB observations suitable for training RGB-based policies. The difficulty lies in balancing the geometric accuracy of scene modifications with the visual plausibility of the resulting images, especially in the presence of occlusions, clutter, and sensor noise. Addressing this problem is crucial for scaling robotic learning to real-world applications where collecting extensive demonstrations is infeasible.

The primary innovation of R2RDreamer is its hybrid approach that combines lightweight 3D scene editing with occlusion-aware 2D video completion. Unlike prior methods that rely on full scene reconstruction or simulation, this framework performs minimal geometric modifications directly in a shared 3D space, preserving the relative geometry of objects and robot actions. It then projects these edited scenes into 2D images, identifying occluded regions through self-occlusion and external occlusion reasoning. The key breakthrough is the use of a deep neural network—based on WAN2.2 architecture—for dense-control video completion, which synthesizes realistic RGB frames conditioned on control signals, maintaining temporal coherence.

This approach effectively decouples the geometric consistency from visual realism, enabling scalable augmentation without demanding complete scene meshes or high-fidelity perception modules. The framework supports both compact 2D policies and larger vision-language models, significantly enhancing the generalization capability of robotic manipulation policies with limited data. Its modular design allows easy integration into existing robotic systems and paves the way for future multi-task, multi-view, and long-horizon learning scenarios.

• �� Scene Editing: In shared 3D space, identify task-relevant object point clouds via segmentation and tracking; apply spatial transformations (translations, rotations) to these point clouds and robot end-effector trajectories, ensuring geometric consistency while expanding spatial diversity.
• �� Occlusion-aware Projection: Render the edited scene into 2D images from the camera viewpoint; identify occluded regions by comparing the projected points with scene geometry, generating masks that distinguish supported and unsupported pixels.
• �� Mask Generation: Construct projection-consistent masks using depth-based lifting and occlusion reasoning, as well as random object-drop masks to simulate occlusion variability, ensuring diverse training samples.
• �� Video Completion: Feed masked control videos into a dense-control image-to-video model based on WAN2.2, conditioned on control signals and optional text prompts; generate temporally coherent RGB frames that fill in occluded regions realistically.
• �� Self-supervised Training: Use original videos and artificially created masks to train the completion model, ensuring it can handle various occlusion scenarios and scene changes.
• �� Data Augmentation: Pair the completed RGB videos with the corresponding edited actions to form augmented demonstration datasets, suitable for training both compact visuomotor policies and larger VLA models.
• �� Experimental Validation: Conduct manipulation tasks with spatial shifts, compare success rates, ablation studies, and analyze the contribution of each component to overall performance.

The experimental setup involves a real robot platform equipped with an RGB-D camera and a parallel-jaw gripper. The tasks include pot-food placement, cup hanging, bridge building, and object covering, designed to evaluate spatial generalization under unseen object positions and configurations. Baseline methods include traditional demonstration, simulation-based augmentation, and existing real-to-real methods like DemoGen and R2RGen. Metrics focus on success rate in held-out configurations, visual quality of augmented data, and policy robustness. The training set comprises limited demonstrations (1, 5, 15, 30), with the augmented data generated via R2RDreamer. Ablation studies remove individual components—3D editing, occlusion-aware projection, or video completion—to assess their impact. Results demonstrate that the full framework significantly outperforms baselines, especially in complex scenarios with non-rigid objects and occlusions, validating its effectiveness in real-world robotic manipulation.

Quantitative results show that with only one source demonstration, success rates improve from 13% to over 40% across tasks, surpassing baseline methods by a large margin. Incorporating all components yields the highest success, with an average increase of 25% in spatial generalization. Ablation experiments reveal that removing occlusion-aware projection or video completion reduces success by approximately 20%, highlighting their importance. Visual assessments confirm that the augmented videos accurately recover occluded regions and maintain temporal consistency, leading to more robust policies. The framework demonstrates strong transferability across different policy types, including diffusion-style and vision-language-action models, indicating broad applicability.

This framework is highly applicable in scenarios requiring rapid adaptation to new environments, such as industrial automation, warehouse logistics, and service robots. It enables robots to learn from minimal demonstrations and adapt to unseen object arrangements or viewpoints, reducing data collection costs. The approach can be integrated into existing robotic systems with RGB-D sensors and standard manipulation hardware. Its ability to generate realistic, geometrically consistent augmented data makes it suitable for training high-performance policies in dynamic, cluttered, or occluded environments. Long-term, this method could facilitate autonomous robots capable of continual learning and adaptation in unstructured settings, significantly advancing the deployment of intelligent robotic agents.

The current approach relies on accurate segmentation and tracking; failures in these modules, especially in cluttered or fast-moving scenes, can impair occlusion reasoning and visual completion. The video completion model, although effective, may produce artifacts in highly complex or extreme occlusion scenarios, affecting policy robustness. Additionally, the method is primarily validated on short-horizon tasks; extending to long-horizon, multi-step tasks requires further improvements in temporal coherence and multi-modal integration. Computational costs associated with scene editing and video synthesis also pose challenges for real-time deployment. Addressing these limitations will be crucial for broader adoption in real-world robotic systems.