WIP: Stochastic drive policy (Option A) — colored-Gaussian exploration on drives#6
Open
jmxpearson wants to merge 3 commits into
Open
WIP: Stochastic drive policy (Option A) — colored-Gaussian exploration on drives#6jmxpearson wants to merge 3 commits into
jmxpearson wants to merge 3 commits into
Conversation
jmxpearson
marked this pull request as ready for review
June 29, 2026 17:09
…of tract/envelope Re-applies the drive-policy experiment onto the tract-on-spectral codebase so the branch carries the FULL feature set (vocal tract, envelope head of the Mamba encoder, freeze/warmup schedules, all entry-script options) with the stochastic drive policy as the only delta. Supersedes the earlier port that was built on the pre-tract spectral-loss-coldstart base and was missing those features. Drive policy (model.sample_drives): adds reparameterized colored-Gaussian noise to the deterministic drives (omega/gamma/kernel-weights only; envelope/tract left deterministic), scaled RELATIVE to each drive's own RMS (sigma is a fraction, not absolute -- the drives span ~75x in scale so an absolute sigma destabilizes gamma). Per-drive learnable-or-frozen sigma (--drive-noise-trainable), external schedulable noise_scale with const/linear/geom warmup, optional omega-gamma correlation, and an entropy bonus. Gated behind drive_noise=False so deterministic checkpoints are unchanged. Plumbed through model_cv, train (save/load flags + TB Drive/* logging), and the entry-script CLI/manifest, composing with the tract/envelope freeze schedules already present here. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RoxdXxdopD6qpBpxtZi1pC
jmxpearson
force-pushed
the
stochastic-drive-policy
branch
from
June 30, 2026 03:46
3521196 to
a73e72f
Compare
jmxpearson
changed the base branch from
spectral-loss-coldstart
to
tract-on-spectral
June 30, 2026 03:46
…on-spectral worktree The monitor snippet hardcoded sys.path.insert to a sibling worktree, whose load_model predates drive_noise and would fail to load a drive_noise checkpoint. Rely on PYTHONPATH (the driver sets it to this branch root) so scoring uses THIS branch's drive_noise-aware load_model + tract/envelope model. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RoxdXxdopD6qpBpxtZi1pC
…aturation) The figure used un-normalized log10|STFT| with a fixed vmax=-2, so any frequency content above ~1e-4 amplitude saturated to the brightest colour. The target washed out (no visible structure) and a ~100x-quiet autonomous rollout could still paint a bright line, misreading as loud. Now colour = 20*log10(|STFT| / target_peak) over [-80, 0] dB, with the SAME target-peak reference for every row (target/auto/source). The target peaks at 0 dB and its structure fills the range; a rollout that is N dB quieter sits N dB lower on the same amplitude-faithful scale (e.g. ~-40 dB for a 100x-quiet auto), so brightness now tracks amplitude. cmap -> magma for dB readability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RoxdXxdopD6qpBpxtZi1pC
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
2 participants
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.
WIP / stacked on #5. Branched off the spectral-rollout PR worktree to try a stochastic-drive-policy refinement. Base is `spectral-loss-coldstart` so this reads as a refinement of #5; the diff narrows once #5's local commits are pushed/merged.
What
Promote the deterministic drive head to a reparameterized Gaussian policy so the poly model explores in drive-space during spectral-rollout training (motivated by "the model isn't exploring enough"). This is Option A: colored-Gaussian additive noise with a learnable per-drive variance and an externally schedulable exploration multiplier.
How
model.sample_drives():d(t) <- d(t) + noise_scale · σ_d · ξ_d(t), whereξ_dis white noise pushed through the existing zero-phase Gaussian drive low-pass and renormalized to unit marginal variance → temporally-local exploration at the control rate (not per-sample white jitter the in-loop filter would erase).σ_d = exp(log_σ_d)is a learnable per-drive parameter;noise_scaleis an external, schedulable multiplier. Reparameterization → pathwise gradients into the σ's. Optional ω–γ correlation via onetanhparam.drive_entropy(): anti-collapse entropy bonus.noise_scale_for_step(): const|linear|geom exploration schedule by global step.drive_noise=False→ deterministic checkpoints keep an identicalstate_dict; eval/autonomy is bit-for-bit unchanged.Design caveat
Pure reparameterized gradient on a reconstruction reward drives σ down (noise usually worsens the fit), so the learnable variance alone collapses → less exploration. To actually increase exploration use
--lam-entropy > 0and/or a scheduled/heldnoise_scale. Both are wired in.Status
noise_scale=0, pathwise gradients reach the σ/correlation params, ω–γ correlation tracks target, entropy = mean log-σ).--rollout-backend eager(Pascal 1080 Ti constraint).Note on diff
Carries the in-progress state of #5 (mrstft
sc/logmcomponent split, per-component weighted TB logging,autonomy_scorereturn_trajectories) because the policy logging is interleaved with it. Those will land via #5.Example
```
python examples/train_poly_spectral_blk445.py --data-dir <...> --drive-noise
--drive-noise-init 0.05 --lam-entropy 1e-3 --drive-noise-corr
--drive-noise-scale-start 1.0 --drive-noise-scale-end 0.2 --drive-noise-schedule geom
--rollout-backend eager
```
🤖 Generated with Claude Code