Document seed-conditioned models and results

README.md

@@ -102,3 +102,83 @@ Note: `me-full_25000_*` checkpoints were trained with flax 0.11 using a per-path
 `nnx.Param` layout in `TPWithWeightsAndBiases` (`kin_flow/net/module/fctp.py`) —
 restore with a matching code state; they are not compatible with the original
 Param-of-list layout.
+
+---
+
+# Seed-conditioned models
+
+Extension of the multi-embodiment model with an **equivariant seed point** — a 3D point
+marking *which* object to grasp — after SeededGrasp ([arXiv 2607.20207](https://arxiv.org/abs/2607.20207)),
+expressed in irreps so SE(3) equivariance is preserved by construction.
+
+These are fine-tuned from `me-full_25000_120`. The seed enters the scene encoder as one
+virtual graph edge from every node to the seed point, contributing an ℓ=1 direction vector
+and ℓ=0 distance — the two quantities the encoder already consumes. Only ~10k new
+parameters (`unet.seed_mlp`) on an 18.8M-parameter model.
+
+## Models
+
+| Folder | Fine-tune | Epoch | Best for | seeded `targeted` |
+|---|---|---|---|---|
+| `me-seed-anneal_25000_4` | annealed LR | 4 | **best overall** · Panda · Allegro | 0.4582 (Panda) |
+| `me-seed-anneal_25000_6` | annealed LR | 6 | VX300 | 0.4339 |
+| `me-seed-anneal_25000_8` | annealed LR | 8 | DexEE | 0.1599 |
+| `me-seed_25000_2` | constant LR | 2 | Shadow Hand | 0.3502 |
+
+`me-seed-anneal_25000_4` has the highest mean across all five grippers (0.3289) and is the
+recommended default. The `-anneal` runs use a learning rate annealed 8e-5 → 0; the original
+schedule held a constant 8e-5 (`decay_steps` was 10,000 against a 200k-step run) which
+degraded results over training.
+
+## Evaluation
+
+`targeted = yield × hit × SR|target` — successful grasps on the *intended* object per raw
+sample drawn. Measured on 119 held-out cluttered scenes (4 objects each), 100 grasps per
+scene, both conditions scored on identical grasps (same RNG, same seed point, differing only
+in whether the seed mask is on).
+
+| Gripper | DoF | unseeded | seeded | + collision projection | total gain |
+|---|---|---|---|---|---|
+| Panda | 2 | 0.2067 | 0.4611 | **0.5190** | +151% |
+| VX300 | 2 | 0.2027 | 0.4410 | **0.4752** | +134% |
+| DexEE | 12 | 0.0914 | 0.1578 | **0.1597** | +75% |
+| Allegro | 16 | 0.1261 | 0.2601 | **0.3220** | +155% |
+| Shadow Hand | 22 | 0.1510 | 0.3060 | **0.3813** | +152% |
+| **mean** | | **0.1556** | **0.3252** | **0.3714** | **+139%** |
+
+Seed conditioning raises the fraction of grasps aimed at the intended object from chance
+(~0.28 for ~3.5 objects/scene) to 0.61–0.79. Collision projection is a separate, training-free
+post-processing step that recovers the collision-free yield seeding costs.
+
+**Supporting numbers** — 46 evaluations across 15 checkpoints and 5 grippers. Seeded and
+unseeded results are identical to 4 decimal places on the pre-fine-tune checkpoint, where the
+seed pathway is zero-initialised, confirming the conditioning is inert at initialisation.
+
+## Caveats
+
+- Test clouds are **analytically sampled** from MuJoCo geometry, not rendered scans (the
+  renderer segfaults on our hardware). Validated to 2 mm geometric agreement but carrying a
+  measured −2.7 pp systematic offset — valid for A/B comparison, **not** directly comparable
+  to the `me-full_25000_*` numbers above.
+- Performance saturates within ~4,000 steps; epochs 2–8 are statistically indistinguishable
+  and longer training does not help.
+- DexEE is an outlier (+75%, `SR|target` 0.26 vs 0.55–0.77 elsewhere): it aims correctly but
+  still fails to hold objects.
+
+## Loading
+
+Same as the `me-full_*` checkpoints, but the model must be built with seed conditioning so
+`unet.seed_mlp` exists:
+
+```python
+# model config needs seed_zero_init; then
+model = Trainer.get_model_from_checkpoint(model, "<path>/me-seed-anneal_25000_4",
+                                          partial=True)
+
+# inference: seed is (xyz, mask) in world metres; mask=0 => unconditional
+se3, dof = inference("Flow", model, sample, num_samples, cfg,
+                     seed=(seed_xyz, np.float32(1.0)))
+```
+
+`partial=True` restores only the parameters present in a checkpoint, which is also how these
+were created from the unconditional model.