Add a learned grasp-yaw demo with differentiable IK

[yuecideng]

Description

Train a conditional grasp-yaw predictor from downstream joint-space cost through DifferentiableIKLayer. A frozen network then chooses an angle for new current-joint/target-position pairs using one network evaluation and one fixed-step IK solve. Grasp position and downward approach remain fixed.

The example uses the learning layers already merged in #806. It adds a training/evaluation runner and a frozen-model Viser viewer. Numerical angle optimization and grid search are integrated reference baselines. Generated results, checkpoints, screenshots and validation artifacts are kept out of Git; the PR branch has been rebuilt without their historical versions. No production API changes.

The comparison includes current heading, one learned absolute angle, an IK-trained MLP, a teacher-supervised MLP with the same initialization/minibatches, per-query Adam optimization, and a 129-angle grid. The IK-trained method does not consume teacher labels; label generation is charged separately to the supervised comparison.

Evaluation

Three training seeds; 1,024 training tasks; 128 independent validation tasks; 128 cases in each held-out split. Only validation task loss selects checkpoints. Training uses 24 IK steps and all final methods use the same 64-step budget, current-joint seed, and strict physical-limit/pose checks.

Held-out split	Learned mean paired cost reduction vs heading	Strict IK success
IID	3.85%	128/128 for all three seeds
Disjoint Cartesian-Y region	3.57%	128/128 for all three seeds
Shifted initial joints	16.82%	128/128 for all three seeds

Average cost gaps to the sampled grid are below 0.2% for each split/seed. The teacher-supervised network performs similarly; this does not establish superiority of end-to-end training. Per-target failures and paired-valid counts are retained in the full run artifacts.

On the RTX PRO 5000 72GB, each IK-trained model took about 28–31 seconds. Supervised fitting took about 0.5 seconds plus 30 seconds to generate training labels (shared across its runs). At batch 64, complete learned inference measured 86.9 ms, fully batched grid search 86.2 ms, and per-query gradient optimization 7.09 s. The grid uses a separate 16,384-candidate inference chunk to avoid an artificial small-batch penalty. The learned method does not outperform the one-dimensional grid in latency here. Incremental Torch peak memory was about 1.03 MB versus 37.74 MB, respectively.

These are local kinematic tasks. Costs measure joint endpoint displacement; collision, contact, dynamics, executable trajectories, and global optimality are not validated.

Preview and usage

conda activate py311
pip install -r examples/learning/differentiable_ik/requirements.txt
python -m examples.learning.differentiable_ik.learn_grasp_yaw --device cuda
python -m examples.learning.differentiable_ik.view_learned_yaw \
  --run outputs/learn_grasp_yaw --port 8101

The viewer supports held-out/model selection, current-posture inspection, automatic frozen inference after moving the cylinder, and optional online reference search. Training weights never change in the viewer. Stale work is discarded after inputs change.

To reproduce training and evaluation:

python -m examples.learning.differentiable_ik.learn_grasp_yaw --device cuda

The README describes data distributions, baselines, CPU smoke commands, timing scope and artifacts. All reports, plots, checkpoints, replay data and validation results are generated locally under the selected output directory. Run training before opening the viewer; no pretrained bundle is committed.

Type of change

• New feature: standalone learning example and interactive demonstration.

Validation

All checks use conda py311:

• Complete CUDA experiment with three training seeds and all six methods; full CPU smoke run.
• Float64 gradcheck through the complete 24-step policy/IK objective for current joints and target position; independent finite differences for a network parameter. All checked parameter/input gradients are finite and nonzero.
• CPU checkpoint prediction parity, deterministic frozen predictions, preserved input gradients, unchanged input arrays, independent FK success/pose-invariant checks, disjoint/unique datasets, and unreachable/reference-fallback checks.
• Explicit label-free fitting and rejection of missing labels for supervised training.
• Browser checks: model/split selection, current-posture toggle, automatic prediction, optional reference search, reset, stale-result rejection, and no page errors. The concurrency check uses a longer 200-update reference to keep work in flight; saved evaluation and timing use 50 updates.
• Viewer smoke against locally generated output: actual meshes, frozen predictions, unchanged weights and HTTP lifecycle.
• Existing compute/layer/shared-model/context regressions: 71 passed, 4 GPU-specific tests skipped in that CPU invocation. The example itself was trained and evaluated on CUDA.
• Black 26.3.1 (black .), compile checks, git diff --check, API coverage (2628/2628) and context checks passed. Example validation is executable; no tutorial-importing pytest modules were added.

Checklist

• Format the repository with Black 26.3.1.
• Update the guide, README and learning-layers context.
• Pass public API coverage; no new production exports.
• Verify learning, gradients, persistence, numerical correctness and interaction.
• Document example-only dependencies and measured limitations.

#806 is merged. This branch is rebuilt directly on main, so its diff contains only the example and related documentation/artifacts.

Artifact cleanup: the final PR contains eight source/configuration/documentation files and no new test files or generated artifacts. Its branch history is a single commit on the existing main baseline.

[greptile-apps]

Confidence Score: 5/5

[Low impact] The PR appears safe to merge, with a non-blocking gap in saved online-search failure records.

Findings

1.  Search failure flags are lost ▶

Fix with agent prompt

### Issue 1
examples/learning/differentiable_ik/learn_grasp_yaw.py:739
`run` saves `online_found` only for the first eight cases per split in `replay.npz`. The remaining flags are lost because `paired_results.npz` contains only the final 64-step results. A heading fallback can pass that longer solve, so final success cannot reveal that the online search failed. Save `online_found` for every query alongside the paired results.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Adds an example-local grasp-yaw predictor trained through differentiable IK.

• The learning runner trains a network to choose grasp yaw through IK.
• The viewer compares frozen grasp-yaw predictions with heading and references.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Current joints and grasp position] --> B[Yaw predictor]
  B --> C[Fixed-position downward target]
  C --> D[Differentiable IK]
  D --> E[Joint displacement and pose loss]
  E -->|Training only| B
  B --> F[Frozen viewer inference]

Loading

_{Reviews (7) · Last reviewed commit: "feat(examples): learn grasp yaw through ..." · Reviewed by Greptile}

[greptile-apps]

Search failure flags are lost

run saves online_found only for the first eight cases per split in replay.npz. The remaining flags are lost because paired_results.npz contains only the final 64-step results. A heading fallback can pass that longer solve, so final success cannot reveal that the online search failed. Save online_found for every query alongside the paired results.

Prompt To Fix With AI

This is a comment left during a code review.
Path: examples/learning/differentiable_ik/learn_grasp_yaw.py
Line: 739

Comment:
**Search failure flags are lost**

`run` saves `online_found` only for the first eight cases per split in `replay.npz`. The remaining flags are lost because `paired_results.npz` contains only the final 64-step results. A heading fallback can pass that longer solve, so final success cannot reveal that the online search failed. Save `online_found` for every query alongside the paired results.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.