Research

Training a Unified Grasp Policy Across 20 Object Categories

Data curation, category balancing, and augmentation strategies that let a single policy generalize from bearings to brackets to connectors.

Training data visualization across diverse industrial object categories

The natural instinct when building a grasping system for multiple object categories is to train a separate policy per category, or at least to train a policy on a balanced mixture and hope it holds together. Both approaches have problems. Per-category policies do not share any geometric understanding across categories, which means you need fresh training data for every new category you want to support. Mixed training with naive balancing often results in a policy that works well on the largest categories and degrades on smaller, underrepresented ones.

Building a unified grasp policy that genuinely generalizes across 20 industrial part categories required us to rethink three things: how we curate training data to maintain within-category diversity without causing between-category imbalance, how we structure the augmentation pipeline to cover the range of sensor and surface variability we see in real cells, and how we prevent the policy from taking shortcuts that work on the training distribution but break on novel parts. This post covers what we learned from that process.

What "20 Categories" Actually Covers

The 20 categories in our training set span the dominant geometry types in discrete manufacturing assembly: cylindrical parts (shafts, bushings, bearings, threaded fasteners), prismatic parts (brackets, housings, blocks, profiles), planar parts (plates, connectors, clips, gaskets), and compound parts (flanged components, stepped shafts, parts with bosses or recesses). We also include a small but important category of soft or compliant parts, primarily o-rings and rubber grommets, which require different contact reasoning than metal parts.

The categories are defined by dominant surface geometry type, not by function. A bearing outer race and a shaft sleeve might serve different functions but belong to the same cylindrical geometry category and share grasp-relevant surface properties. This definition makes the category structure more stable across different manufacturing domains: automotive, electronics, and general precision machining all contain representatives of the same geometry types, even though the specific parts differ.

Category coverage was determined by what actually appears in the factory cells we have studied and what is represented in available CAD mesh databases. Some categories, particularly standard fasteners, have far more available mesh data than others, like non-standard stamped brackets from specific industry applications. This natural imbalance in available data drove most of the curation decisions.

Data Curation: Within-Category Diversity vs. Between-Category Balance

The first curation problem is within-category diversity. A category like "cylindrical parts" that contains 500 meshes of nearly identical M6 hexagon bolts and 30 meshes of more varied bushing types is not usefully diverse. The policy will learn to handle M6 bolts well and will generalize poorly to bushings, because the embedding space for that category is dominated by the bolt geometry.

We measured within-category diversity using a curvature histogram distance metric: for each category, we computed the pairwise distances between all mesh instances in terms of their surface curvature histograms, and calculated the mean and minimum pairwise distance. Categories with low minimum pairwise distance contained near-duplicate meshes, which we pruned. Categories with low mean pairwise distance needed geometric augmentation rather than additional meshes.

The second curation problem is between-category balance. After within-category pruning, our cylindrical category still contained roughly three times as many mesh instances as our compound-body category. Naive uniform sampling across instances would have the policy seeing cylindrical parts at three times the rate of compound parts during training. We used a category-level sampling schedule rather than instance-level uniform sampling: at each training batch, each category contributes an equal number of sample instances regardless of how many total instances that category contains. Within a category, sampling is uniform across instances.

This approach has a cost: the effective training signal from the larger categories is diluted, because you see fewer unique instances from them per epoch. We accepted this tradeoff because the alternative, letting larger categories dominate, produces a policy that performs well on its best-represented categories and fails silently on others. Silent failure on an underrepresented category in production is worse than slightly reduced performance across all categories.

Augmentation Pipeline for Industrial Sensor Conditions

Industrial depth sensors present a specific set of challenges that tabletop benchmark augmentation pipelines are not designed for. The main ones are: specular reflection dropout (structured light sensors lose returns on mirror-finish metal surfaces at certain angles), edge dropout (returns near sharp geometric edges are systematically missing), and additive noise from ambient light and sensor temperature drift.

We augment at the point cloud level, after depth image generation from the mesh but before patch tokenization. For each training instance, we apply a randomized combination of: specular dropout, which removes points in regions where the surface normal angle exceeds a threshold relative to the sensor axis; edge dropout, which removes points within a fixed spatial distance of detected sharp edge features on the mesh; Gaussian noise on point positions; and global density variation, which randomly subsamples the point cloud to simulate varying sensor distances and integration times.

The parameters for each augmentation type were calibrated against real sensor data from our evaluation hardware. We measured the empirical distribution of dropout rates and noise magnitudes on a set of real part scans and set the augmentation parameter ranges to cover that distribution. This calibration step is important: if the augmentation range is too narrow, it doesn't cover the real-world variability; if it is too wide, it adds unrealistic noise configurations that teach the model to be robust to things it will never see in deployment.

One augmentation we deliberately do not apply is color jitter or texture augmentation, because we are working with depth-only inputs. This is a constraint that actually simplifies the augmentation design: the only source of visual variation in a point cloud is geometry and noise, and we can calibrate both directly against real sensor behavior.

Preventing Shortcut Learning

A common failure mode in multi-category grasp policy training is shortcut learning: the policy learns to predict grasps based on category-level cues (this object is a bearing, so grasp around the outer diameter) rather than actual local geometry. Shortcut policies fail on novel instances that belong to the same category but have unusual geometry for that category, and they fail badly on compound-geometry objects that do not fit neatly into a single category.

We identify shortcut behavior through a specific evaluation protocol: for each category in our validation set, we compute the grasp success rate on easy instances (meshes close to the category centroid in our embedding space) and on hard instances (meshes near category boundaries). A policy that has learned genuine geometric reasoning should have a smaller performance gap between easy and hard instances than one that has learned category shortcuts. Large gaps in this evaluation are a signal that the policy is relying on category identity rather than local geometry.

When we saw these gaps in earlier training runs, the fix was not to add more hard instances to the training data. It was to remove category labels from the policy input entirely. The policy never sees a category identifier; it only sees the object-centric point cloud representation. Without an explicit category signal, the policy cannot take the shortcut of pattern-matching on category identity. It is forced to reason from geometry.

The Compound-Body Problem

Compound-body parts remain the most challenging category. A stepped shaft with a flange has a cylindrical body, a disk-shaped flange, and a stepped diameter transition. The globally optimal grasp might be on the cylindrical shaft, on the flange edge, or at the step depending on the mass distribution of the specific part and the downstream placement task. None of these grasp regions looks like the canonical geometry for any single category, and the policy must reason about the part's full geometry to select among them.

We handle this in the policy architecture through a multi-resolution representation that captures both local surface geometry and global part shape. The local patches give the policy information about contact quality at candidate grasp sites; the global shape representation gives it information about the part's extent and symmetry, which is needed to reason about approach direction relative to the part's principal axes. Compound parts benefit more from the global shape representation than simple-geometry parts do, because the global geometry is what distinguishes the candidate grasp sites from each other.

Our current validation numbers for compound-body parts are lower than for the other categories. This is an honest statement about where the system currently is: compound-geometry is harder, our training data for it is sparser, and the policy architecture is still being improved for this case. We do not want to overstate the generality of what we have built. What we can say is that the category-balanced, geometry-grounded training approach we described here is the foundation on which handling compound parts well will eventually rest.

More from RLWRLD

Read more on robot manipulation

Technical posts on manipulation, point-cloud learning, sim-to-real, and industrial deployment.