Aligning Policy Encoders with a Shape Teacher

Madelon Bernardy; Mohammad Mahdi Derakhshani; David M Knigge; Riccardo Valperga; Artem Moskalev

cs.RONov 5, 2026

Aligning Policy Encoders with a Shape Teacher

Authors: Madelon Bernardy, Mohammad Mahdi Derakhshani, David M Knigge, Riccardo Valperga, Artem Moskalev

Organizations: University of Amsterdam · New Theory AI

Abstract

3D-aware representations can improve robotic manipulation by providing geometric information that is robust to changes in viewpoint, object position, and occlusion. However, using 3D inputs at deployment requires additional sensing and processing, making image-based policies more practical. We ask whether an image-based policy can acquire 3D shape information during training while retaining a purely image-based interface at deployment. We propose shape-teacher alignment, a training-time mechanism that uses a pretrained 3D shape representation as an alignment target for an image-based policy encoder. Specifically, we use Point-MAE features from individual objects as targets for corresponding image patches, transferring geometric structure to the policy without changing its deployment inputs or architecture. On LIBERO, shape-teacher alignment produces image representations that are more structured around object shape and corresponding parts, while substantially accelerating policy convergence and improving success rates. For pretrained policies, convergence is accelerated by 4.7x and 5.6x on the Object and Spatial suites, respectively, without requiring the shape teacher or depth input at deployment.

Figures & tables

Figure 1: Patch similarity visualization. The red square marks the query patch and the green square its most similar patch.

From pre-trained
	Same object	Similar-shape object	Unrelated
Baseline	10.2	8.2	81.6
Shape teacher aligned	8.9	35.1	56.0
From scratch
Baseline	17.8	10.7	71.5
Shape teacher aligned	1.3	36.4	62.2

Table 1: Percentage of query patches whose most similar patch falls on the same object, another object with similar shape, or an unrelated region. For each query patch, the match is the patch with the highest cosine similarity among the evaluated images.

Panel of Figure 2 — Figure 2: Object-level image representations on LIBERO-Object for the from-scratch policy. Colors correspond to object categories . Shape-teacher alignment produces clearer object-level structure than the baseline.

Panel of Figure 3 — Figure 3: Policy learning with and without shape-teacher alignment across LIBERO-Object and LIBERO-Spatial, for pretrained and from-scratch policies. Shape-teacher alignment accelerates convergence in both training regimes. Shaded regions show $\pm 1$ standard deviation over three seeds.

Panel of Figure 4 — Figure 4: Zero-shot robustness to target occlusion and displacement for the pretrained policies. Annotations show the relative change in success rate from standard evaluation. Error bars show $\pm 1$ standard deviation over three seeds.

		Baseline	PointMAE	DINOv2
Success rate (ep. 100)	Object, from pretrained	$0.85\pm 0.00$	$0.90\pm 0.04$	$0.91$
	Object, from scratch	$0.10\pm 0.05$	$0.53$	$0.75$
	Spatial, from pretrained	$0.71\pm 0.02$	$0.82\pm 0.05$	$0.85$
	Spatial, from scratch	$0.20\pm 0.06$	$0.62$	$0.35$
Convergence speed-up	Object, from pretrained	–	$4.7{\times}$	$8.6{\times}$
	Object, from scratch	–	$2.3{\times}$	$\geq 11.1{\times}$

Table 2: Success rate at epoch 100 and convergence speed-up for alignment with PointMAE and DINOv2. Convergence speed-up is measured relative to the corresponding no-alignment baseline. Standard deviation over three seeds are reported for Baseline and PointMAE alignment.

Figure 5: Patch-level similarity after alignment with DINOv2 and PointMAE. DINOv2 produces more object-oriented similarity patterns, whereas PointMAE produces greater similarity across corresponding parts of similarly shaped objects.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Overview of the methodology. On the left, the pipeline to retrieve Point-MAE targets is visualized. On the right, the action-prediction pipeline is visualized.

Shape	Objects
Drink Carton	Milk, Orange Juice
Box	Cream Cheese, Butter, Chocolate Pudding
Tin	Tomato Sauce, Alphabet Soup
Bottle	BBQ Sauce, Salad Dressing, Ketchup

Appendix

Table 3: Similarity groups used for cluster analysis

Panel of Figure 7 — Figure 7: The convergence of each model on the LIBERO-Object and LIBERO-Spatial datasets for different alignment weights $\lambda$ , for both pre-trained and from scratch configurations. Shaded regions show $\pm 1$ standard deviation over three seeds.

From pre-trained
	Same object	Similar-shape object	Unrelated
Baseline	10.2	8.2	81.6
Shape teacher aligned	8.9	35.1	56.0
DINOv2 aligned	31.7	5.7	62.6
From scratch
Baseline	17.8	10.7	71.5

Appendix

Table 4: Percentage of query patches whose most similar patch falls on the same object, another object with similar shape, or an unrelated region, including DINOv2 [ Oquab et al., 2024 ] results. For each query patch, the match is the patch with the highest cosine similarity among the evaluated images.

Panel of Figure 8 — Figure 8: The convergence of each model on the LIBERO-Object and LIBERO-Spatial datasets for different alignment weights $\lambda$ , for both pre-trained and from scratch configurations, for alignment with DINOv2. Shaded regions show $\pm 1$ standard deviation over three seeds.

Panel of Figure 9 — Figure 9: Displacements of the target object per task, for both suites. Orange marks the placements of the original dataset; blue and green mark the two out-of-distribution rings, ring 1 and ring 2, respectively.

Visuomotor Policy Learning

Visuomotor policy learning trains robots to turn camera observations, often with language and proprioception, into actions for manipulation and navigation.

See all 296 papers

R3D: Revisiting 3D Policy Learning
Apr 16, 2026Zhengdong Hong, Shenrui Wu, Haozhe Cui +83D ViTs Visuomotor Policy Learning
Zhejiang University · ShanghaiTech University
Sparse2Act: Learning Action-Aligned Sparse 3D Representations for Cross-Domain Robot Manipulation
Jun 10, 2026Yu Guo, Chang Yu, Siyu Ma +4Representation Learning Generalization in Robotic Manipulation
University of California, Los Angeles · University of California, San Diego · University of Utah
VolumeDP: Modeling Volumetric Representation for Manipulation Policy Learning
Mar 18, 2026Tianxing Zhou, Feiyang Xue, Zhangchen Ye +3Visuomotor Policy Learning Robotic Manipulation
1IIIS, Tsinghua University · 2Galaxea AI

Apr 16, 2026J/K move · Enter open · S save

R3D: Revisiting 3D Policy Learning

Zhengdong Hong, Shenrui Wu, Haozhe Cui, Boyi Zhao+7

Zhejiang University · ShanghaiTech University

3D ViTs Visuomotor Policy Learning Diffusion Policy Imitation Learning Robot Imitation Learning

Open paper AskPDF arXiv

3D policy learning promises superior generalization and cross-embodiment transfer, but progress has been hindered by training instabilities and severe overfitting, precluding the adoption of powerful 3D perception models. In this work, we systematically diagnose these failures, identifying the omission of 3D data augmentation and the adverse effects of Batch Normalization as primary causes. We propose a new architecture coupling a scalable transformer-based 3D encoder with a diffusion decoder, engineered specifically for stability at scale and designed to leverage large-scale pre-training. Our approach significantly outperforms state-of-the-art 3D baselines on challenging manipulation benchmarks, establishing a new and robust foundation for scalable 3D imitation learning. Project Page: https://r3d-policy.github.io/

Abstract

Visuomotor Policy Learning

Explore similar work

R3D: Revisiting 3D Policy Learning

Sparse2Act: Learning Action-Aligned Sparse 3D Representations for Cross-Domain Robot Manipulation

VolumeDP: Modeling Volumetric Representation for Manipulation Policy Learning