PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalizable Robotic Insertion with World Models

Nicklas Hansen, Iretiayo Akinola, Yijie Guo, Jie Xu, Bingjie Tang, Hao Su, Xiaolong Wang, Abhishek Gupta, Dieter Fox, Yashraj Narang

arXiv:2609.28258v1cs.ROcs.CVcs.LG

TL;DR

Existing assembly approaches often use specialist policies for individual insertion tasks, making generalization across objects a central challenge. This paper trains a visual world model using wrist-camera observations and robot proprioception. Zero-shot success on unseen assemblies rises with training-set size, reaching up to 56% after training on 90 assemblies.

  • Problem

    Robotic assembly approaches often rely on specialist policies for individual object pairs, while obtaining controllers with both dexterity and generalization remains challenging.

    Source: §I · p. 1

  • Method

    InsertionWM learns generalist insertion policies with visual world models trained from wrist-camera observations and robot proprioception.

    Source: §I · p. 1

  • Results

    Zero-shot success on unseen assemblies increases with training-set size, reaching up to 56% after training on 90 assemblies.

    Source: B. Results · p. 6

  • Takeaways & Limitations

    The work demonstrates the potential of visual model-based RL for broader and more performant generalist multi-task assembly policies.

    Source: §I · p. 2

    The system performs poorly on assemblies with significant visual occlusion.

    Source: §VII · p. 7

Abstract

from arXiv · show

Robotic assembly in high-mixture settings requires adaptable systems that can handle diverse parts, yet current approaches typically rely on policies specialized to each insertion task. Although this can reach high success rates, it makes the process of deploying systems for new problems tedious and time consuming. We present a framework for generalizable insertion using world models that combine robot proprioceptive information with raw visual observations captured by a wrist-mounted camera. Our model-based approach trains a single world model on up to 90 insertion tasks with geometrically diverse parts, achieving 56% zero-shot success on unseen objects with unknown geometry compared to just 7% with a model-free baseline. Importantly, performance improves as more objects are included in the training dataset, demonstrating strong scalability. Lastly, finetuning the generalist model on held-out objects significantly enhances data-efficiency compared to training from scratch and, in some cases, achieves better asymptotic performance. To our knowledge, this is the first system capable of assembling unseen objects in an entirely data-driven manner, and thus represents a significant step toward scalable, generalizable robotic assembly systems.

I. INTRODUCTION

Specialist policies and hand-crafted controllers make robotic assembly difficult to scale across object geometries and scenes. InsertionWM uses visual world models to train generalist insertion policies, with zero-shot success on unseen assemblies improving as training assemblies are added.

I. INTRODUCTION

  • Specialist policies and hand-crafted controllers struggle to generalize across object geometries and scene configurations, limiting scalability and adaptability.Learning-based specialist systems also require extensive training and often part-specific tuning for each object pair.Source: §I · p. 1
  • InsertionWM trains a generalist policy using world models that combine wrist-camera visual observations with robot proprioception, without object-specific priors.The framework targets zero-shot insertion of unseen objects with entirely unknown geometry.Source: §I · pp. 1–2
  • Predictive world models provide dense self-supervised transition signals and decision-time planning for multi-task learning with a shared visual encoder.The paper identifies sample inefficiency and task interference as challenges for model-free policy gradients in this setting.Source: §I · p. 2
  • Up to 56% zero-shot success on unseen assemblies with unknown geometry is achieved after training on 90 diverse assemblies.Performance improves as more assemblies are added, and finetuning on held-out assemblies improves data-efficiency and sometimes asymptotic performance.Source: §I · p. 2
  • The contribution is to demonstrate visual model-based RL's potential for broader, more performant generalist assembly policies and to study practical design choices.The authors characterize the work as a practical guide rather than a new algorithm.Source: §I · p. 2

II. RELATED WORK

Prior assembly research includes simulation-trained specialist policies and deployments on real hardware, while much robotics RL addresses tasks less demanding in precision. This work focuses on algorithmic learning for precise, vision-based insertion of unseen objects and applies model-based RL to assembly.

II. RELATED WORK

  • Insertion is challenging because parts can have intricate geometries and asymmetries, require precise control, and are generally small.Prior learning-based assembly work has used simulation environments designed to mimic real setups, with policies later deployed on hardware.Source: §II · p. 2
  • AutoMate trains PPO specialist policies online for individual object pairs in simulation and reports an exploratory distillation into a blind generalist.Its setup uses curated assembly assets, a parallelizable simulator, and a shaped reward encouraging top-down insertion.Source: §II · p. 2
  • This work targets precise visual policies for zero-shot insertion of unseen objects with unknown geometry, building on TD-MPC2 model-based RL.The authors describe it as the first successful application of model-based RL algorithms to robotic assembly.Source: §II · p. 2

III. PRELIMINARIES

The paper studies simulated insertion with raw wrist-camera depth and robot proprioception, and formulates control as a partially observable problem. It builds on TD-MPC2, which learns a latent world model and plans actions using that model.

III. PRELIMINARIES

  • The simulated setup uses a Franka Panda with a parallel-jaw gripper, a wrist-mounted RGB-D-like camera, and 200 assets forming 100 object pairs.The hardware configuration mirrors platforms used in related real-world work, but the experiments here are in simulation.Source: §III · p. 2
  • The insertion task is a partially observable MDP, approximating state with 96×96 depth observations and 24-dimensional robot proprioception.Actions are normalized 6-DoF delta end-effector poses, and the objective is to maximize cumulative discounted reward.Source: §III · pp. 2–3
  • TD-MPC2 learns a latent world model from interaction data and uses MPPI with a learned policy prior for sampling-based trajectory optimization at inference.Its end-to-end objective combines joint-embedding prediction, reward prediction, and temporal-difference loss without decoding raw observations.Source: §III · p. 3
  • The model comprises an encoder, latent dynamics, reward predictor, terminal-value predictor, and policy prior; collected trajectories are stored in a FIFO replay buffer.These components encode states, predict next latent states, rewards, future returns, and actions.Source: §III · p. 3
  • The training objective jointly optimizes world-model components over a fixed horizon, weighting near-term predictions more heavily.The objective includes joint-embedding, reward, and value prediction terms; reward and value targets use two-hot discretization with cross-entropy, and stop-gradient blocks gradients through latent prediction targets.Source: §III · p. 3
  • The dataset comprises 100 assemblies that vary greatly in geometry; 10 are reserved as a held-out test set, and the remaining 90 are used for training, with both sets drawn from the same data distribution.Source: Figure 2 · p. 3

IV. DATASET AND ENVIRONMENT

The experiments use 100 geometrically diverse two-part assemblies in a randomized simulation setup, with wrist-camera depth and proprioceptive observations and a shaped insertion reward.

IV. DATASET AND ENVIRONMENT

  • The asset set contains 100 geometrically diverse two-part assemblies, and the method can be applied to larger datasets without algorithmic changes.Insertion strategies vary across assemblies, motivating geometry-aware policies.Source: §IV · pp. 3–4
  • The simulation uses a Franka Panda with randomized socket and plug starts, plus a per-assembly curriculum that adjusts initial states to current success.The socket begins upright for approximately top-down insertion, and the plug is grasped via motion planning before each trajectory.Source: §IV · p. 4
  • Observations combine 96×96 wrist-camera depth images with 24 proprioceptive values, including joint angles, noisy pose estimates, and a goal position.Visual observations can significantly improve insertion of small, thin objects.Source: §IV · p. 4
  • The reward combines distance, imitation, and binary success terms, with an interpenetration penalty for rare simulation errors.The imitation term encourages behavior from demonstrations derived by reversing procedurally generated disassembly trajectories.Source: §IV · p. 4

V. GENERALIZABLE ROBOTIC INSERTION

InsertionWM uses a world model to control insertion from raw wrist-camera depth and robot proprioception. Its multimodal encoder fuses visual and proprioceptive features, while training relies on visual information to distinguish objects.

V. GENERALIZABLE ROBOTIC INSERTION

  • InsertionWM takes proprioception and raw wrist-camera depth as input and outputs delta end-effector pose commands, without CAD models or object IDs at deployment.The commands are executed via planning with the learned world model.Source: §V · p. 4
  • The multimodal encoder uses a shallow depth ConvNet and a two-layer proprioception MLP, then fuses their outputs into a latent state.The encoder is used for both single-object and multi-object world models.Source: §V · pp. 4–5
  • Wrist-camera depth observations provide rich visual information that lets the model infer object geometry in a data-driven manner.The authors state that adding visual observations is crucial to framework performance.Source: Figure 4 · p. 5, §V · p. 5
  • Generalist models train across parallel environment instances with randomized scenes and rely on visual information, rather than object IDs, to distinguish assemblies.The described setup uses one instance per unique assembly when training multi-task models.Source: §V · p. 5

VI. EXPERIMENTS

The evaluation covers both single-assembly specialists and multi-assembly generalists, using average insertion success across available assets as its key metric.

VI. EXPERIMENTS

  • The evaluation includes both single-assembly specialist policies and multi-assembly generalist policies.The experiments use the setup described in Section IV.Source: §VI · p. 5
  • Average insertion success rate across all available assets is the key performance metric.Source: §VI · p. 5
  • The approach is validated through experimental evaluation of the two policy settings.Source: §VI · p. 5

A. Experimental Details

The experiments compare specialist and generalist policies using held-out assemblies and AutoMate as a baseline, and report both success and training requirements.

A. Experimental Details

  • The last 10 of 100 assemblies form a fixed, in-distribution held-out test set, while generalists train on the remaining 90.When training on subsets, assemblies are selected by ascending identifier.Source: A. Experimental Details · p. 5
  • Generalist rates average across assemblies and 100 trials per assembly; specialist rates average across five independent training runs and 100 trials per run.Source: A. Experimental Details · p. 5
  • AutoMate is the primary baseline; unlike its generalist pipeline, InsertionWM trains online from raw visual inputs and needs no per-object CAD at deployment.The study also compares blind models, alternate camera placement, curriculum variants, and a generalist trained with expert rollouts.Source: A. Experimental Details · p. 5
  • All world models have 5M parameters; specialists train for 2M steps, AutoMate for 50M, and generalists for 7M.Source: A. Experimental Details · p. 5
  • 79.78% overall success for our specialist method versus 70.31% for AutoMate across 100 geometries; higher occlusion tends to accompany lower success.Source: Figure 5 · p. 6
  • Training takes 40 hours for a visual specialist and three days for a 90-object generalist on an NVIDIA RTX 3090 GPU.Source: A. Experimental Details · p. 6

B. Results

The model-based approach scales across assemblies and generalizes zero-shot to unseen objects, while specialist performance also exceeds AutoMate overall. Finetuning improves data-efficiency, and ablations identify vision and curriculum as useful design choices.

Single-assembly specialists

  • 79.78% success across 100 specialist assemblies versus AutoMate’s 70.31%, with performance varying by socket geometry.The approach struggles with large sockets, where visual occlusion is likely, and performs significantly better on small, thin sockets.Source: B. Results · p. 6

Multi-assembly generalists

  • 55.0% success on 90 seen assemblies versus AutoMate’s 15.6% on 80, and up to 56% zero-shot success on unseen assemblies after training on 90.At that scale, unseen-assembly performance reaches in-domain performance; AutoMate’s training and test splits are unknown, and its generalist uses CAD-model point-cloud embeddings.Source: B. Results · p. 6

Finetuning a generalist on held-out assemblies

  • Finetuning the 90-assembly generalist on held-out assemblies improves data-efficiency and sometimes exceeds scratch-trained specialists’ asymptotic performance.The latter occurs for assemblies 1036 and 1041; Figure 6 plots success rate against environment steps for both variants.Source: B. Results · p. 6, Figure 6 · p. 7

Analysis & ablations

  • Vision is critical to both specialist and generalist performance, while large plugs that partly occlude sockets are a failure-prone case.The ablations compare vision with proprioception-only control and an alternative camera placement.Source: B. Results · pp. 6–7, Figure 7 · p. 7
  • Curricula improve specialist and generalist performance across the full state distribution, but the particular curriculum matters less.Expert demonstrations improve performance only marginally and are not used because the approach scales poorly with larger datasets.Source: B. Results · p. 7

VII. CONCLUSION

The system achieves zero-shot insertion of unseen objects with unknown geometry, with success improving as training data grows. Important remaining boundaries include visual occlusion, limited assembly datasets, and challenges in real-hardware deployment.

VII. CONCLUSION

  • 56% average zero-shot success on held-out tasks after training on 90 insertion tasks, with task success improving as training data increases.The conclusion also notes poor performance with significant visual occlusion and a slight decline in in-domain performance as training objects increase.Source: §VII · p. 7
  • Real-hardware deployment remains challenging because of simulation-to-real discrepancies, inaccurate depth sensing for small objects, and real-time visual-processing demands.Noisy commodity depth blurs small-part edges, and contact is difficult to simulate faithfully.Source: §VII · p. 7
  • The authors identify broader pose generalization and active perception that reduces geometric uncertainty before acting as promising directions.They also point to limited availability of assembly datasets for training generalists.Source: §VII · pp. 7–8
Loading 2609.28258v1…