Repository navigation
feat(datasets): train 0-shot and 1-shot in one run via zero_shot_prob - #562
Conversation
A paired sample is demonstration + rollout, so a run is either wholly 1-shot or wholly 0-shot. `zero_shot_prob` is the fraction of samples whose demonstration half is blanked, letting one run teach both. Blanking zeroes the demonstration's camera frames AND sets `timestep_is_pad` across that half. The pad flag is load-bearing: without it the TTT inner loop would keep updating its fast weights on the blank frames (#558 zeroes `eta` only where that flag is set), which is worse than either mode. The loader raises instead of degrading quietly when the flag is absent. The mode is a pure function of (seed, index), drawn off an RNG stream separate from the pair draw so that index -> mode is stable across ranks and epochs and enabling the feature does not shift which episodes a given index pairs. Defaults to 0.0, where the feature is inert and pair draws are byte-identical to before.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
[claude-review] summary for commit f18f95f No blocking issues found. The head commit is docs-only and resolves the two remaining nits from the previous round, both verified against the implementation:
All earlier findings across the six review rounds (missing CHANGELOG entry, mode-seed/pair-seed collision at |
…aw seed The two seeds differed only by their multiplier, so they coincide at seed 0 -- the default, since the factory passes `cfg.seed or 0`. Pair draws stayed bit-identical (the generators are separate objects), but the mode was derived from the same engine state as the episode choice, correlating which episodes a sample pairs with whether it is blanked. Add a constant offset that cannot collide for any seed >= 0. Also probe for `timestep_is_pad` at construction instead of on the first blanked sample: at 0 < p < 1 the lazy guard surfaced a few steps into the run, inside a dataloader worker, far from the misconfiguration. Adds the CHANGELOG entry, a test pinning that the factory passes the kwarg through, and a test that the two seed streams stay distinct at seed 0.
There was a problem hiding this comment.
Reviewed 0c8e9e9. All four findings from the previous review round were addressed (CHANGELOG entry, construction-time probe, mode-seed offset, factory wiring test) — thanks. No blocking issues this round; two suggestions inline, both non-blocking. Summary in the [claude-review] comment.
The previous offset was 2^40, which torch silently discards: its CPU generator consumes only the low 32 bits of a seed, so `manual_seed(i)` and `manual_seed(i + 2**40)` return the identical stream. The seed-collision fix was therefore a no-op, and the test that claimed to pin it re-derived the arithmetic from copied multipliers, so it compared constants to themselves and passed anyway. Use an odd offset below 2^32. Both properties are load-bearing: under 2^32 so it survives truncation, and odd so a collision is unsatisfiable for EVERY seed -- it would need `seed * 6_919_000 + offset = 0 (mod 2**32)`, and the first term is always a multiple of 8. The pins now exercise the generator itself, so a revert to either a too-large or an even offset fails. Separately: a blanked sample now keeps the rollout's own instruction instead of the key's `ambiguous_prompt`. That override exists because the demonstration carries the answer; applying it to a sample whose demonstration was erased leaves neither, and the target becomes unlearnable.
|
Both round-2 suggestions addressed in 61d3041. The tautological-pin catch found a live bug. Rewriting that test to exercise the generator instead of re-deriving the arithmetic showed the previous fix was itself a no-op: torch's CPU generator consumes only the low 32 bits of a seed, so the The offset is now Prompt interaction: a blanked sample now keeps the rollout's own instruction rather than the key's |
There was a problem hiding this comment.
Reviewed 61d3041. Both round-2 suggestions were addressed, and the rework of the tautological pin genuinely fixed a live bug — I verified empirically that manual_seed(i) and manual_seed(i + 2**40) return identical streams on torch CPU, so the previous 1 << 40 offset really was a no-op, and the new 0x9E3779B9 offset survives truncation. No blocking issues; one nit inline. Summary in the [claude-review] comment.
…acks them The for-any-seed collision test hardcoded the multipliers' difference, so changing either multiplier would leave it passing against a stale constant while the oddness argument it rests on quietly stopped holding. Name both multipliers and derive the difference from them, in the test and in the offset's own derivation. Values are unchanged, so pair draws stay bit-identical.
There was a problem hiding this comment.
Reviewed ad678d0. The multiplier-naming refactor is correct — I verified the values are unchanged (pair draws stay bit-identical), the offset is odd and below 2^32, the difference is even, and the collision pin now derives from the source constants so a multiplier change can't pass against a stale number. All prior-round findings remain addressed. One nit inline; no blocking issues. Note: no Python environment was available on this runner, so I verified the tests statically rather than executing them — cpu_test.yml is the authoritative run.
…ng values The prose still quoted the literal multipliers after they were named, so a multiplier change that kept the difference even would pass every test while the comments silently misdescribed the seeds.
There was a problem hiding this comment.
Reviewed 288dbff (the docs-only follow-up to ad678d0). The source-side literal cleanup is correct, and I re-verified the collision arithmetic against the named constants (difference 6,919,000 is even, offset 0x9E3779B9 is odd and below 2^32, no colliding seed in the first 10^6). I also traced timestep_is_pad downstream: it gates only the TTT inner-loop eta (ttt_layer.py:643), and since attention runs per timestep, the un-blanked state/actions in a blanked demo half cannot reach the rollout's prediction — so zeroing only the camera keys is sound, not an oversight. No blocking issues; two residual nits inline.
… Raises The class docstring is the constructor's only documentation and did not cover `zero_shot_prob`, its range check, or the construction-time `timestep_is_pad` probe -- while the PR checklist's docstring box is checked. Documents `n_demos` too, which shared the gap. Also drops the literal multiplier difference from the collision test's docstring, and states the property the assert actually pins (even, not a multiple of 8).
There was a problem hiding this comment.
Reviewed 39972ab (docs-only follow-up). Verified each new docstring claim against the implementation per CLAUDE.md rule 8: the (n_demos + 1) * sequence_length shape claim, the blanked-sample-keeps-rollout-prompt claim (__getitem__'s and not blanked guard), and the three new Raises conditions all match the code, and the rewritten collision-test docstring now states exactly what the asserts pin (odd offset, even derived difference). No blocking issues — two doc-accuracy nits inline.
The Raises section said a key needs "fewer than two episodes" to raise, but the check is `len(v) < n_demos + 1` -- wrong for n_demos > 1, and the same commit had just documented n_demos above it. The CHANGELOG still said the multipliers' difference is "divisible by 8" after the source and the test were deliberately weakened to "even" to stop pinning a stale-prone literal.
There was a problem hiding this comment.
Reviewed f18f95f (docs-only). It corrects exactly the two stale claims flagged as nits in the previous round, and I verified both against the implementation per CLAUDE.md rule 8:
- The
Raisesdocstring now says a key needsn_demos + 1episodes — matches thelen(v) < needcheck atsrc/opentau/datasets/paired_sequence.py:189-194. - The CHANGELOG now claims the multiplier difference is even rather than divisible by 8.
_MODE_SEED_MULTIPLIER - _PAIR_SEED_MULTIPLIER = 6,919,000is even,_MODE_SEED_OFFSET = 0x9E3779B9is odd and below 2³², sos·diff + offset ≡ 0 (mod 2³²)is unsatisfiable for every seed — the argument holds, and the test attests/datasets/test_paired_sequence.py:721-723pins the same (weaker, non-stale-prone) properties the prose now states, so source, test, and CHANGELOG agree.
No findings. The PR is clean at this commit.
What this does
A paired sample is
demonstration + rollout, so today a run is either wholly 1-shot or wholly 0-shot. This addsDatasetMixtureConfig.zero_shot_prob: the fraction of paired samples whose demonstration half is blanked, so a single run can train both demo-present and demo-absent behaviour.Blanking does two things:
timestep_is_padacross that half.The pad flag is load-bearing.
TTTMLPLayer.forwardzeroes the inner-loopetaonly where that flag is set (#558), so without it the inner loop would keep updating its fast weights on the blank frames — worse than either mode, and silent._blank_demosraises rather than degrading quietly when the flag is absent (it is only emitted in sequence mode).loss_maskis alreadyFalseacross every demonstration, so it is unchanged, and sample shapes are unchanged, so batching is unaffected.The mode is a pure function of
(seed, index)drawn off an RNG stream separate from the pair draw. That keepsindex -> modestable across ranks and epochs (so distributed runs and resumes reproduce), and means enabling the feature does not shift which episodes a given index pairs.Defaults to
0.0, where the feature is inert.How it was tested
Added
TestZeroShotProbintests/datasets/test_paired_sequence.py(8 tests; the stub gained an opt-intimestep_is_padso every existing test keeps its exact sample shape). The suite is 50 passed.Covered:
0.0is inert — nothing blankedtimestep_is_padset across that half,loss_maskstillFalsethere, and the rollout half retains its contenttimestep_is_padis absentAlso exercised end to end against a real paired dataset through the normal loader, rather than the stub alone:
10070/20000 = 0.5035atzero_shot_prob=0.50/20000blanked atzero_shot_prob=0.05000pair draws byte-identical to the pre-change implementation,0differingThat last one is the backwards-compatibility argument: this is additive, not a behaviour change for existing n-shot runs.
How to checkout & try? (for the reviewer)
Enable it on a pairing run by setting the fraction of samples that should arrive without a demonstration:
{ "dataset_mixture": { "pair_episodes": true, "n_demos": 1, "zero_shot_prob": 0.5 } }Checklist
Note: Before submitting this PR, please read the contributor guideline.