Graspmax — GeoMatch & GeoMatch++

Graspmax contains two geometry-aware contact prediction models for dexterous robotic grasping, trained on the CMapDataset across 5 robot end-effectors.

---

Models

GeoMatch

Dual GCN encoder (object + robot surface) → projection layers → 5 autoregressive MLP modules → BCE contact map prediction.

Based on: Geometry Matching for Multi-Embodiment Grasping (NeurIPS 2024)

GeoMatch++

Extends GeoMatch with a morphology encoder (third GCN over the robot's kinematic-tree graph) and a DCP-style cross-attention transformer that fuses object geometry with robot morphology before contact prediction. Pretrained GeoMatch encoders are frozen; only the morphology encoder, transformer, projection heads, and AR modules are trained.

Based on: GeoMatch++: Morphology-Aware Grasping via Correspondence Learning

---

Architecture Comparison

Component	GeoMatch	GeoMatch++
Object GCN encoder	3 layers × 256 → 512	Same, frozen
Robot surface GCN	3 layers × 256 → 512	Same, frozen
Morphology encoder	—	NEW GCN(9 → 256×3 → 512)
Cross-attention	—	NEW DCP transformer (512-dim, 4 heads, 1 layer)
Projection heads	Linear(512→64) × 2	Same, re-initialised
AR keypoint modules	5× MLP	Same, re-initialised
Total params	~1.9M	~6.4M (5.8M trainable)

---

Training Details

GeoMatch

Setting	Value
Dataset	CMapDataset (ContactDB + YCB)
Training samples	41,871
End-effectors	EZGripper, Barrett, Robotiq 3-Finger, Allegro, ShadowHand
Batch size	256
Optimizer	Adam (β₁=0.9, β₂=0.99)
Learning rate	1e-4
Epochs	200
Hardware	AMD Instinct MI300X (192 GB HBM3), ROCm 6.2.4
Training time	22.18 hours
Precision	FP32
Final val loss	0.435

GeoMatch++

Setting	Value
Initialisation	Pretrained GeoMatch encoders (frozen)
Trainable params	~5.8M
Batch size	32 per GPU × 8 GPUs = 256 effective
Optimizer	Adam (β₁=0.9, β₂=0.99)
Learning rate	5e-5
Epochs	150
Hardware	8× AMD Instinct MI300X, ROCm 6.2.4 (DDP)
Training time	~2.8 hours
Precision	FP32
Final val loss	0.350 (↓19% vs GeoMatch)
Final val accuracy	0.940

GeoMatch++ Training Curves

Epoch	Val Loss	Val Accuracy
0	0.465	0.999
25	0.370	0.880
50	—	—
89	0.362	0.902
100	—	—
140	—	—
149	0.350	0.940

---

Checkpoints

GeoMatch

File	Epoch	Notes
checkpoint_epoch50.pth	50	Early convergence
checkpoint_epoch100.pth	100	Mid-training
checkpoint_epoch150.pth	150	Near-converged
geomatch_final.pth	200	Final model (recommended)

GeoMatch++

File	Epoch	Notes
geomatch_pp_checkpoint_epoch50.pth	50	Early convergence
geomatch_pp_checkpoint_epoch100.pth	100	Mid-training
geomatch_pp_checkpoint_epoch140.pth	140	Near-converged
geomatch_pp_final.pth	149	Final model (recommended)

---

Usage

GeoMatch

import torch, sys
sys.path.append(".")
import config
from models.geomatch import GeoMatch

model = GeoMatch(config).cuda()
model.load_state_dict(torch.load("geomatch_final.pth", map_location="cuda"))
model.eval()

with torch.no_grad():
    contact_map, keypoint_probs = model(
        obj_pc,               # [B, 2048, 3]
        robot_pc,             # [B, 6, 3]
        robot_key_point_idx,  # [B, 6]
        obj_adj,              # [B, 2048, 2048]
        robot_adj,            # [B, 6, 6]
        xyz_prev,             # [B, 6, 3]
    )
# contact_map:    [B, 2048, 6, 1]
# keypoint_probs: [B, 2048, 5, 1]

GeoMatch++

import torch, sys
sys.path.append(".")
import config
from models.geomatch_pp import GeoMatchPP

model = GeoMatchPP(config).cuda()
model.load_state_dict(torch.load("geomatch_pp_final.pth", map_location="cuda"))
model.eval()

with torch.no_grad():
    contact_map, keypoint_probs = model(
        obj_pc,               # [B, 2048, 3]   object point cloud
        robot_pc,             # [B, 6, 3]      robot surface points
        robot_key_point_idx,  # [B, 6]         keypoint indices
        obj_adj,              # [B, 2048, 2048] object adjacency
        robot_adj,            # [B, 6, 6]      robot adjacency
        xyz_prev,             # [B, 6, 3]      previous keypoint positions
        morph_features,       # [B, 32, 9]     morphology node features
        morph_adj,            # [B, 32, 32]    morphology adjacency
    )
# contact_map:    [B, 2048, 6, 1]
# keypoint_probs: [B, 2048, 5, 1]

Morphology graphs (morph_features, morph_adj) are pre-built per robot using preprocess_morphology.py and stored in gnn_morphology_new.pt.

---

Repository Structure

models/
  geomatch.py      # GeoMatch model (dual GCN + AR modules)
  geomatch_pp.py   # GeoMatch++ model (+ morphology encoder + DCP transformer)
  gnn.py           # Graph Convolutional Network
  mlp.py           # MLP building block
config.py          # Hyperparameters for both models

---

Citation

@inproceedings{geomatch2024,
  title     = {Geometry Matching for Multi-Embodiment Grasping},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2024},
}

@article{geomatch_pp2024,
  title   = {GeoMatch++: Morphology-Aware Grasping via Correspondence Learning},
  journal = {arXiv preprint arXiv:2412.18998},
  year    = {2024},
}

---

License

Original GeoMatch code © 2023 DeepMind Technologies Limited, licensed under the Apache License 2.0.
GeoMatch++ extension and checkpoints trained by Dimios45 as part of the Graspmax project.