Upload README.md with huggingface_hub

README.md

@@ -10,63 +10,124 @@ tags:
   - cross-attention
 ---
 
-# Graspmax — GeoMatch & GeoMatch++
+# Graspmax — GeoMatch v2 · GeoMatch++ · GeoMatch v1
 
-**Graspmax** contains two geometry-aware contact prediction models for dexterous robotic grasping, trained on the CMapDataset across 5 robot end-effectors.
+**Graspmax** contains geometry-aware contact prediction models for dexterous robotic grasping,
+trained on the CMapDataset across 5 robot end-effectors (EZGripper, Barrett, Robotiq 3-Finger,
+Allegro, ShadowHand).
+
+> **⚠️ Version notice:** GeoMatch v1 and GeoMatch++ were trained with a corrupted
+> `robot_keypoints.json` (2× scale factor and wrong shadowhand axis-swap stage).
+> **Use GeoMatch v2 for any new work.** v1 and GeoMatch++ are kept for reproducibility only.
+
+---
+
+## Models at a Glance
+
+| Model | Status | File prefix | Val loss | Val acc |
+|---|---|---|---|---|
+| **GeoMatch v2** | ✅ **Recommended** | `geomatch_v2_*` | **1.594** | **0.695** |
+| GeoMatch++ | ⚠️ Deprecated (built on v1 encoders) | `geomatch_pp_*` | 0.350 | 0.940 |
+| GeoMatch v1 | ⚠️ Deprecated (corrupted keypoints) | `geomatch_final / checkpoint_epoch*` | 0.435 | 0.959 |
+
+> The lower loss/higher accuracy of v1 and GeoMatch++ are an artefact of training on corrupted
+> keypoints — the 2× scale inflated keypoint distances making the contact maps geometrically
+> trivial to predict. v2 trains on correct geometry and is the only model that produces valid
+> IK targets during grasp generation.
 
 ---
 
-## Models
+## Architecture
+
+### GeoMatch (v1 and v2 share the same architecture)
 
-### GeoMatch
-Dual GCN encoder (object + robot surface) → projection layers → 5 autoregressive MLP modules → BCE contact map prediction.
+Dual GCN encoder (object + robot surface) → L2-normalised embeddings → linear projection heads
+(512→64) × 2 → 5 autoregressive MLP modules → per-keypoint BCE contact map prediction.
 
 > Based on: [*Geometry Matching for Multi-Embodiment Grasping*](https://openreview.net/forum?id=oyWkrG-LD5) (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*](https://arxiv.org/abs/2412.18998)
+Extends GeoMatch with a **morphology encoder** (GCN over the robot kinematic-tree graph, 9D node
+features, 32 nodes) and a **DCP-style cross-attention transformer** that fuses object geometry with
+robot morphology before contact prediction. Pretrained GeoMatch v1 encoders are frozen.
 
----
+> Based on: [*GeoMatch++: Morphology-Aware Grasping via Correspondence Learning*](https://arxiv.org/abs/2412.18998)
 
-## Architecture Comparison
+### Component Comparison
 
-| Component | GeoMatch | GeoMatch++ |
+| Component | GeoMatch v1 / v2 | 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) |
+| Object GCN encoder | 3 layers × 256 → 512, trainable | Same, **frozen** (from GeoMatch v1) |
+| Robot surface GCN | 3 layers × 256 → 512, trainable | Same, **frozen** (from GeoMatch v1) |
+| Morphology encoder | — | **NEW** GCN(9 → 256×3 → 512), trainable |
 | 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)** |
+| **Total params** | **~1.9M** | **~6.4M (~5.8M trainable)** |
+
+---
+
+## What Changed in v2 (Keypoint Bug Fix)
+
+GeoMatch requires a `robot_keypoints.json` that defines canonical 3D keypoint positions for each
+robot in rest-pose world space. The v1 keypoints had two bugs:
+
+**Bug 1 — 2× scale factor:** The generation script applied `world_pos *= 2.0`, citing
+`HandModel`'s `hand_scale=2.0` class default. However, every actual call site passes
+`hand_scale=1.0`, overriding that default. Because the scale was applied before the inverse-FK
+projection that `HandModel.get_canonical_keypoints()` uses (`T⁻¹[2p;1] ≠ 2·T⁻¹[p;1]`), the
+distortion was not uniform — it grew with each link's distance from the kinematic root, corrupting
+both training labels and inference IK targets.
+
+**Bug 2 — ShadowHand axis-swap at wrong stage:** The `[x, -z, y]` axis permutation for ShadowHand
+was applied to the final world-space `world_pos` (after FK). The reference implementation
+(`gripper_utils.py`) applies it to raw mesh points in link-local space *before* the visual-origin
+transform. Rotation and axis permutation do not commute, so the wrong stage produced scrambled
+keypoint positions for any ShadowHand link with a non-zero visual-origin rotation.
+
+Both bugs were confirmed by comparing `generate_keypoints_json.py` against `gripper_utils.py` and
+verified by observing that v1 ShadowHand tip keypoints had y-values of ~−0.84 m (outside any
+physical hand envelope) versus the corrected ~0.01 m.
 
 ---
 
 ## Training Details
 
-### GeoMatch
+### GeoMatch v2 ✅ (Recommended)
 
 | Setting | Value |
 |---|---|
-| Dataset | CMapDataset (ContactDB + YCB) |
-| Training samples | 41,871 |
+| Dataset | CMapDataset (ContactDB + YCB), fixed keypoints |
 | 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 |
+| Training time | 8.58 hours |
 | Precision | FP32 |
-| **Final val loss** | **0.435** |
+| **Final val loss** | **1.594** |
+| **Final val accuracy** | **0.695** |
 
-### GeoMatch++
+#### GeoMatch v2 Training Curves
+
+| Epoch | Val Loss | Val Accuracy |
+|---|---|---|
+| 0   | 1.935 | 0.205 |
+| 25  | 1.731 | 0.563 |
+| 50  | 1.675 | 0.580 |
+| 100 | 1.649 | 0.632 |
+| 150 | 1.603 | 0.656 |
+| 199 | 1.594 | 0.695 |
+
+---
+
+### GeoMatch++ ⚠️ (Deprecated — built on GeoMatch v1 encoders)
 
 | Setting | Value |
 |---|---|
-| Initialisation | Pretrained GeoMatch encoders (frozen) |
+| Initialisation | Pretrained **GeoMatch v1** encoders (frozen) |
 | Trainable params | ~5.8M |
 | Batch size | 32 per GPU × 8 GPUs = 256 effective |
 | Optimizer | Adam (β₁=0.9, β₂=0.99) |
@@ -75,46 +136,72 @@ Extends GeoMatch with a **morphology encoder** (third GCN over the robot's kinem
 | 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** |
+| **Final val loss** | **0.350** *(artefact of corrupted training data)* |
+| **Final val accuracy** | **0.940** *(artefact of corrupted training data)* |
 
-### GeoMatch++ Training Curves
+#### 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 |
 
 ---
 
+### GeoMatch v1 ⚠️ (Deprecated — corrupted keypoints)
+
+| Setting | Value |
+|---|---|
+| Dataset | CMapDataset (ContactDB + YCB), **corrupted keypoints** |
+| 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** *(artefact of corrupted training data)* |
+| **Final val accuracy** | **0.959** *(artefact of corrupted training data)* |
+
+---
+
 ## 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 v2 ✅ (Use these)
+
+| File | Epoch | Val Loss | Notes |
+|---|---|---|---|
+| `geomatch_v2_epoch50.pth`  | 50  | 1.675 | Early convergence |
+| `geomatch_v2_epoch100.pth` | 100 | 1.649 | Mid-training |
+| `geomatch_v2_epoch150.pth` | 150 | 1.603 | Near-converged |
+| `geomatch_v2_final.pth`    | 199 | 1.594 | **Final model (recommended)** |
+
+### GeoMatch++ ⚠️ (Deprecated)
 
-### 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)** |
+| `geomatch_pp_final.pth`               | 149 | Final (deprecated) |
+
+### GeoMatch v1 ⚠️ (Deprecated)
+
+| 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 (deprecated) |
 
 ---
 
 ## Usage
 
-### GeoMatch
+### GeoMatch v2 (Recommended)
+
 ```python
 import torch, sys
 sys.path.append(".")
@@ -122,23 +209,24 @@ import config
 from models.geomatch import GeoMatch
 
 model = GeoMatch(config).cuda()
-model.load_state_dict(torch.load("geomatch_final.pth", map_location="cuda"))
+model.load_state_dict(torch.load("geomatch_v2_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]
+        obj_pc,               # [B, 2048, 3]   object point cloud
+        robot_pc,             # [B, 6, 3]      robot surface points (6 keypoints)
+        robot_key_point_idx,  # [B, 6]         keypoint indices into robot_pc
+        obj_adj,              # [B, 2048, 2048] object adjacency (sparse COO)
+        robot_adj,            # [B, 6, 6]      robot adjacency
+        xyz_prev,             # [B, 6, 3]      previous keypoint positions
     )
-# contact_map:    [B, 2048, 6, 1]
-# keypoint_probs: [B, 2048, 5, 1]
+# contact_map:    [B, 2048, 6, 1]  — per-object-point × per-keypoint contact probability
+# keypoint_probs: [B, 2048, 5, 1]  — autoregressive keypoint contact probabilities
 ```
 
-### GeoMatch++
+### GeoMatch++ (Deprecated — kept for reproducibility)
+
 ```python
 import torch, sys
 sys.path.append(".")
@@ -151,20 +239,18 @@ 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
+        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]
         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`.
+Morphology graphs are pre-built per robot using `preprocess_morphology.py` → `gnn_morphology_new.pt`.
 
 ---
 
@@ -172,11 +258,12 @@ Morphology graphs (`morph_features`, `morph_adj`) are pre-built per robot using
 
 ```
 models/
-  geomatch.py      # GeoMatch model (dual GCN + AR modules)
+  geomatch.py      # GeoMatch model (shared by v1 and v2)
   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
+config.py          # Hyperparameters for all models
+generate_keypoints_json.py  # Fixed keypoint generator (used for v2 training data)
 ```
 
 ---
@@ -202,4 +289,4 @@ config.py          # Hyperparameters for both models
 ## License
 
 Original GeoMatch code © 2023 DeepMind Technologies Limited, licensed under the Apache License 2.0.  
-GeoMatch++ extension and checkpoints trained by [Dimios45](https://huggingface.co/Dimios45) as part of the Graspmax project.
+GeoMatch++ extension, v2 training, and all checkpoints produced by [Dimios45](https://huggingface.co/Dimios45) as part of the Graspmax project.