fix: teacher-force the offset loss in JointPolicyObjective#258
Open
felixmaximilian wants to merge 9 commits into
Open
fix: teacher-force the offset loss in JointPolicyObjective#258felixmaximilian wants to merge 9 commits into
felixmaximilian wants to merge 9 commits into
Conversation
…set/_sample_codes Behavior-preserving: forward/predict outputs are byte-identical (regression test against the pre-refactor implementation lands with the fix commit). Exposes the full (b, g, c, action_dim) offset table for diagnostics. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The offset head was supervised at codes sampled from the policy's own logits (sample_codes=True during training). Given the features, a sampled code is independent of the GT action, so under L1 the offset table's optimum is offset[c] -> median(action|state) - decode(c) for every c: it cancels the code selection and pulls all codes toward the per-dimension conditional median, producing simultaneous gas+brake at argmax decoding in bimodal states. - teacher_force_offset: bool = True gathers the offset at target_codes so each entry only sees residuals of actions that quantized to it; false reproduces the previous behavior (A/B is a config flip via config/experiment/yaak/control_transformer/finetune.yaml) - the sampled-code reconstruction error (the old loss value) is logged as a gradient-free metric (metric/offset_sampled_recon) for curve comparability across the A/B - tests: gather correctness, exact regression vs the pre-refactor _predict (seeded multinomial included), gradient routing (offset loss -> offset_head only; focal -> code_head only; tokenizer frozen), teacher_force_offset=false parity with the old loss Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…lt off) Defense-in-depth for the offset head: when offset_scale is set, the gathered offset becomes scale * tanh(x / scale) - identity for |x| << scale (sub-cell residuals pass through unchanged), saturating at +-scale (cross-cell excursions capped). Off by default so the teacher-forcing A/B measures the fix in isolation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- offset_diag.py: shared harness (checkpoint loading with export-time argmax decoding, hydra-composed 3hz dataloaders, per-batch tensor iterator exposing features/code_logits/offset table/GT codes) - offset_cancellation_probe.py (1a+1b): cancellation ratio R = Var_c(decoded+offset)/Var_c(decoded) by entropy quartile, and the sampled-vs-GT contamination audit - pedal_conflict_probe.py (1c): model-side predicted gas+brake conflict rate at the sensor-data thresholds, by entropy quartile, vs GT floor - residual_unimodality_probe.py (1d): within-code residual histograms + tail-mass flag on the action tokenizer - pedal_conflict_stats.py: point CONFIG_ROOT at this checkout (was a hardcoded rmind-main sibling) and print the previously-computed-but- dropped conflict totals Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…heckpoint Freezes everything except offset_head and retrains it from the bugged checkpoint's weights under both supervision variants on cached features, then evaluates cancellation ratio, pedal-conflict rate, and offset magnitudes side by side. Fast (hours) real-data evidence that teacher forcing alone decontaminates the offset table while continued sampled-code training does not. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
surface: decode all 65,536 combos per tokenizer, report conflict fraction (unweighted + real-usage-weighted at sensor and 0.02 thresholds), GT-decode rate, policy argmax hit rate, and usage concentration. Key finding: the conflict-filtered tokenizers have the same decode surface as the dirty one (42-47% of combos at 0.02) - most combos are off-manifold, so training-data filtering cannot constrain them. tier2b: co-retrain code+offset heads on cached frozen features against a clean codebook, teacher-forced vs sampled; the cancellation mechanism reproduces identically (R 0.86-0.95 vs 0.01), so the bug is codebook-independent. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…probe Collect raw pred/GT pedal values so one pass reports, per side: gas/brake activation rates at sensor/0.02/0.05 thresholds (overall and by entropy bucket) and conflict rates at all three thresholds. A median-collapsed offset head suppresses minority-mode activations, so a predicted gas-activation deficit vs GT is the majority-mode lock-in signature the A/B should show on the bug arm and not the fix arm. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- uv format in CI floats to the newest ruff (0.15.21), which reformats inference/backends/tensorrt.py (whitespace-only) — apply it - repo-wide 'ruff check' auto-fixes (import order in train.py, logger callbacks, tensorrt.py) plus targeted noqas for the five remaining pre-existing violations (__init__.py eval resolver, vq.py acronym import / arg count / private codebook access) No functional changes; all files outside the fix's scope were only touched to keep 'just check' green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
felixmaximilian
added a commit
that referenced
this pull request
Jul 12, 2026
Required by the PR #258 hparams_jq interpolation; true = the fixed teacher-forced offset loss, matching the finetune experiment default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Under argmax any mode-committing policy shows a gas-activation deficit vs GT in ambiguous states, so argmax activation rates cannot separate majority-mode lock-in (bug) from legitimate mode commitment. Sampled decoding is the discriminator: a healthy multimodal head draws modes in proportion and matches GT activation rates; a median-collapsed head cannot, regardless of decoding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
JointPolicyObjectivesupervises the offset head at codes sampled from the policy's own logits (sample_codes=Trueduring training). A sampled code is independent of the action it is graded against, so every code's offset entry trains on the full conditional action distribution, and under L1 its optimum is the per-dimension conditional median — for every code. The offset table converges to a negative copy of the codebook (measured: regression slope −0.97, corr −0.99 against the decodes) and cancels the code selection entirely (cancellation ratio R ≈ 0.02–0.05 onmodel-2gqxhjod:v9, i.e. offsets erase 95–99% of cross-code output variance).Consequences at argmax export: coordinate-wise medians of multimodal action distributions press gas and brake simultaneously in ambiguous states (5.1% of val frames vs 0% in GT on the same frames), and minority behaviors are emitted never, not rarely (majority-mode lock-in — e.g. launching from rest). Data cleaning cannot fix this; the mechanism reproduces identically on conflict-filtered tokenizers.
Fix
Teacher-force the offset loss: gather the offset at
target_codes(already computed for the focal loss), so each code's entry only sees residuals of actions that actually quantized to it. Inference paths are unchanged — the offset table is still indexed by predicted codes at serving.teacher_force_offset: bool = Truekeeps the pre-fix behavior reachable (=false) so the A/B is a config flip, wired throughpolicy_finetune.yaml+ the finetune experiment configmetric/offset_sampled_recon) for train-curve comparabilityoffset_scalesoft bound (scale·tanh(x/scale), default off): makes codebook-scale cancellation unrepresentable — defense in depth, kept out of the A/B_predictsplit into_heads/_sample_codes/_gather_offset(behavior-preserving; bit-exact regression test vs the old implementation, seeded multinomial included)src/rmind/scripts/: cancellation probe (1a+1b), pedal-conflict probe (1c, incl. activation rates + multi-threshold), residual-unimodality probe (1d), offset-head-only retrain, codebook conflict surfaceEvidence (condensed)
Offset-head-only retrain of the bugged checkpoint repairs it in place (R 0.02 → 0.79–0.87) while a sampled-code control stays collapsed. Full 10-epoch A/B in flight on tokenizer
model-y74asdtd:v9: baseline rpiav4ay vs fix zcglgth4 (identical seed/data/config, only the flag differs).Deep dive: running report (Notion) · interactive explainer — live simulator + real codebook maps · full evidence page — toy A/B + checkpoint probes
Tests
tests/test_joint_policy.py— 15 passed: gather correctness (hand-computed), bit-exact_predictregression vs the pre-refactor implementation, gradient routing (offset loss →offset_headonly; focal →code_headonly; tokenizer frozen),teacher_force_offset=falseparity with the old loss, metric detachment,offset_scalebound/identity.🤖 Generated with Claude Code