Implement stage2_model.stage1_context = "sampled" (gitea #41)

Stage 2 was trained on ground-truth stage-1 outcomes but deployed on
sampled ones, and in a rollout that gap compounds over every step of
every track — the same train/inference gap teacher_forcing="scheduled"
already closes within stage 2, just never applied at the stage
boundary. "sampled" was declared in the schema but rejected loudly by
validate_config as unimplemented; this lands the real implementation.

Mirrors the existing scheduled-sampling precedent rather than a hard
switch: new stage2_model.ctx_p_start/ctx_p_end (defaults 1.0 -> 0.0)
linearly ramp P(condition on ground truth) from epoch 0 to the final
epoch, so stage 2 doesn't chase a wildly moving stage-1 target early in
training. Per the plan discussed with the user: the sample is drawn
from stage 1's sampling_model() (EMA weights when present, matching
what inference actually deploys), mixed per example via a Bernoulli
draw (never blended within a row), and validation always uses the
ground truth regardless of the schedule. Fixes a latent bug the same
pattern would otherwise have hit: every sampler in giant/sample.py
flips its model to .eval() with no restore, so sampling from the raw
(non-EMA) stage-1 model mid-step now explicitly restores its .training
flag afterward to avoid silently corrupting stage 1's own training mode
for the rest of the epoch.

validate_config now enforces stage1_context in {"truth", "sampled"},
requires both stages active for "sampled" (nothing to sample from
otherwise), range-checks ctx_p_start/ctx_p_end, and rejects the
ctx_p_start = ctx_p_end = 1.0 configuration as an unadvertised no-op
identical to "truth".

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tests/test_config.py

@@ -780,13 +780,79 @@ def test_validate_config_bad_stop_sampling_rejected():
         assert "stop_sampling" in str(e)
 
 
-def test_validate_config_stage1_context_sampled_not_implemented():
+def test_validate_config_stage1_context_sampled_accepted_with_both_stages_active():
+    """gitea #41: 'sampled' is now implemented, so DEFAULT_CONFIG's
+    stage1_model/stage2_model.active = true (both) must let it through."""
     cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_bad_stage1_context_rejected():
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "sampled" in str(e)
+        assert "stage1_context" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage1_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage1_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage1_model.active" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage2_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage2_model.active" in str(e)
+
+
+@pytest.mark.parametrize("key", ["ctx_p_start", "ctx_p_end"])
+@pytest.mark.parametrize("value", [-0.1, 1.1])
+def test_validate_config_ctx_p_out_of_range_rejected(key, value):
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            f"stage2_model.{key}": value,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert key in str(e)
+
+
+def test_validate_config_stage1_context_sampled_always_truth_rejected_as_noop():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.ctx_p_start": 1.0,
+            "stage2_model.ctx_p_end": 1.0,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "ctx_p_start" in str(e) and "ctx_p_end" in str(e)
 
 
 def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():

tests/test_config_consumed_keys.py

@@ -54,13 +54,6 @@ _EXCLUDED_FILES = ("giant/model/_legacy.py",)
 # it. If a key here starts showing up as consumed, the fix landed and this
 # entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
 _KNOWN_UNUSED = {
-    "stage2_model.stage1_context": (
-        "issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the "
-        "ground-truth stage-1 output; 'sampled' is now rejected loudly by "
-        "validate_config (not silently accepted), but the key still isn't "
-        "read by any build/train consumer file since only 'truth' can pass "
-        "validation — see Issue 16 for the real implementation"
-    ),
     "stage2_model.autoregressive.order": (
         "gitea #30 — validate_config now checks order is 'energy_desc', but "
         "nothing in the build/train/rollout consumer whitelist reads the "

tests/test_train.py

@@ -5,7 +5,7 @@ import csv
 import math
 import tempfile
 from pathlib import Path
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -21,6 +21,7 @@ from giant.constants import (
 )
 from giant.data.dataset import StepBatch
 from giant.model.network import Stage2Autoregressive, build_critics, build_models
+from giant.sample import sample_stage1 as trainers_sample_stage1
 from giant.training import (
     FlowDDPMStageTrainer,
     StageSpec,
@@ -831,3 +832,127 @@ def test_wgan_physical_omits_grad_norm_slice_columns():
         header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
         assert "stage2/train/grad_norm_type_slice" not in header
         assert "stage2/train/grad_norm_cont_slice" not in header
+
+
+# --- stage2_model.stage1_context = "sampled" (gitea #41) --------------------
+
+
+def _sampled_ctx_cfg(ema_decay=0.999):
+    cfg = _base_cfg()
+    cfg["stage1_model"]["generator"] = "flow"
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["stage1_context"] = "sampled"
+    cfg["stage2_model"]["ctx_p_start"] = 0.0
+    cfg["stage2_model"]["ctx_p_end"] = 0.0
+    cfg["train"]["ema_decay"] = ema_decay
+    return cfg
+
+
+def _build_sampled_trainers(cfg):
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    return build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+
+
+def test_build_stage_trainers_attaches_stage1_only_under_sampled():
+    trainers = _build_sampled_trainers(_sampled_ctx_cfg())
+    assert trainers["stage2"].stage1_source is trainers["stage1"]
+    assert trainers["stage1"].stage1_source is None
+
+
+def test_build_stage_trainers_leaves_stage1_source_none_under_truth():
+    """Regression guard for the old silent no-op: 'truth' (the default) must
+    never attach a stage1_source, so _stage1_context short-circuits without
+    ever calling sample_stage1."""
+    cfg = _base_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    assert trainers["stage2"].stage1_source is None
+
+
+def test_stage1_context_sampled_calls_sample_stage1_and_differs_from_truth():
+    cfg = _sampled_ctx_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage1, stage2 = trainers["stage1"], trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert spy.call_count == 1
+    assert spy.call_args.args[0] is stage1.sampling_model()
+    assert not torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_truth_default_never_calls_sample_stage1():
+    cfg = _base_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage2 = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert spy.call_count == 0
+    assert torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_val_epoch_none_uses_ground_truth_even_under_sampled():
+    cfg = _sampled_ctx_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage2 = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None)
+    assert spy.call_count == 0
+    assert torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_sampled_preserves_stage1_training_mode():
+    """Every sampler in giant/sample.py flips its model to .eval() as a side
+    effect with no restore of its own (see sample_flow). Sampling from the
+    RAW stage-1 model (ema_decay=0, so sampling_model() returns self.model,
+    the same weights the stage-1 trainer is actively training on) must not
+    silently leave it in eval mode for the rest of the epoch's stage-1
+    updates."""
+    cfg = _sampled_ctx_cfg(ema_decay=0.0)
+    trainers = _build_sampled_trainers(cfg)
+    stage1, stage2 = trainers["stage1"], trainers["stage2"]
+    stage1.train_mode()
+    assert stage1.model.training
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert stage1.model.training
+
+
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_build_stage_trainers_sampled_step_runs(stage2_generator):
+    """Both trainer subclasses' call sites (FlowDDPMStageTrainer._compute,
+    WGANStageTrainer.step) must run end to end under 'sampled' and produce a
+    finite loss."""
+    cfg = _sampled_ctx_cfg()
+    cfg["stage2_model"]["generator"] = stage2_generator
+    trainers = _build_sampled_trainers(cfg)
+    trainer = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
+    loss_key = "g_loss" if stage2_generator == "wgan" else "loss"
+    assert math.isfinite(stats[loss_key])
+
+
+def test_train_end_to_end_stage1_context_sampled():
+    """Full train() run with stage1_context='sampled' must complete and
+    write a checkpoint + metrics.csv with finite losses throughout."""
+    cfg = _sampled_ctx_cfg()
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert len(rows) == cfg["train"]["epochs"]
+        assert all(math.isfinite(float(r["stage2/train/loss"])) for r in rows)