Reorganize GeoMatch checkpoints into per-version subfolders; add RobotFingerPrint (arXiv:2409.14519) paper reproduction checkpoints

README.md

@@ -8,13 +8,15 @@ tags:
   - contact-prediction
   - morphology
   - cross-attention
+  - cvae
+  - coordinate-space
 ---
 
-# Graspmax — GeoMatch v2 · GeoMatch++ · GeoMatch v1
+# Graspmax — GeoMatch v2 · GeoMatch++ · GeoMatch v1 · RobotFingerPrint
 
-**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).
+**Graspmax** contains geometry-aware contact/coordinate prediction models for dexterous robotic
+grasping, trained on the CMapDataset / GenDexGrasp dataset 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).
@@ -24,11 +26,12 @@ Allegro, ShadowHand).
 
 ## Models at a Glance
 
-| Model | Status | File prefix | Val loss | Val acc |
+| Model | Status | Folder | 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 |
+| **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_v1/` | 0.435 | 0.959 |
+| **RobotFingerPrint** | ✅ Paper reproduction | `robotfingerprint/` | see below | — |
 
 > 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
@@ -54,7 +57,26 @@ robot morphology before contact prediction. Pretrained GeoMatch v1 encoders are
 
 > Based on: [*GeoMatch++: Morphology-Aware Grasping via Correspondence Learning*](https://arxiv.org/abs/2412.18998)
 
-### Component Comparison
+### RobotFingerPrint
+
+A conditional VAE (`GcsCVAE`) that predicts a per-point **Unified Gripper Coordinate Space (UGCS)**
+2D coordinate map over an object's point cloud, conditioned on the object geometry. Unlike GeoMatch's
+discrete per-keypoint contact classification, RobotFingerPrint regresses a continuous (u, v)
+coordinate for every object point, which is what allows a single trained model to transfer grasps
+across grippers with a *different number of fingers* without any manual re-targeting.
+
+- **Encoder**: PointNet-style per-point Conv1d stack over `(object_pc, gt_uv)` → max-pool global
+  feature → linear heads to VAE latent mean/logvar (`encoder_layers_size=[5, 64, 128, 512, 512]`,
+  `latent_size=128`)
+- **Decoder**: per-point features + global object feature + latent code → Conv1d stack
+  (`decoder_decoder_layers_size=[64+512+128, 512, 64]`) → two parallel U/V prediction heads
+  (`uv_layers_size=[64, 32, 1]`)
+- **Loss**: reconstruction (weighted L2 on predicted vs. ground-truth UV coordinates) + annealed
+  KL divergence (weight increased every `ann_per_epochs` epochs following a temperature schedule)
+
+> Based on: [*RobotFingerPrint: Unified Gripper Coordinate Space for Multi-Gripper Grasp Synthesis*](https://arxiv.org/abs/2409.14519) (Khargonkar, Casas, Prabhakaran, Xiang)
+
+### Component Comparison (GeoMatch family)
 
 | Component | GeoMatch v1 / v2 | GeoMatch++ |
 |---|---|---|
@@ -167,34 +189,80 @@ physical hand envelope) versus the corrected ~0.01 m.
 
 ---
 
+### RobotFingerPrint ✅ (Paper reproduction — 4 experiments)
+
+All 4 experiments use the exact recipe published by the paper's authors (this repo's own README):
+`--n_epochs 16 --ann_temp 1.5 --ann_per_epochs 2`, plus code defaults `lr=1e-4`, `batch_size=64`,
+`lw_recon=1000.0`, `lw_kld=0.01`, `attn_alpha=3`, Adam(β₁=0.9, β₂=0.999), seed=42.
+
+| Setting | Value |
+|---|---|
+| Dataset | GenDexGrasp CMapDataset-sqrt_align + RobotFingerPrint UGCS coordinates |
+| Batch size | 64 |
+| Optimizer | Adam (β₁=0.9, β₂=0.999) |
+| Learning rate | 1e-4 (StepLR decay disabled — `decay_lr_freq=1000` > total epochs) |
+| Epochs | 16 |
+| KL annealing | temperature 1.5, weight increased every 2 epochs |
+| Hardware | AMD Instinct MI300X (192 GB HBM3), ROCm 6.2.4, single GPU per run |
+| Precision | FP32 |
+
+| Experiment | Seen grippers | Held-out (unseen) | Training time | Final val recon loss | Final val KLD | Final val overall |
+|---|---|---|---|---|---|---|
+| `fullrobots/` | ezgripper, barrett, robotiq_3finger, allegro, shadowhand | none (main paper result) | 449 s | **0.2495** | 112.57 | 268.72 |
+| `unseen_barrett/` | ezgripper, robotiq_3finger, allegro, shadowhand | barrett | 366 s | **0.2568** | 105.99 | 274.90 |
+| `unseen_ezgripper/` | barrett, robotiq_3finger, allegro, shadowhand | ezgripper | 383 s | **0.2791** | 109.36 | 297.77 |
+| `unseen_shadowhand/` | ezgripper, barrett, robotiq_3finger, allegro | shadowhand | 366 s | **0.2590** | 111.11 | 278.02 |
+
+> "Val recon loss" is the reconstruction term of the CVAE loss (weighted L2 between predicted and
+> ground-truth UGCS coordinates) — lower is better. The 3 unseen-gripper runs measure how well the
+> model's learned coordinate space generalizes to a gripper never seen during training (evaluated
+> only on the 4 remaining seen grippers' validation split, same as the paper's ablation setup — a
+> full cross-gripper zero-shot transfer evaluation requires the downstream grasp generation +
+> IsaacGym stability test pipeline, not covered by this checkpoint alone).
+
+---
+
 ## Checkpoints
 
 ### 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_v2/checkpoint_epoch50.pth`  | 50  | 1.675 | Early convergence |
+| `geomatch_v2/checkpoint_epoch100.pth` | 100 | 1.649 | Mid-training |
+| `geomatch_v2/checkpoint_epoch150.pth` | 150 | 1.603 | Near-converged |
+| `geomatch_v2/final.pth`               | 199 | 1.594 | **Final model (recommended)** |
 
 ### GeoMatch++ ⚠️ (Deprecated)
 
 | 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 (deprecated) |
+| `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 (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) |
+| `geomatch_v1/checkpoint_epoch50.pth`  | 50  | Early convergence |
+| `geomatch_v1/checkpoint_epoch100.pth` | 100 | Mid-training |
+| `geomatch_v1/checkpoint_epoch150.pth` | 150 | Near-converged |
+| `geomatch_v1/final.pth`               | 200 | Final (deprecated) |
+
+### RobotFingerPrint ✅ (final checkpoint only per experiment)
+
+| File | Experiment | Notes |
+|---|---|---|
+| `robotfingerprint/fullrobots/final.ckpt`         | All 5 grippers seen | **Main paper result — recommended** |
+| `robotfingerprint/unseen_barrett/final.ckpt`     | Barrett held out | Generalization ablation |
+| `robotfingerprint/unseen_ezgripper/final.ckpt`   | EZGripper held out | Generalization ablation |
+| `robotfingerprint/unseen_shadowhand/final.ckpt`  | ShadowHand held out | Generalization ablation |
+
+> These are full PyTorch Lightning checkpoints (model weights + hyperparameters, no optimizer/epoch
+> intermediates — only the last training epoch of each run is kept). Only the final epoch is
+> published; intermediate per-epoch checkpoints are not included here.
 
 ---
 
@@ -209,7 +277,7 @@ import config
 from models.geomatch import GeoMatch
 
 model = GeoMatch(config).cuda()
-model.load_state_dict(torch.load("geomatch_v2_final.pth", map_location="cuda"))
+model.load_state_dict(torch.load("geomatch_v2/final.pth", map_location="cuda"))
 model.eval()
 
 with torch.no_grad():
@@ -234,7 +302,7 @@ 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.load_state_dict(torch.load("geomatch_pp/final.pth", map_location="cuda"))
 model.eval()
 
 with torch.no_grad():
@@ -252,17 +320,66 @@ with torch.no_grad():
 
 Morphology graphs are pre-built per robot using `preprocess_morphology.py` → `gnn_morphology_new.pt`.
 
+### RobotFingerPrint
+
+Lightning checkpoints store hyperparameters alongside weights, so the model reconstructs itself
+directly from the `.ckpt` file — no separate config needed:
+
+```python
+import sys
+sys.path.append(".")
+from robotfingerprint.model.grasp_network import GcsGraspModel
+
+model = GcsGraspModel.load_from_checkpoint("robotfingerprint/fullrobots/final.ckpt")
+model.eval().cuda()
+
+with torch.no_grad():
+    # input_pc: [B, N, 3] object point cloud; gt_gcs only used to establish shape during
+    # training — at inference time use model.model.predict(input_pc) for sampling-based prediction
+    pred_uv = model.model.predict(input_pc.cuda())
+# pred_uv: [B, N, 2] predicted Unified Gripper Coordinate Space (u, v) per object point
+```
+
+For the full downstream pipeline (coordinate inference on held-out objects → grasp generation →
+IsaacGym stability testing), see the
+[RobotFingerPrint repository](https://github.com/IRVLUTD/robot-finger-print) `gcs_gdx_inf_cvae.py`
+and `gcs_gdx_grasp_gen.py` scripts, using these checkpoints via `--logdir`/`--ckpt`.
+
 ---
 
 ## Repository Structure
 
 ```
-models/
-  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 all models
+geomatch_v1/         # GeoMatch v1 checkpoints (deprecated, corrupted keypoints)
+  checkpoint_epoch50.pth
+  checkpoint_epoch100.pth
+  checkpoint_epoch150.pth
+  final.pth
+geomatch_pp/         # GeoMatch++ checkpoints (deprecated, built on v1 encoders)
+  checkpoint_epoch50.pth
+  checkpoint_epoch100.pth
+  checkpoint_epoch140.pth
+  final.pth
+geomatch_v2/         # GeoMatch v2 checkpoints (recommended)
+  checkpoint_epoch50.pth
+  checkpoint_epoch100.pth
+  checkpoint_epoch150.pth
+  final.pth
+robotfingerprint/     # RobotFingerPrint (arXiv:2409.14519) paper reproduction
+  fullrobots/final.ckpt
+  unseen_barrett/final.ckpt
+  unseen_ezgripper/final.ckpt
+  unseen_shadowhand/final.ckpt
+  model/
+    grasp_network.py  # GcsGraspModel (Lightning module)
+    modules.py         # GcsCVAE, PointNetCmapEncoder/Decoder
+    loss.py            # GcsLoss (recon + annealed KLD)
+models/                # Shared GeoMatch source (v1/v2/++)
+  geomatch.py
+  geomatch_pp.py
+  gnn.py
+  mlp.py
+config.py              # Hyperparameters for GeoMatch models
 generate_keypoints_json.py  # Fixed keypoint generator (used for v2 training data)
 ```
 
@@ -282,11 +399,24 @@ generate_keypoints_json.py  # Fixed keypoint generator (used for v2 training dat
   journal = {arXiv preprint arXiv:2412.18998},
   year    = {2024},
 }
+
+@article{khargonkar2024robotfingerprint,
+  title   = {RobotFingerPrint: Unified Gripper Coordinate Space for Multi-Gripper Grasp Synthesis},
+  author  = {Khargonkar, Ninad and Casas, Luis Felipe and Prabhakaran, Balakrishnan and Xiang, Yu},
+  journal = {arXiv preprint arXiv:2409.14519},
+  year    = {2024},
+}
 ```
 
 ---
 
 ## License
 
-Original GeoMatch code © 2023 DeepMind Technologies Limited, licensed under the Apache License 2.0.  
-GeoMatch++ extension, v2 training, and all checkpoints produced by [Dimios45](https://huggingface.co/Dimios45) as part of the Graspmax project.
+Original GeoMatch code © 2023 DeepMind Technologies Limited, licensed under the Apache License 2.0.
+GeoMatch++ extension, v2 training, and all GeoMatch checkpoints produced by
+[Dimios45](https://huggingface.co/Dimios45) as part of the Graspmax project.
+
+RobotFingerPrint model source (`robotfingerprint/model/`) is from the original authors'
+[repository](https://github.com/IRVLUTD/robot-finger-print) (MIT-style license, see their repo for
+exact terms); checkpoints in `robotfingerprint/` were trained by
+[Dimios45](https://huggingface.co/Dimios45) reproducing the paper's published recipe.

robotfingerprint/fullrobots/final.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51dba9d45120fbdd3db72b92f90ab9f6c4818ed68b3f0bdf2408df504c86df69
+size 11523772

robotfingerprint/model/grasp_network.py

@@ -0,0 +1,162 @@
+import lightning as L
+import torch
+
+from .modules import GcsCVAE
+from .loss import GcsLoss
+
+
+class GcsGraspModel(L.LightningModule):
+    def __init__(
+        self,
+        learning_rate: float = 1e-4,
+        cmap_loss_wrecon: float = 100.0,
+        cmap_loss_wkld: float = 0.1,
+        cmap_loss_temp: float = 1.5,
+        cmap_loss_ann_per_epoch: int = 2,
+        pred_type: str = "gcs",
+        loss_attn_weight: float = 3,
+        decay_lr_freq: int = 1000,
+    ):
+        super().__init__()
+        assert pred_type in {"cmap", "gcs", "gcs+cmap", "gcs_u", "gcs_v"}
+
+        self.lr = learning_rate
+        self.loss_criterion = GcsLoss(
+            cmap_loss_wrecon,
+            cmap_loss_wkld,
+            cmap_loss_temp,
+            cmap_loss_ann_per_epoch,
+            loss_attn_weight,
+        )
+
+        self.pred_type = pred_type
+        self.decay_lr_freq = decay_lr_freq
+        if pred_type == "gcs":
+            input_size = 5  # B,N,3 pts  and N, 2 gcs uvmap
+            output_size = (
+                64  # B, N, 64 global feats --> gets mapped to B, N, 2 uv preds
+            )
+        else:
+            raise NotImplementedError
+
+        self.model = GcsCVAE(
+            encoder_layers_size=[input_size, 64, 128, 512, 512],
+            decoder_decoder_layers_size=[64 + 512 + 128, 512, 64],
+            uv_layers_size=[64, 32, 1],
+        )
+        self.cmap_loss_ann_per_epoch = cmap_loss_ann_per_epoch
+
+        print(
+            "Hparams:",
+            decay_lr_freq,
+            loss_attn_weight,
+            cmap_loss_ann_per_epoch,
+            cmap_loss_temp,
+            learning_rate,
+        )
+        self.save_hyperparameters()
+
+    def get_inputs(self, batch):
+        input_pc = batch["full_pc"]  # (B, N, 3)
+        target_cmap = batch["cmaps_fullpc"]  # (B, N, 1)
+        target_gcs = batch["gcs_fullpc"]  # (B, N, 2)
+
+        if self.pred_type == "cmap":
+            target_map = target_cmap
+        elif self.pred_type == "gcs_u":
+            target_map = target_gcs[:, :, :1]
+        elif self.pred_type == "gcs_v":
+            target_map = target_gcs[:, :, 1:]
+        elif self.pred_type == "gcs":
+            target_map = target_gcs
+        elif self.pred_type == "gcs+cmap":
+            target_map = torch.cat(
+                (
+                    target_gcs,
+                    target_cmap.unsqueeze(-1) if target_cmap.ndim == 2 else target_cmap,
+                ),
+                dim=-1,
+            )
+        return input_pc, target_map, target_cmap, target_gcs
+
+    def forward(self, input_pc, gt_gcs):
+        return self.model(input_pc, gt_gcs)
+
+    def training_step(self, batch, batch_idx):
+
+        input_pc, target_map, target_cmap, target_gcs = self.get_inputs(batch)
+
+        _, pred_map, means, logvars, z_latent_code = self.model(input_pc, target_map)
+
+        # For training use the attention weights dictated by the gt cmap
+        loss_cmap, loss_recon, loss_kld = self.loss_criterion(
+            means, logvars, target_map, pred_map, target_cmap
+        )
+
+        loss = loss_cmap
+        loss_dict = {
+            "loss/trn/overall": loss.item(),
+            "loss/trn/cmap_rec": loss_recon.item(),
+            "loss/trn/cmap_kld": loss_kld.item(),
+        }
+        self.log_dict(
+            loss_dict,
+            on_step=True,
+            on_epoch=False,
+            logger=True,
+        )
+        self.log("trn_loss", loss_recon.item(), prog_bar=True)
+        return loss
+
+    def on_train_epoch_end(self):
+        if (self.current_epoch + 1) % self.cmap_loss_ann_per_epoch == 0:
+            self.loss_criterion.update_kld_weight()
+
+    def validation_step(self, batch, batch_idx):
+        input_pc, target_map, target_cmap, target_gcs = self.get_inputs(batch)
+        _, pred_map, means, logvars, z_latent_code = self.model(input_pc, target_map)
+
+        loss_cmap, loss_recon, loss_kld = self.loss_criterion(
+            means, logvars, target_map, pred_map, target_cmap
+        )
+
+        loss = loss_cmap
+
+        loss_dict = {
+            "loss/val/overall": loss.item(),
+            "loss/val/cmap_rec": loss_recon.item(),
+            "loss/val/cmap_kld": loss_kld.item(),
+        }
+        self.log_dict(
+            loss_dict,
+            logger=True,
+            on_epoch=True,
+            on_step=False,
+            sync_dist=True,
+        )
+
+    def test_step(self, batch, batch_idx):
+        input_pc, target_map, target_cmap, target_gcs = self.get_inputs(batch)
+        _, pred_map, means, logvars, z_latent_code = self.model(input_pc, target_map)
+        loss_cmap, loss_recon, loss_kld = self.loss_criterion(
+            means, logvars, target_map, pred_map
+        )
+        loss = loss_cmap
+
+        loss_dict = {
+            "loss/tst/overall": loss,
+            "loss/tst/cmap_rec": loss_recon,
+            "loss/tst/cmap_kld": loss_kld,
+        }
+        self.log_dict(loss_dict, logger=True, sync_dist=True)
+
+    def predict_step(self, batch, batch_idx):
+        input_pc, _, _, _ = self.get_inputs(batch)
+        return self.model.predict(input_pc)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, betas=(0.9, 0.999))
+        lr_scheduler = torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=self.decay_lr_freq, gamma=0.2
+        )
+        return [optimizer], [lr_scheduler]

robotfingerprint/model/loss.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class GcsLoss(nn.Module):
+    def __init__(
+        self,
+        lw_init_recon: float,
+        lw_init_kld: float,
+        ann_temp: float,
+        ann_per_epochs: int = 1,
+        attn_weight: float = 3,
+    ):
+        super(GcsLoss, self).__init__()
+        self.lw_recon = lw_init_recon
+        self.lw_kld = lw_init_kld
+        self.ann_temp = ann_temp
+        self.iter_counter = 0
+        self.ann_per_epochs = ann_per_epochs
+        self.use_attn = attn_weight > 1
+        self.attn_alpha = attn_weight
+
+    def forward(self, means, logvars, uv_coords_gt, uv_coords_pred, cmap_gt=None):
+        """
+        :param means:
+        :param logvars:
+        :param uv_coords_gt: (B, N, 2+)
+        :param uv_coords_pred: (B, N, 2+)
+
+        :return:
+            loss, loss_reconstruction, loss_KL_Divergence
+        """
+        loss_kld = (
+            -0.5 * torch.sum(1 + logvars - means.pow(2) - logvars.exp(), dim=-1).mean()
+        )
+
+        if self.use_attn and cmap_gt is not None:
+            square_error = torch.square(uv_coords_gt - uv_coords_pred).sum(dim=-1)
+            attention_weights = torch.exp(cmap_gt.squeeze(-1) * self.attn_alpha)
+            square_error = square_error * attention_weights
+            # now we (B, N) sum-> (B, ) div-> (B,) .mean() -> (scalar)
+            loss_recon = torch.sqrt(
+                (square_error.sum(dim=-1) / attention_weights.sum(dim=-1))
+            ).mean()
+        else:
+            loss_recon = torch.sqrt(
+                torch.square(uv_coords_gt - uv_coords_pred).sum(dim=-1)
+            ).mean()
+
+        loss = self.lw_kld * loss_kld + self.lw_recon * loss_recon
+        return loss, loss_recon, loss_kld
+
+    @staticmethod
+    def metrics(self, uv_coords_gt, uv_coords_pred):
+        return F.smooth_l1_loss(uv_coords_gt, uv_coords_pred, reduction="mean")
+
+    def update_kld_weight(self):
+        self.lw_kld *= self.ann_temp
+
+    def apply_iter(self):
+        # NOT USED with Pytorch Lightning mode
+        if self.iter_counter % self.ann_per_epochs == self.ann_per_epochs - 1:
+            self.lw_kld *= self.ann_temp
+        self.iter_counter += 1

robotfingerprint/model/modules.py

@@ -0,0 +1,289 @@
+import torch
+import torch.nn as nn
+
+class PointNetCmapEncoder(nn.Module):
+    def __init__(self, layers_size=[4, 64, 128, 512]):
+        super(PointNetCmapEncoder, self).__init__()
+        self.layers_size = layers_size
+        self.conv_layers = nn.ModuleList()
+        self.bn_layers = nn.ModuleList()
+        self.activate_func = nn.ReLU()
+
+        for i in range(len(layers_size) - 1):
+            self.conv_layers.append(nn.Conv1d(layers_size[i], layers_size[i + 1], 1))
+            self.bn_layers.append(nn.BatchNorm1d(layers_size[i + 1]))
+            nn.init.xavier_normal_(self.conv_layers[-1].weight)
+
+    def forward(self, x):
+        # input: B * N * 4
+        # output: B * latent_size
+        x = x.transpose(1, 2)
+        for i in range(len(self.conv_layers) - 1):
+            x = self.conv_layers[i](x)
+            x = self.bn_layers[i](x)
+            x = self.activate_func(x)
+        x = self.bn_layers[-1](self.conv_layers[-1](x))
+        x = torch.max(x, 2, keepdim=True)[0]
+        x = x.view(-1, self.layers_size[-1])
+        return x
+
+
+class PointNetCmapDecoder(nn.Module):
+    def __init__(
+        self,
+        global_feat_size=512,
+        latent_size=128,
+        pointwise_layers_size=[3, 64, 64],
+        global_layers_size=[64, 128, 512],
+        decoder_layers_size=[64 + 512 + 128, 512, 64, 64, 1],
+    ):
+        super(PointNetCmapDecoder, self).__init__()
+        assert global_feat_size == global_layers_size[-1]
+        assert (
+            decoder_layers_size[0]
+            == latent_size + global_feat_size + pointwise_layers_size[-1]
+        )
+
+        self.global_feat_size = global_feat_size
+        self.latent_size = latent_size
+        self.pointwise_layers_size = pointwise_layers_size
+        self.global_layers_size = global_layers_size
+        self.decoder_layers_size = decoder_layers_size
+
+        self.pointwise_conv_layers = nn.ModuleList()
+        self.pointwise_bn_layers = nn.ModuleList()
+        self.global_conv_layers = nn.ModuleList()
+        self.global_bn_layers = nn.ModuleList()
+        self.activate_func = nn.ReLU()
+
+        for i in range(len(pointwise_layers_size) - 1):
+            self.pointwise_conv_layers.append(
+                nn.Conv1d(pointwise_layers_size[i], pointwise_layers_size[i + 1], 1)
+            )
+            self.pointwise_bn_layers.append(
+                nn.BatchNorm1d(pointwise_layers_size[i + 1])
+            )
+            nn.init.xavier_normal_(self.pointwise_conv_layers[-1].weight)
+
+        for i in range(len(global_layers_size) - 1):
+            self.global_conv_layers.append(
+                nn.Conv1d(global_layers_size[i], global_layers_size[i + 1], 1)
+            )
+            self.global_bn_layers.append(nn.BatchNorm1d(global_layers_size[i + 1]))
+            nn.init.xavier_normal_(self.global_conv_layers[-1].weight)
+
+        self.decoder_conv_layers = nn.ModuleList()
+        self.decoder_bn_layers = nn.ModuleList()
+        self.sigmoid = nn.Sigmoid()
+
+        for i in range(len(decoder_layers_size) - 1):
+            self.decoder_conv_layers.append(
+                nn.Conv1d(decoder_layers_size[i], decoder_layers_size[i + 1], 1)
+            )
+            self.decoder_bn_layers.append(nn.BatchNorm1d(decoder_layers_size[i + 1]))
+            nn.init.xavier_normal_(self.decoder_conv_layers[-1].weight)
+
+        # self.h2_decoder_conv_layers = nn.ModuleList()
+        # self.h2_decoder_bn_layers = nn.ModuleList()
+        # for i in range(len(decoder_layers_size) - 1):
+        #     self.h2_decoder_conv_layers.append(
+        #         nn.Conv1d(decoder_layers_size[i], decoder_layers_size[i + 1], 1)
+        #     )
+        #     self.h2_decoder_bn_layers.append(nn.BatchNorm1d(decoder_layers_size[i + 1]))
+        #     nn.init.xavier_normal_(self.h2_decoder_conv_layers[-1].weight)
+
+    def forward(self, x, z_latent_code):
+        """
+        :param x: B x N x 3
+        :param z_latent_code: B x latent_size
+        :return:
+        """
+        bs = x.shape[0]
+        npts = x.shape[1]
+
+        pointwise_feature = x.transpose(1, 2)
+        for i in range(len(self.pointwise_conv_layers) - 1):
+            pointwise_feature = self.pointwise_conv_layers[i](pointwise_feature)
+            pointwise_feature = self.pointwise_bn_layers[i](pointwise_feature)
+            pointwise_feature = self.activate_func(pointwise_feature)
+        pointwise_feature = self.pointwise_bn_layers[-1](
+            self.pointwise_conv_layers[-1](pointwise_feature)
+        )
+
+        global_feature = pointwise_feature.clone()
+        for i in range(len(self.global_conv_layers) - 1):
+            global_feature = self.global_conv_layers[i](global_feature)
+            global_feature = self.global_bn_layers[i](global_feature)
+            global_feature = self.activate_func(global_feature)
+        global_feature = self.global_bn_layers[-1](
+            self.global_conv_layers[-1](global_feature)
+        )
+        global_feature = torch.max(global_feature, 2, keepdim=True)[0]
+        global_feature = global_feature.view(bs, self.global_feat_size)
+
+        global_feature = torch.cat([global_feature, z_latent_code], dim=1)
+        global_feature = global_feature.view(
+            bs, self.global_feat_size + self.latent_size, 1
+        ).repeat(1, 1, npts)
+        pointwise_feature = torch.cat([pointwise_feature, global_feature], dim=1)
+        # pointwise_feature_h2 = pointwise_feature.clone()
+        for i in range(len(self.decoder_conv_layers) - 1):
+            pointwise_feature = self.decoder_conv_layers[i](pointwise_feature)
+            pointwise_feature = self.decoder_bn_layers[i](pointwise_feature)
+            pointwise_feature = self.activate_func(pointwise_feature)
+        pointwise_feature = self.decoder_bn_layers[-1](
+            self.decoder_conv_layers[-1](pointwise_feature)
+        )
+
+        # for i in range(len(self.h2_decoder_conv_layers) - 1):
+        #     pointwise_feature_h2 = self.h2_decoder_conv_layers[i](pointwise_feature_h2)
+        #     pointwise_feature_h2 = self.h2_decoder_bn_layers[i](pointwise_feature_h2)
+        #     pointwise_feature_h2 = self.activate_func(pointwise_feature_h2)
+        # pointwise_feature_h2 = self.h2_decoder_bn_layers[-1](
+        #     self.h2_decoder_conv_layers[-1](pointwise_feature_h2)
+        # )
+
+        ### pointwise_feature shape B x out_size x N
+        # pointwise_feature = (
+        #     self.sigmoid(pointwise_feature).view(bs, npts, -1).squeeze(-1)
+        # )
+        # Keep this without sigmoid, since we might do additional transforms
+        # return pointwise_feature.view(bs, npts, -1).squeeze(-1)
+        return pointwise_feature  # shape (bs, -1, npts)
+        # output =  torch.cat((pointwise_feature, pointwise_feature_h2), dim=1).view(bs, npts, -1)
+        # return output
+        # return self.sigmoid(output)
+
+
+class GcsCVAE(nn.Module):
+    def __init__(
+        self,
+        latent_size=128,
+        encoder_layers_size=[5, 64, 128, 512],
+        decoder_global_feat_size=512,
+        decoder_pointwise_layers_size=[3, 64, 64],
+        decoder_global_layers_size=[64, 128, 512],
+        decoder_decoder_layers_size=[64 + 512 + 128, 512, 64, 64, 1],
+        num_coarse_pts=2048,
+        uv_layers_size=None,
+    ):
+        # NOTE:
+        # encoder_layers_size[0]  is 5 instead of 4 since our input is (obj_pc, obj_uv_coords)
+        super(GcsCVAE, self).__init__()
+        self.num_coarse = num_coarse_pts
+        self.latent_size = latent_size
+
+        self.cmap_encoder = PointNetCmapEncoder(layers_size=encoder_layers_size)
+        self.cmap_decoder = PointNetCmapDecoder(
+            latent_size=latent_size,
+            global_feat_size=decoder_global_feat_size,
+            pointwise_layers_size=decoder_pointwise_layers_size,
+            global_layers_size=decoder_global_layers_size,
+            decoder_layers_size=decoder_decoder_layers_size,
+        )
+
+        self.encoder_z_means = nn.Linear(encoder_layers_size[-1], latent_size)
+        self.encoder_z_logvars = nn.Linear(encoder_layers_size[-1], latent_size)
+
+        self.pred_uv = False
+        if uv_layers_size:
+            # uv_layers_size = [64, 64, 1]
+            self.pred_uv = True
+            num_layers = len(uv_layers_size) - 1
+
+            self._pred_u_layers = nn.ModuleList()
+            self._pred_v_layers = nn.ModuleList()
+            for i in range(num_layers):
+                curr_size, next_size = uv_layers_size[i], uv_layers_size[i + 1]
+
+                self._pred_u_layers.append(nn.Conv1d(curr_size, next_size, 1))
+                nn.init.xavier_normal_(self._pred_u_layers[-1].weight)
+                self._pred_u_layers.append(nn.BatchNorm1d(next_size))
+
+                self._pred_v_layers.append(nn.Conv1d(curr_size, next_size, 1))
+                nn.init.xavier_normal_(self._pred_v_layers[-1].weight)
+                self._pred_v_layers.append(nn.BatchNorm1d(next_size))
+
+                if i < num_layers:
+                    self._pred_u_layers.append(nn.ReLU())
+                    self._pred_v_layers.append(nn.ReLU())
+
+            self.pred_u_net = nn.Sequential(*self._pred_u_layers)
+            self.pred_v_net = nn.Sequential(*self._pred_v_layers)
+
+    def forward(self, obj_pts, gcs_gt):
+        """
+        :param obj_pts: B, N, 3
+        :param gcs_gt: B, N, 2
+        :return:
+        """
+
+        bs = obj_pts.shape[0]
+        npts = obj_pts.shape[1]
+        obj_cmap = torch.cat(
+            (
+                obj_pts,
+                gcs_gt.unsqueeze(-1) if gcs_gt.ndim == 2 else gcs_gt,
+            ),
+            dim=-1,
+        )
+        means, logvars = self.forward_encoder(object_cmap=obj_cmap)
+        z_latent_code = self.reparameterize(means=means, logvars=logvars)
+        cmap_values = self.forward_decoder(obj_pts, z_latent_code).view(bs, npts, -1)
+        return obj_pts, cmap_values, means, logvars, z_latent_code
+
+    def predict(self, object_pts):
+        """
+        Test time prediction of contact maps from randomly sampled latent vectors on a given input
+
+        Input:
+         object_pts: (B, N, 3) tensor
+
+        Returns:
+         gcs_values: (B, N, 2) tensor of contact map values for each batched obj pc
+        """
+        bsize = object_pts.shape[0]
+        z_samples = torch.randn(
+            bsize, self.latent_size, device=object_pts.device
+        ).float()
+        return self.inference(object_pts, z_samples)
+
+    def inference(self, object_pts, z_latent_code):
+        """
+        :param object_pts: B x N x 3
+        :param z_latent_code: B x latent_size
+        :return:
+        """
+        cmap_values = self.forward_decoder(object_pts, z_latent_code)
+        return cmap_values
+
+    def reparameterize(self, means, logvars):
+        std = torch.exp(0.5 * logvars)
+        eps = torch.randn_like(std)
+        return means + eps * std
+
+    def forward_encoder(self, object_cmap):
+        cmap_feat = self.cmap_encoder(object_cmap)
+        means = self.encoder_z_means(cmap_feat)
+        logvars = self.encoder_z_logvars(cmap_feat)
+        return means, logvars
+
+    def forward_decoder(self, object_pts, z_latent_code):
+        """
+        :param object_pts: B x N x 3
+        :param z_latent_code: B x latent_size
+        :return:
+        """
+        cmap_values = self.cmap_decoder(object_pts, z_latent_code)
+        bs = cmap_values.shape[0]
+        npts = cmap_values.shape[-1]
+        if self.pred_uv:
+            u = self.pred_u_net(cmap_values).view(bs, npts, -1)
+            v = self.pred_v_net(cmap_values).view(bs, npts, -1)
+            # u = torch.sigmoid(u)
+            # v = torch.sigmoid(v)
+            output = torch.cat([u, v], dim=-1).view(bs, npts, -1)
+            return output
+        else:
+            return cmap_values.view(bs, npts, -1)
+

robotfingerprint/unseen_barrett/final.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1e725f0e6bf613e55b88ddd4e862c01486bf2e9c1e73ca12bdb7f2def1d3314
+size 11523772

robotfingerprint/unseen_ezgripper/final.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd50763647da294b6e988fdd4c4ae691b7e3ccf4c0bfc29991c9ab82cdd66dce
+size 11523772

robotfingerprint/unseen_shadowhand/final.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db094eaeb2b39e4872ca229c1101e6afdfb0f8398feec9a2997b9077a4043a52
+size 11523772