Claude Code audit & bug scan (3D keypoints, style transfer, IK)

[sibocw]

🤖 Generated by Claude Code (Opus 4.8). Automated audit + bug scan of the research code (production orchestration intentionally out of scope; shared scientific logic it calls into was audited). Findings produced by reading the code and fanning out review subagents; each item is spot-verified against source. Status lines reflect a maintainer triage pass.

Revision 2 (2026-06-23): reconciled against a second, independent multi-agent audit pass and verified key claims empirically on the real atomic-batch data (bulk_data/pose_estimation/atomic_batches/4variants/*/*_labels.h5). Corrections/additions are tagged [rev2]. Net effect: one finding upgraded (I1-A is an active, ~10–18° IK error — not latent — but belongs under Issue 3, not Issue 1), two downgraded with measurements (I1-C, I1-D), one factual fix (I2-B), and several confirmations + one new footgun (I1-G).

Scope & method

Audited the 3D-keypoints model, shared backbone, body-seg model (for contrast), camera unprojection, inference paths, synthetic-data pipeline, style-transfer module, and inverse kinematics. [rev2] additionally measured label statistics (depth range, OOB-drop rate, 2D-label orientation) and numerically reproduced the camera/coordinate math on real labels. Three reported symptoms drove the audit:

1. 3D keypoints model performs worse than the body-seg model (surprising).
2. Higher-resolution style-transferred images did not improve pose predictions.
3. Inverse kinematics sometimes gives incorrect results.

Status legend: 🔧 fix · 🔎 investigate · 📝 TODO (separate) · ✅ acknowledged/won't-change-now · ❓ pending confirmation

---

Root-cause summary [rev2]

• Issue 1 — depth is the weak axis, two compounding reasons. Architecturally the depth head global-average-pools away all spatial info (I1-B). Geometrically depth is nearly unobservable: measured per-frame depth spread ≈2.8 mm at ~100 mm through a 5° FOV (near-orthographic), std 0.38 mm, occupying ~25–42 of 64 bins. With focal_scaling≈5314, 1 px lateral = 0.019 mm but 1 depth-bin = 0.064 mm — to match 1 px of lateral precision, depth must be right to 0.3 of a bin; a whole leg segment (~0.4 mm) spans only ~6 depth bins. So "3D" = a healthy 2D heatmap + a barely-observable z predicted by the weakest head. (The x-y pathway and 2D label convention are clean — verified.)
• Issue 2 — resolution is pinned at 256-in / 128-heatmap in three independent places (style-transfer GAN at 256, dataloader crop-not-resize to 256, heatmap = input/2) and stored in lossy H.264 4:2:0 / CRF-15 MP4. Any one defeats higher resolution; the global-pool depth head is resolution-invariant on top.
• Issue 3 — IK is handed bad targets, then amplifies them. (a) I1-A systematically distorts the targets (~10–18° angle error, measured, not removed by AlignPose); (b) noisy depth (Issue 1) distorts z; then (c) warm-start propagation, non-mirrored bounds, NaN aborts, and redundant-DOF seed-dependence add the intermittency ("sometimes").

---

Cross-cutting context

• Resolution coupling is implicit and unguarded across extraction (target_image_size), training (input_image_size=256), heatmap stride (input/2), style-transfer (image_side_length), and unprojection (camera_rendering_size). Nothing asserts these agree. Several findings below are facets of this.
• Scripts currently broken on claude-audit, independent of the science: I3-A (dof_name_lookup_canonical_to_nmf), the inference c= kwarg, and I2-C — fix first so paths run at all.
• [rev2] Package import is broken in a clean checkout: pose/data/synthetic/atomic_batch.py:8 imports the private _default_ffmpeg_params_for_video_writing from pvio.io, which the installed pvio no longer exports → ImportError chains through pose/data/synthetic/__init__.py and pose/contrast/__init__.py, making the data + contrast packages unimportable. Depend on a public pvio API (or pass ffmpeg params explicitly) and pin pvio.

---

Issue 1 — 3D keypoints worse than body-seg

I1-B — 🔧 Depth head global-average-pools away all spatial information (primary cause)

Where: pose/keypoints3d/model.py:356,376 — nn.Sequential(nn.AdaptiveAvgPool2d(1), conv, …, fc, reshape).

The depth head collapses the 128×128 decoder map to 1×1 before the conv, then one FC regresses depth bins for all 32 keypoints from a single global 64-d vector — no per-keypoint spatial grounding. The xy pathway (heatmaps) and the body-seg head are fully spatial; depth is the asymmetric weak link, consistent with "3D worse than seg." [rev2] This is compounded by the geometric near-unobservability of depth (see Root-cause summary): even a perfect head must resolve sub-0.1 mm differences. Fix: sample features at each keypoint's predicted location (gather d0 at the soft-argmax xy, or heatmap-weighted pooling per keypoint) before the depth FC. (Maintainer: agreed this is a problem.)

I1-D — 🔎 Depth range is narrow; whole-sample dropping is real but rarely fires

Where: production_models/keypoints3d/v2/configs/model_architecture_config.yaml (depth_min:-102, depth_max:-98); OOB mask is .any(dim=1) over keypoints (model.py:931-943) + oob_treatment: drop.

The range is fine at inference (maintainer confirmed). During training, any frame with one keypoint outside [−102,−98] drops the entire frame from the loss (.any(dim=1)). [rev2] Measured: with the production window only 0.18% of frames are dropped (and 0% at the far edge) — so it does not starve the depth head as originally feared. The genuine problem is the opposite: depth has almost no dynamic range (≈2.8 mm spread, std 0.38 mm), so ~⅓ of the 64 bins (−99.5 → −98) are never used and the window is mis-centered. Make training more robust + use the bins: per-keypoint masking instead of whole-sample drop; re-center/narrow the bin range to the occupied span (or fewer bins); log the per-epoch drop fraction.

I1-G — 🔧 [rev2] new Default depth window is disjoint from the data → silent zero-loss training

Where: pose/keypoints3d/config.py:16,18 defaults depth_min=-70, depth_max=-63.

Real keypoint_pos[...,2] is [-102.3, -99.5] mm (measured) — no overlap with the [-70,-63] default. Any run that doesn't override it has every keypoint OOB → oob_treatment:"drop" drops every sample → total loss returns 0 (model.py:983-987,1053-1064), i.e. training silently does nothing. Production yamls override to [-102,-98] so shipped models are safe, but this is a landmine. Fix: delete/replace the default with a data-derived range, and assert at startup that the configured window brackets the observed label range (one assert also covers I1-D).

I1-C — 🔎 Train loader crops to input_image_size instead of resizing; labels not rescaled — latent

Where: pose/data/synthetic/atomic_batch.py:277 (selection = frame[:n_rows, start_col:end_col]); assert at :250-252 only catches too-narrow tiles. Extraction target_image_size defaults to None ("save at source resolution"; preextract_atomic_batches.py:21-22,63-68, in-file example uses original_image_size=(464,464)).

If atomic batches were extracted at 464 without target_image_size, the loader keeps only the top-left 256×256 while keypoint labels stay in 464-space → most go OOB → dropped → poisoned subset. [rev2] Resolved for current data: the on-disk 4variants tiles are measured at 256×256 (frames are 256-tall; keypoint_pos xy ∈ [0,256)), so the crop is a no-op today — this is latent, not active. It remains (a) a footgun if anyone re-extracts at higher resolution and (b) the reason Issue 2 can't work via the loader. Fix: assert stored tile size == input_image_size and fail loudly, or actually resize the crop and scale keypoint_pos[...,:2] by the same factor.

I1-E — ✅ Keypoint heatmaps at half input resolution (acknowledged, not changing now)

Where: pose/keypoints3d/model.py:193-206 (heatmaps off d0 at input/2 = 128) vs pose/bodyseg/model.py:56-62 (a final_upsampler to full 256). Caps xy precision relative to seg; noted, deliberately deferred. [rev2] Compounds with the σ=2 heatmap-px Gaussian target (config.py:92, = 4 input-px) and the sparse-vs-dense supervision gap (seg supervises every pixel; heatmaps put signal on ~0.7% of pixels) — together these explain why a dense segmentation task is far better-conditioned than 2.5D keypoint regression from the same backbone.

I1-F — 📝 TODO (separate): confirm inference input range is [0,1]

Where: run_keypoints3d_inference.py:74-77 uses SimpleVideoCollectionLoader(..., transform=Resize(...)) (pvio). Backbone applies ImageNet norm assuming inputs already in [0,1] (pose/common.py:62-91); training divides by 255 (atomic_batch.py:285, verified to yield [0,1]). If the inference loader yields uint8/[0,255], normalization is wildly off. Verify the loader emits float [0,1].

---

Issue 2 — higher-resolution style transfer didn't help

I2-A — ❓ Resolution is frozen before the pose model sees it (multiple chokepoints)

(Maintainer: mark as potential problems, pending further confirmation.)

• Pose side: trainer/inference run at fixed 256 (data_config input_image_size:[256,256]; run_keypoints3d_inference.py:74); atomic-batch resolution is baked at extraction and then cropped not resized at load (I1-C). Higher-res sources are downsampled/cropped to 256 regardless.
• Style-transfer side (non-tiled): style_transfer/cut_inference.py:98-107 uses preprocess="resize_and_crop", load_size=crop_size=image_side_length (=256 default in run_inference.py:39) → every frame hard-resized to 256² and the generator emits 256². run_inference.py also ignores the model's trained crop_size/preprocess from train_options.json (unlike the tiled script).
• [rev2] Heatmap ceiling (= I1-E): even with a larger input, the heatmap is structurally input/2, so spatial precision is capped; and the depth head is resolution-invariant (global avg-pool, I1-B). [rev2] Storage: atomic batches are stored as lossy H.264 4:2:0 / CRF-15 MP4 (pvio.io defaults), subsampling chroma 2× — eroding the high-frequency detail extra resolution would add.

To actually benefit from higher resolution, raise it consistently end-to-end (target_image_size at extraction, input_image_size at training, plus a higher-res heatmap/final-upsampler and a spatially-grounded depth head, plus crop_size parsed from trained CUT options at inference, plus less-lossy storage).

I2-B — ✅/📝 Tiled high-res inference can inject seam/blend artifacts (results manually checked OK; reporting suggestion only)

Where: style_transfer/tiled_inference.py, run_tiled_inference.py.

• [rev2] Correction: the default is weight_type="uniform" (tiled_inference.py:203, run_tiled_inference.py:87) → an all-ones weight map (:38-39), i.e. a plain 50/50 average in overlaps with no feathering → potential hard seams where two independently InstanceNorm-normalized tiles meet.
• The non-default cosine (0.5−0.5cos(2π·t)) and pyramid (1−|2t−1|) windows (:51-65) are zero at the tile border → on a 50%-overlap grid the seam pixel gets ≈0 weight from one tile, so they don't blend at the seam either (and don't form a partition of unity). gaussian is nonzero at borders but also not unity-summing.
• Mirror-padded out-of-frame tiles (:99-102,118-128) contribute hallucinated content to genuine edge pixels.
• randomize_seams=True by default (run_tiled_inference.py:89) moves seams every frame → temporal flicker.

Suggestion (not urgent since outputs look OK): use a partition-of-unity feather window nonzero at the seam; zero-weight mirror-padded regions; default randomize_seams=False.

I2-C — 🔧 Broken/hardcoded extraction script

Where: style_transfer/scripts/extract_dataset_bodypart.py:421 calls list_nmf_simulations_and_num_frames(dir, filename) but the function (extract_dataset.py:58-60) requires a third num_workers arg → TypeError on launch (verified). Same file :165-167 hardcodes SCALING = 900.0/256.0 for spotlight keypoint coords — silently wrong if resolution changes. Fix both.

---

Issue 3 — inverse kinematics sometimes wrong

I1-A — 🔧 [rev2] confirmed active, relocated here pred_xy (256-px) unprojected with intrinsics built for 464-px renders → distorted IK targets

Where: pose/keypoints3d/scripts/run_keypoints3d_inference.py:32,74,101-110 (and the same pattern in production/spotlight/keypoints3d.py:38,66-67,132); intrinsics in pose/camera.py:57-63.

Inference resizes frames to inference_image_size=256 (so pred_xy ∈ [0,256), matching the 256-space training labels) but builds CameraToWorldMapper with camera_rendering_size=(464,464) and feeds the 256-space pred_xy into it — no rescale.

[rev2] Why it does NOT explain Issue 1: the trained model never touches the mapper, and the qualitative eval (test_keypoints3d_models.py:153-178) is self-referential — it builds the mapper from original_image_size and unprojects both predictions and labels with the same mapper, so any miscalibration cancels in the comparison. That's why "3D looks only a bit off" in eval.

[rev2] Why it IS a real Issue-3 bug (measured, not latent): the unprojection feeding IK uses the 464 mapper on 256 coords. The wrong principal point (231.5 vs 127.5) and focal (5314 vs 2932) produce an anisotropic distortion: in-plane x,y are scaled 0.55× while depth z is exact (verified numerically on real labels: x,y ratio 0.552, z ratio 1.000). The ~2 mm principal-point translation and the uniform in-plane scale are absorbed by AlignPose (per-leg coxa-translation + single isotropic scale_factor, seqikpy/alignment.py:467-485) — but AlignPose is angle-preserving, so the residual anisotropy passes straight into the joint angles. Measured on ground-truth keypoints, the injected joint-angle error is median ≈12° at CTr, ≈19° at FTi, ≈2.5° at TiTa (overall mean 11.5°; 65% of joints >5°, 49% >10°). The correct 256 mapper yields anatomically sane geometry (≈7 mm-wide, 2.5 mm-deep fly), confirming 256 is the right sensor size. (Supersedes the rev1 "latent / subtle / ❓" disposition — an experiment was run.)

Fix: build the mapper at the resolution pred_xy lives in (rendering_size=inference_image_size), or rescale pred_xy *= camera_rendering_size/inference_image_size first; then add a resolution-consistency assert. Cheap, high-leverage, and unblocks fair evaluation of the other IK items below.

I3-A — 🔧 dof_name_lookup_canonical_to_nmf referenced but deleted → IK save crashes

Where: pose/keypoints3d/scripts/run_inverse_kinematics.py:752 and production/spotlight/keypoints3d.py:263. _save_seqikpy_output is on the main IK path (run_inverse_kinematics.py:619, no try/except). [rev2] Confirmed undefined in neuromechfly/constants.py (only keypoint_name_lookup_* and joint_segments_nmf_to_canonical exist).

Git history: defined in 5f8b6ee, removed in 3c9449a ("compatibility with flygym v2"), call sites never updated → AttributeError at save on the current branch. (Working results must predate the refactor or come from the notebook.) Fix: restore an ordered canonical→NMF DOF mapping for the 7 DOFs (ThC_yaw, ThC_pitch, ThC_roll, CTr_pitch, CTr_roll, FTi_pitch, TiTa_pitch) in neuromechfly/constants.py, or replace the call with that explicit ordered list (must match the Angle_{leg}_{dof} keys seqikpy emits so DOF packing order is correct).

I3-B — 🔎 Joint bounds not left/right mirror-consistent (investigate)

Where: neuromechfly/constants.py:288-336 (author flagged with # ?). [rev2] Verified specifics:

• ThC_yaw: RM/RH = (-45°,45°) vs LM/LH = (-45°,90°) (:306,322) — not a mirror (R should be (-90°,45°)); right mid/hind clipped 45° tighter.
• CTr_pitch: RF and RM = (-270°,10°) vs their L counterparts (-180°,10°) (:293,309) — pitch should match L/R; -270° is physically implausible and admits wrapped/ambiguous solutions. (RH/LH are both -180, fine.)
• ThC_roll: mid/hind mirror correctly (R=(-180,10) ↔ L=(-10,180)), but front does not (RF=(-135°,10°) vs LF=(-10°,90°); mirror of L would be (-90°,10°)).

A too-tight bound clips the true angle on one side; a too-loose/implausible bound lets the solver lock onto a wrapped pose → intermittent, side/posture-dependent errors. Investigate the NMF model's true range of motion and derive R bounds as the proper mirror of L (R=(-L_hi,-L_lo) for roll/yaw; identical for pitch).

I3-C — 🔧 Warm-start propagates one bad solve across many frames (implement gate/reset)

Where: invkin.run_seqikpy (invkin.py:177-184) → seqikpy LegInvKinSeq seeds frame t with frame t−1's solution. One bad frame (noisy/occluded keypoint, near-singular extension, knee-flip) becomes the seed for the next → errors persist/drift. [rev2] With parallel_over_time=True there's an additional variant: each time-chunk's frame 0 re-seeds from the static nmf_initial_angles and overlapping chunks are linearly blended, so adjacent chunks can settle in different valid configurations of the redundant 3-DOF ThC and the blend produces a physically wrong transient at every boundary (and results vary with worker/chunk count — a direct "sometimes" mechanism). Position-only IK of the redundant ThC is inherently seed-dependent (left/right roll ambiguity → flips). Implement: detect high final residual or large frame-to-frame jumps and re-solve from nmf_initial_angles (and/or multi-start); seed each chunk from the neighbor's overlap solution (or run serial) instead of blending; optionally outlier-reject/low-pass the 3D input before IK. (Re-evaluate after I1-A is fixed — IK is currently handed distorted targets.)

I3-D — 🔧 NaN hard-assert aborts an entire recording (improve)

Where: invkin.py:37 — assert not np.isnan(data_block).any(). One NaN/occluded keypoint kills the whole trial with an opaque error. There is also no confidence gating — conf_xy/conf_depth are computed but never used to drop low-quality keypoints before length-estimation/IK. Improve: drop/interpolate NaN frames per leg (or solve only valid frames); gate on confidence.

I3-E — 🔧 align_fwdkin_xyz_to_rawpred_xyz indexes constrained array with the raw file's keypoint order (improve)

Where: run_inverse_kinematics.py:520-526 → invkin.py:189-265. Identical by construction today, but if the raw file's keypoint order ever differs, per-leg translations apply to the wrong keypoints silently. Improve: pass/use keypoints_order_constrained and assert it equals the raw order.

---

Also noted (low priority)

• Inference __main__ is broken: run_keypoints3d_inference.py:290 passes c=inference_image_size (no such kwarg → TypeError), and its __main__ reads config keys absent from production/spotlight/config.yaml.
• Norm inconsistency: DecoderBlock uses BatchNorm2d (pose/common.py:173,175) while heads use GroupNorm and the model docstring argues for GroupNorm (small batches). Effective batch is large here so not fatal; a sim-to-real running-stats liability worth aligning. [rev2] Shared identically by both keypoints and body-seg, so not a cause of the keypoints-vs-seg gap.
• [rev2] Temperature mismatch: soft-argmax decodes with xy/depth_temperature=0.8 but the loss uses plain softmax/log-softmax (τ=1.0). Measured coordinate bias is tiny (~0.05 heatmap-px, ~67 µm depth); align for cleanliness, not a primary cause.

Verified clean (don't re-chase)

• Heatmap ↔ label coordinate round-trip is consistent at 256² (model.py:610-627,688-697,1042); padding path doesn't trigger (256 divisible by 32). Fragile only for non-square/odd input sizes (the model.py:616-621 orig/padded rescale is a latent bug there, inert at 256²).
• [rev2] 2D label convention is x=col, y=row, origin top-left, in input-image pixels (verified against body_seg_maps foreground and decoded MP4: ~95% of keypoints land on foreground under this interpretation vs ~53% under the transpose). No x/y swap.
• [rev2] Camera unprojection math (pose/camera.py) is self-consistent — K⁻¹·[u·d, v·d, d] exactly inverts the dm_control projection given a matching sensor size. The only defect is the size mismatch (I1-A).
• [rev2] Contrastive InfoNCE multi-positive fix (PR Fix contrastive loss #41) is correct (pose/contrast/model.py:126-194): -logsumexp(log_softmax(logits)[:, :n_pos]) correctly credits all positive columns; row/label ordering matches collapse_batch (variant-major); verified numerically (identical positives → loss≈0.017, random→4.38). The historical single-positive bug is gone.
• scale_template_by_leg_segment_sizes size_new = size_dict[prev_keypoint_name] (invkin.py:114) is correct — both poseforge and seqikpy name a segment after its proximal keypoint; no off-by-one.
• NMF↔canonical keypoint mapping, per-leg proximal→distal order, make_symmetric + total-length summation are correct.
• Depth soft-label (model.py:886) is a properly normalized discretized Gaussian.
• Style-transfer channel order / [-1,1]↔[0,1] denorm is clean (no BGR swap); grayscale→3ch replication verified.
• concat_atomic_batches shares one label set across variants correctly; effective batch = train_batch_size × n_variants as documented; collate alignment verified.

Suggested pickup order for the next session [rev2 updated]

1. I1-A (build mapper at the prediction resolution / rescale pred_xy; add consistency assert) — now confirmed to inject ~10–18° into IK; cheapest high-leverage fix. Plus the I3-A restore and the c= kwarg fix to make IK paths run at all.
2. I1-B (spatially-grounded depth head) — primary cause of "3D worse than seg"; pair with re-centering the depth bins (I1-D) and fixing the disjoint default + adding the range assert (I1-G).
3. I3-C (warm-start gate/reset; don't blend chunks), then I3-D/E (NaN handling, confidence gating, order assert), then I3-B (bounds RoM investigation) — re-evaluate after I1-A is fixed.
4. I1-C (assert stored tile size == input_image_size; add a real resize+label-rescale path) — unblocks any future higher-resolution work (Issue 2), together with the GAN/heatmap/storage ceilings in I2-A.
5. I1-F (inference input-range check), I2-C (fix broken extraction script), [rev2] fix the pvio private-symbol import.

---

Rev 3 (2026-06-24) — fixes implemented + remaining human verification

Fixes for these findings were implemented by Claude Code (upon user instruction) as PRs against flygymv2: #49 (I1-A camera + c=), #50 (I1-B spatial depth head + I1-D/I1-G guards, opt-in), #51 (pvio import + I1-C loader guard + I2-C + the stale test_atomic_batch_dataset.py path), #52 (I3-A..E; I3-B now data-derived from the observed dof_angles RoM, superseding the mirror-only attempt).

What still needs a human:

• End-to-end IK self-consistency (IK robustness: restore deleted DOF map + NaN/order/bounds guards (I3-A..E) #52): scripts/verify_ik_selfconsistency.py (needs flygym+seqikpy, Data pipeline: fix pvio import break + loader resolution guard (I1-C) + extraction script (I2-C) #51 merged) recovers joint angles from stored keypoints and compares to ground-truth dof_angles; expect mean per-DOF error ≪ 15°. Confirms I3-A + I1-A + the solve at once.
• I3-B bounds are now derived from the training RoM ±10° (a data-driven check showed 10/42 prior bounds clipped real data). Verify against NeuroMechFly anatomical limits for production; RH_ThC_roll shows wrapping (−202°) needing upstream unwrapping.
• DOF order: IK output is [yaw,pitch,roll,…], ground truth [pitch,roll,yaw,…] — consume by name (dof_names_per_leg attr).
• Spatially-grounded depth head (opt-in) + depth-range guards (I1-B/I1-D/I1-G) #50 depth head is opt-in (depth_head_type="spatial") and requires a fresh training run to evaluate; default behavior unchanged.

[sibocw]

@stimpfli FYI this is the Claude Code audit. Proposed fixes are in the five PRs above. I haven't checked these very carefully but I asked Claude to generate scripts to verify things that are unclear, and retained things that I thought probably made sense.