Skip to content

WIP: Stochastic drive policy (Option A) — colored-Gaussian exploration on drives#6

Open
jmxpearson wants to merge 3 commits into
tract-on-spectralfrom
stochastic-drive-policy
Open

WIP: Stochastic drive policy (Option A) — colored-Gaussian exploration on drives#6
jmxpearson wants to merge 3 commits into
tract-on-spectralfrom
stochastic-drive-policy

Conversation

@jmxpearson

Copy link
Copy Markdown
Member

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 ξ_d is 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_scale is an external, schedulable multiplier. Reparameterization → pathwise gradients into the σ's. Optional ω–γ correlation via one tanh param.
  • drive_entropy(): anti-collapse entropy bonus.
  • The spectral reward is evaluated on sampled drives; the TF anchor stays on the mean drives (behavioral cloning, keeps the policy mean near the data).
  • noise_scale_for_step(): const|linear|geom exploration schedule by global step.
  • Fully gated behind drive_noise=False → deterministic checkpoints keep an identical state_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 > 0 and/or a scheduled/held noise_scale. Both are wired in.

Status

  • Compiles clean; unit-tested (unit-variance temporally-correlated noise, no-op when noise_scale=0, pathwise gradients reach the σ/correlation params, ω–γ correlation tracks target, entropy = mean log-σ).
  • Not yet run on GPU. Needs a smoke run with --rollout-backend eager (Pascal 1080 Ti constraint).

Note on diff

Carries the in-progress state of #5 (mrstft sc/logm component split, per-component weighted TB logging, autonomy_score return_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

@jmxpearson
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
jmxpearson force-pushed the stochastic-drive-policy branch from 3521196 to a73e72f Compare June 30, 2026 03:46
@jmxpearson
jmxpearson changed the base branch from spectral-loss-coldstart to tract-on-spectral June 30, 2026 03:46
claude added 2 commits June 30, 2026 08:34
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants