Source-linked AI summary
KeyGen: Unsupervised Keypoint based Object-Centric Representations for Category-Level Policy Generalization
Shuxin Cao, Liquan Wang, Masoud Moghani, Benjamin Joffe, Animesh Garg
TL;DR
Behavior cloning often struggles to transfer across diverse unseen objects, and exhaustive data collection across deployment variability is infeasible. KeyGen learns canonicalized semantic 3D keypoints and uses them with object-centric geometry to condition a diffusion policy; the paper reports improved generalization across instances and effectiveness in real-world tasks.
Problem
Behavior cloning often fails to transfer from training objects to novel instances with varied shape, size, material, and pose, while exhaustive data collection across deployment variability is infeasible.
Source: §I · p. 1Method
KeyGen learns semantic 3D keypoints in a canonical frame and conditions a diffusion behavior-cloning policy on them and object-centric geometry to generate trajectories.
Source: §I · pp. 1–2Results
Consistent 3D keypoints preserve functional geometry across instances and improve generalization on seen and unseen objects, while real-world pushing and pouring remain effective under real-world sensing and dynamics.
Source: §I · p. 2, B. Object-wise Generalization within a… · p. 6Takeaways & Limitations
The photorealistic benchmark and planning-driven pipeline support controlled evaluation of instance transfer, data efficiency, and robustness to geometric variation.
Source: §I · pp. 1–2
Abstract
from arXiv · showhide
Generalization in robotic manipulation requires policies to perform tasks across diverse unseen object instances that vary in shape, size, and pose. However, conventional behavior cloning (BC) methods often overfit to instance-specific geometry and appearance, limiting transfer to novel objects. We introduce KeyGen, a framework that learns canonicalized semantic 3D keypoints from point clouds and uses them as structured object-centric representations for policy learning. A visuomotor diffusion policy conditions on these keypoints together with object-centric geometry to predict full manipulation trajectories, enabling consistent geometric correspondence across object instances. To evaluate category-level generalization, we construct a photorealistic simulation benchmark with three manipulation tasks and a planning-driven data generation pipeline that produces expert trajectories across diverse object instances. Experiments show that KeyGen significantly outperforms prior methods on both seen and unseen objects under pose variation, scales effectively with additional demonstrations per object, maintains robustness to object rescaling, and achieves strong performance in both simulation and real-world manipulation.
I. INTRODUCTION
Category-level manipulation remains difficult because behavior cloning often fails to transfer across within-category variation, while available robotic datasets cannot exhaustively cover deployment variability. KeyGen addresses this with canonicalized semantic 3D keypoints, a conditioned diffusion policy, and a planning-driven benchmark for controlled evaluation.
I. INTRODUCTION
- Behavior cloning often fails to transfer from training objects to novel instances, while exhaustive data collection across deployment variability is infeasible.Objects within a category can vary in shape, size, material, and pose.Source: §I · p. 1
- KeyGen learns semantic 3D landmarks from multi-view RGB-D and conditions a diffusion behavior-cloning policy on keypoints and object-centric geometry to generate manipulation trajectories.The representation canonicalizes orientation and is designed to remain consistent across shape, pose, and appearance variation.Source: §I · pp. 1–2
- Evaluation covers transfer to unseen instances, increased demonstrations per object, zero-shot rescaling, and real-world pushing and pouring.The real-world experiments show the keypoint-conditioned policy remains effective under real-world sensing and dynamics.Source: §I · pp. 1–2
- The paper introduces a photorealistic simulation benchmark and planning-driven pipeline that generates expert trajectories across diverse object instances.The benchmark supports controlled evaluation of instance transfer, data efficiency, and scale robustness.Source: §I · pp. 1–2
II. RELATED WORK
Prior work pursues category-level transfer through part-based, appearance-feature, keypoint, and equivariant representations, each with stated limitations. KeyGen instead learns compact semantic 3D keypoints from multi-view RGB-D in a canonical frame for downstream policy learning.
Representations for Few-Shot Generalization in Manipulation
- Part-based priors can lack consistent 3D correspondences under category variation or clutter, while dense pretrained visual features can miss functional structure for contact-rich control.Keypoint-based one-shot imitation may instead depend on test-time planning or trajectory warping.Source: §II · p. 2
- Field-based methods can require dense sensing or costly construction and inherit limitations of static appearance-driven features; equivariant methods may require inference-time optimization or struggle in cluttered multi-object scenes.These concerns are presented as limitations of the cited approaches, not as universal failures.Source: §II · p. 2
Unsupervised 3D Keypoint Detection
- Unsupervised keypoint approaches use canonical-pose assumptions, equivariant learning, rotation alignment, or canonical-pose prediction, while supervised point-cloud annotation is expensive and ambiguous.KeyGen targets raw multi-view RGB-D by normalizing pose with a pretrained canonical orientation model before learning consistent keypoints.Source: §II · p. 2
III. KEYGEN: APPROACH
KeyGen formulates category-level manipulation as behavior cloning from demonstrations and uses semantic 3D keypoints with object-centric geometry to condition continuous trajectory prediction. Its keypoint detector first canonicalizes object pose, then learns landmarks using paired scans and objectives for spread, surface saliency, and cross-view consistency.
III. KEYGEN: APPROACH
- The policy learns to match expert action distributions from demonstrations and conditions continuous trajectory prediction on instance-consistent keypoints and object-centric geometry.The formulation targets transfer from limited training objects to novel instances within the same category.Source: §III · p. 2
A. 3D Keypoint Detection
- A category-specific classifier selects a canonical orientation from 24 discrete rotations, and the point cloud is normalized before keypoint learning.The classifier is trained on ShapeNet orientations and kept fixed during keypoint learning.Source: A. 3D Keypoint Detection · p. 2
- Two randomly sampled scans are aligned to a shared canonical frame and processed by the detector to produce keypoints parameterized as convex combinations of input points.The detector operates on paired views of the same object to learn view-consistent landmarks.Source: A. 3D Keypoint Detection · p. 2, Figure 3 · p. 3
- The weighted training objective combines separation, heuristic, and consistency losses to encourage spatial diversity, salient surface coverage, and cross-view semantic alignment.The losses support well-spread, surface-adherent, view-invariant keypoints.Source: A. 3D Keypoint Detection · p. 3, Figure 3 · p. 3
- At inference, a single point cloud is canonicalized, keypoints are predicted in that frame, and the landmarks are reprojected to the observation frame.This yields pose-invariant, semantically consistent 3D landmarks for downstream policies.Source: A. 3D Keypoint Detection · p. 4
B. Policy Learning
The policy encodes point clouds, keypoints, and robot proprioception as conditioning features for a diffusion model over action trajectories. At inference, it iteratively denoises a noisy initialization to produce a full trajectory.
B. Policy Learning
- The observation combines an object-centric point cloud, inferred 3D keypoints, and robot proprioception, each encoded independently for policy conditioning.The point cloud and keypoints use separate encoders, while robot state becomes a proprioception token.Source: B. Policy Learning · p. 4
- Training perturbs a clean action trajectory with Gaussian noise and trains a network to predict that noise from the noisy trajectory, conditioning features, and diffusion timestep.The noise prediction uses an ℓ2 denoising objective.Source: B. Policy Learning · p. 4
- At test time, the model starts from Gaussian noise and iteratively denoises it to generate an entire action trajectory over the full horizon.The reverse updates use scheduler-defined scalars and Gaussian noise.Source: B. Policy Learning · p. 4
IV. BENCHMARK & DEMOS
The benchmark provides a photorealistic, configurable setting for pose-sensitive category-level manipulation and generates expert demonstrations from object-centric geometric anchors.
IV. BENCHMARK & DEMOS
- The Isaac Lab environment evaluates mug pouring, knife collecting, and bowl stacking across sampled object instances, poses, and physical parameters.It also supports scanned meshes and controls camera settings, scale, backgrounds, textures, and lighting for domain randomization.Source: §IV · p. 4
- A compact FSM maps canonical mesh anchors to world-space waypoints, then executes them with collision-checked IK and spline retiming.Planner fallback, guard conditions, and success predicates support execution across shape and size variation.Source: §IV · p. 4
- The pipeline records four-view RGB-D, object-centric point clouds, proprioception, and actions as paired observations and expert trajectories.The stated release includes the environment, FSM pipeline, object packs and configs, logging scripts, seeds, and standardized protocols.Source: §IV · p. 4
V. EXPERIMENTS
The experiments assess category-level transfer, data efficiency, and scale robustness in simulation, then compare KEYGEN with baselines in real-world pushing and pouring.
V. EXPERIMENTS
- KEYGEN achieves consistently stronger real-world performance than ACT, DP3, and GenDP, improving success by roughly 15%.The real-world tasks are pushing-based shoe alignment and pouring.Source: §V · p. 4
- The simulation evaluation studies unseen-instance transfer under pose variation, increasing demonstrations per object, and robustness to object-scale shifts.It compares five imitation and diffusion-policy baselines under matched sensing and training budgets and includes ablations of keypoints and object-centric cropping.Source: §V · p. 4
- Figure 4 contrasts KEYGEN rollout examples across the three tasks with baseline failures involving grasp alignment, placement, and orientation.The caption identifies these as common baseline failure modes.Source: Figure 4 · p. 5
A. Experimental Setup
The setup tests three pose-sensitive simulated tasks and two real-world tasks, with matched evaluation protocols and task-specific sensing and training conditions.
A. Experimental Setup
- Simulation covers mug pouring, knife collecting, and bowl stacking, with keypoints trained on 30 same-category ShapeNet instances.Policies use 100 self-generated trajectories from one training instance per task and are evaluated on 12 novel instances, with 20 rollouts per object.Source: A. Experimental Setup · p. 5
- Real-world tests cover shoe alignment and pouring, training on two instances with 60 trajectories per instance and testing on three unseen objects.RGB from fixed and wrist ZED cameras is paired with estimated depth, SAM2 segmentation, and an object-centric point cloud; the keypoint extractor is fine-tuned with additional cup and shoe meshes.Source: A. Experimental Setup · p. 5
- Simulation compares KEYGEN against ACT, DP3, 3DDA, BAKU, and P3-PO; real-world comparisons use ACT, DP3, and GenDP.Within each setting, methods share sensing and training budgets, and performance is measured by task success rate under matched protocols.Source: A. Experimental Setup · p. 5
B. Object-wise Generalization within a Category
Under hard pose variation, baselines drop sharply while KEYGEN remains stable, and KEYGEN's scale robustness is strongest away from extreme rescaling limits.
B. Object-wise Generalization within a Category
- In hard settings, baselines drop sharply while KEYGEN remains stable on seen objects; the pattern also holds for unseen objects, where baselines show grasp and rotation errors.On seen objects in easy settings, all methods exceed 85% on pouring and knife; KEYGEN’s consistent 3D anchors preserve functional geometry across instances.Source: B. Object-wise Generalization within a… · pp. 5–6
- KEYGEN stays high under moderate scale shifts across three tasks and degrades mainly at the smallest and largest tested scales.The scale test trains at 1.0 and evaluates zero-shot at isotropic multipliers 0.5, 0.7, 0.9, 1.0, and 1.2.Source: Figure 5 · p. 6
- Table I reports seen- and unseen-instance success rates across tasks, with easy settings varying position and hard settings varying position and orientation.Stacking involves negligible pose variation; the caption describes single-object training and 20 rollouts per configuration.Source: Table I · p. 6
- Table II compares ablations that remove object-centric cropping, omit keypoints, or replace 3D keypoints with 2D image-based keypoints.The caption states that success rates are reported on seen and unseen objects across three tasks.Source: Table II · p. 6
C. Scalability with Demonstrations per Object
With the training object set fixed, KEYGEN’s held-out-object success continues to improve as demonstrations per object increase, unlike baselines that slow, flatten, or change little.
C. Scalability with Demonstrations per Object
- The data-efficiency evaluation fixes the training set at five object instances and raises demonstrations per object from 10 to 100 to test category-level transfer versus overfitting, reporting pouring and collect-knife results.Source: C. Scalability with Demonstrations per… · p. 6
- KEYGEN gains steadily on held-out objects as demonstrations increase and continues improving at the highest data regime.The figure varies trajectories per object from 10 to 100 for pouring and collect knife.Source: Figure 6 · p. 6, C. Scalability with Demonstrations per… · p. 6
- ACT grows slowly after early gains, 3DDA improves then flattens, and DP3 changes little as demonstrations increase.KEYGEN’s widening gap is attributed to canonicalized keypoints providing cross-instance correspondences that let extra trajectories refine a transferable geometric prior.Source: Figure 6 · p. 6, C. Scalability with Demonstrations per… · p. 6
D. Robustness to Object-Scale Shifts
Under zero-shot isotropic rescaling, KEYGEN remains robust across three tasks, with task-specific degradation mainly at extreme scale shifts for pouring.
D. Robustness to Object-Scale Shifts
- Pouring performs best near nominal scale and degrades mainly at extreme sizes, while knife insertion stays high and stacking remains strong at nominal and larger scales.The evaluation trains at nominal scale and tests isotropic rescaling across all three tasks.Source: D. Robustness to Object-Scale Shifts · p. 6
- Real-world rollouts show generalization across cups and shoes with varied geometry and appearance, and KEYGEN has the highest success on seen and unseen objects in both tasks.The real-robot tasks are pouring and shoe alignment; the largest gains are on unseen objects.Source: Figure 7 · p. 7, Figure 8 · p. 7
- Relative anchors such as a handle, spout, blade axis, or rim support proportional grasp and motion adaptation across size changes.The paper contrasts these geometric anchors with reliance on absolute dimensions or appearance cues.Source: D. Robustness to Object-Scale Shifts · p. 7
E. Ablation Study
KEYGEN achieves stronger real-world transfer to unseen objects than the tested baselines. Its ablations highlight the role of 3D keypoints and object-centric preprocessing, and the conclusion frames canonicalized keypoints as structured priors for action generation.
E. Ablation Study
- Removing keypoints preserves strong seen-object performance but hurts unseen transfer; removing object-centric cropping degrades results, while 2D keypoints perform worse at higher runtime.Source: E. Ablation Study · p. 7
F. Real World Experiment Analysis
- KEYGEN reaches roughly 0.75−0.85 success on held-out objects versus ∼0.63−0.65 for the best baseline, an absolute instance-transfer gain of about 15%.It also achieves roughly 0.85−0.90 success on training instances across pouring and shoe alignment.Source: F. Real World Experiment Analysis · p. 7
VI. CONCLUSION
- KEYGEN estimates canonical orientation and predicts normalized-frame keypoints, producing pose-invariant representations used as structured priors for action generation.The conclusion reports strong performance on seen and unseen objects across three manipulation tasks.Source: §VI · p. 7