code/train/main.py

# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Main file for Sudoku GPT experiments."""

import os
import sys

from absl import app
from absl import flags
from absl import logging

from clu import platform

import jax
import tensorflow as tf
import wandb

import ml_collections
from ml_collections import config_flags

from train import train_and_evaluate
from train import train_backtrack

import pdb



sys.dont_write_bytecode = True

logging.set_verbosity(logging.INFO)

FLAGS = flags.FLAGS

_WORKDIR = flags.DEFINE_string(
    'workdir',
    None,
    'Directory to store model data.')
_EXP_NAME = flags.DEFINE_string(
    'exp_name',
    None,
    'Experiment name.')
_CKPT_LOC = flags.DEFINE_string(
    'ckpt_loc',
    None,
    'Directory to restore model.')

config_flags.DEFINE_config_file(
    'config',
    None,
    'File path to the training hyperparameter configuration.',
    lock_config=True)
flags.mark_flags_as_required(['workdir', 'exp_name'])


def get_config():
    """Get the default hyperparameter configuration.

    Returns:
    A ConfigDict object.
    """
  
    # Common configuration for all experiments.
    config = ml_collections.ConfigDict()

    # Dataset choice
    config.dataset = 'sudoku'

    # Sequence order
    config.seq_order = "solver-order"           ## Choices = ["fixed", "solver-order", "random"]

    # Training related parameters
    config.max_steps = int(os.environ.get("SUDOKU_MAX_STEPS", 100000))
    config.dtype = jax.numpy.bfloat16
    config.minibatch_size = int(os.environ.get("SUDOKU_MINIBATCH", 64))

    # Continuous latent thoughts (ATC / Coconut style) + curriculum.
    # Set SUDOKU_LATENT_SLOTS=0 for the no-latent control baseline (plain
    # transformer + difficulty curriculum, no recurrence, no candidate head).
    config.num_latent_slots = int(os.environ.get("SUDOKU_LATENT_SLOTS", 6))
    config.latent_token_id = 10      # placeholder id for inactive slots
    # Recurrent latent feedback. 1 = full ATC/Coconut recurrence (slot k built
    # from slot k-1's hidden). 0 = NO recurrence: the K stage slots + shared
    # candidate head + curriculum + BCE supervision are all kept, but the latent
    # hidden state is never fed back, so each stage is an independent parallel
    # readout with no carried state (the "stagewise, no-recurrence" control).
    config.recurrent_latent = int(os.environ.get("SUDOKU_RECURRENT", 1))
    # Truncated BPTT over latent CoTs. 0 = full backprop through every
    # recurrent pass. T>0 = forward still builds all passes, but gradients
    # only flow through the last T thoughts (earlier z_j are stop-gradient).
    config.bptt_trunc = int(os.environ.get("SUDOKU_BPTT_TRUNC", 0))
    # Backtracking (stage-replay) training: when 1, use the train_backtrack loop
    # which interleaves frontier-stage batches with replays of earlier stages at
    # their MATCHED recurrence depth (num_passes=t for stage t) to mitigate
    # forgetting. backtrack_prob = fraction of steps that are replay steps.
    config.backtrack = int(os.environ.get("SUDOKU_BACKTRACK", 0))
    # Replay strategy: "prob" = fixed-probability replay of a random earlier stage
    # (backtrack_prob). "adaptive" = deficit-driven controller that only replays a
    # stage once its held-out accuracy has regressed below its graduation value by
    # backtrack_margin, repairs the most-deficient stage first with a FULL
    # frontier-style step, and returns once it recovers or hits the step cap.
    # Nesting backtrack, in the standard loop (NOT train_backtrack.py, whose
    # controller keys on candidate-set accuracy and drops the nest gate).
    # Replays the worst earlier latent step that is under the nest bar, using
    # that stage's own instances. nest_backtrack_ratio = share of steps spent
    # there while the deficit is mild; below nest_backtrack_hard the repair
    # takes every step until the step recovers.
    config.nest_backtrack = int(os.environ.get("SUDOKU_NEST_BT", 0))
    config.nest_backtrack_ratio = float(
        os.environ.get("SUDOKU_NEST_BT_RATIO", 0.25))
    config.nest_backtrack_hard = float(
        os.environ.get("SUDOKU_NEST_BT_HARD", 0.85))
    # Cap on a run of exclusive repair, in steps. 0 = uncapped, which let a
    # stage-3 run sit at 100% replay for 82k steps without clearing the bar.
    config.nest_backtrack_max_excl = int(
        os.environ.get("SUDOKU_NEST_BT_MAXEXCL", 0))
    # 1 = servo the replay share on the signed (nest deficit - hinge excess)
    # error instead of flipping between 100% repair and 100% frontier. The
    # share stays inside [NEST_BT_RATIO, NEST_BT_EXCLMAX], so the graduated
    # steps keep a maintenance share and the frontier is never starved.
    config.nest_backtrack_servo = int(
        os.environ.get("SUDOKU_NEST_BT_SERVO", 0))
    config.nest_backtrack_gain = float(
        os.environ.get("SUDOKU_NEST_BT_GAIN", 1.0))
    config.nest_backtrack_excl_max = float(
        os.environ.get("SUDOKU_NEST_BT_EXCLMAX", 1.0))
    config.backtrack_mode = os.environ.get("SUDOKU_BACKTRACK_MODE", "prob")
    config.backtrack_prob = float(os.environ.get("SUDOKU_BACKTRACK_PROB", 0.5))
    config.backtrack_margin = float(os.environ.get("SUDOKU_BACKTRACK_MARGIN", 0.03))
    config.backtrack_max_repair_steps = int(
        os.environ.get("SUDOKU_BACKTRACK_MAX_REPAIR_STEPS", 6000))
    # After leaving a repair episode, require this many frontier steps before
    # another repair can trigger (prevents thrashing / starves-the-frontier).
    config.backtrack_min_frontier_steps = int(
        os.environ.get("SUDOKU_BACKTRACK_MIN_FRONTIER_STEPS", 0))
    # While the frontier stage's own level accuracy is below this target, use a
    # stricter (larger) repair trigger so stage-6 training is prioritized until
    # it is itself "good". 0 disables.
    config.backtrack_frontier_target_acc = float(
        os.environ.get("SUDOKU_BACKTRACK_FRONTIER_TARGET_ACC", 0.0))
    # Optional explicit graduation refs for inherited stages when resuming at
    # stage > 1. Comma-separated floats for stages 1..start-1, e.g.
    # "0.462,0.487,0.534,0.597,0.636". Empty = seed from first eval (old behavior).
    config.backtrack_grad_acc_seed = os.environ.get("SUDOKU_GRAD_ACC_SEED", "")
    # Max fraction of steps (since resume) that may be spent in repair mode.
    # 0 disables the cap. Typical fair setting: 0.25.
    config.backtrack_max_repair_fraction = float(
        os.environ.get("SUDOKU_BACKTRACK_MAX_REPAIR_FRACTION", 0.0))
    # Soft graduation refs: each eval, if current < grad, decay
    #   grad <- (1-d)*grad + d*current
    # so chronic mild regression is forgiven; sharp drops still trigger.
    # 0 disables (hard bars). Typical: 0.05–0.1.
    config.backtrack_grad_decay = float(
        os.environ.get("SUDOKU_BACKTRACK_GRAD_DECAY", 0.0))
    # After this absolute step, force frontier-only (no new repairs).
    # 0 disables. Use for "BT early, then freeze".
    config.backtrack_freeze_after_step = int(
        os.environ.get("SUDOKU_BACKTRACK_FREEZE_AFTER", 0))
    # Frontier steps sample uniformly over ALL unlocked levels (standard-loop
    # data mix) instead of only the frontier level. 1 = fair "standard+repairs"
    # recipe (default); 0 = legacy frontier-level-only behavior.
    config.backtrack_frontier_mix = int(
        os.environ.get("SUDOKU_BACKTRACK_FRONTIER_MIX", 1))
    config.curriculum_start_stage = int(os.environ.get("SUDOKU_START_STAGE", 1))
    # Stage s = reasoning DEPTH s: s latent slots active, wave snapshots 1..s
    # supervised. Difficulty is NOT gated -- every puzzle is available from
    # step 0. Cap max_stage (e.g. =start_stage) to hold a single depth.
    config.curriculum_max_stage = int(os.environ.get("SUDOKU_MAX_STAGE", 12))
    # Promotion is gated on the candidate-set accuracy of the DEEPEST active
    # slot, restricted to cells that changed from the previous snapshot.
    # Instance arm: both gates required. mass = P(digit in raw candidate set S);
    # spread = H(p|S)/log|S| on |S|>=2 cells. S is the wave-solver candidate
    # set, not the filtered-instance support.
    config.promote_acc_threshold = float(os.environ.get("SUDOKU_PROMOTE_ACC", 0.85))
    config.promote_spread_threshold = float(
        os.environ.get("SUDOKU_PROMOTE_SPREAD", 0.85))
    # Entropy ceiling: fraction of |S|>=2 cells needing H(p) <= log|S|. Uniform
    # over S is the max-entropy distribution supported on S, so exceeding the
    # ceiling proves mass sits outside S. Necessary, not sufficient. 0 = off.
    config.promote_hbound_threshold = float(
        os.environ.get("SUDOKU_PROMOTE_HBOUND", 0.0))
    # The single-criterion gate on CE(q||p) vs log|S|, where q is the instances'
    # own per-cell digit distribution. When enabled it REPLACES the mass/spread/
    # hbound gates, which stay as diagnostics: it penalizes leak and collapse
    # together, and unlike a spread bar it is satisfiable at every stage.
    #
    #   mode "mean" (default): promote when the STAGE AVERAGE clears the bar,
    #     mean CE(q||p) <= mean log|S| - margin, over every |S|>=2 cell. This is
    #     the stage-level reading of "the value cross-entropy has been driven
    #     below the uniform-candidate entropy", i.e. the thing training
    #     minimizes. Set SUDOKU_PROMOTE_CEQ_MARGIN to demand slack in nats.
    #   mode "frac": promote when at least SUDOKU_PROMOTE_CEQ of the individual
    #     cells clear it. Strictly harder: the average can pass while many cells
    #     fail, because near-uniform cells sit right on the ceiling and any leak
    #     tips them over.
    #   mode "hinge": promote when mean over puzzles of max(0, CE_i - log|S|_i)
    #     is <= SUDOKU_PROMOTE_HINGE. Same comparison as "mean" but resolved per
    #     puzzle before averaging, so a puzzle still above its own bar can no
    #     longer be cancelled by puzzles carrying slack. By Jensen this is >= the
    #     "mean" violation, hence strictly stronger; demanding exactly 0 is
    #     usually infeasible because puzzles whose constraints barely bite have
    #     H(q) ~ log|S| and keep the floor above 0 (see per_stage_val_hinge_floor,
    #     which is logged next to it for exactly this calibration).
    #   mode "kl": promote when mean over puzzles of max(0, CE_i - H(q)_i) is
    #     <= SUDOKU_PROMOTE_KL, i.e. KL(q||p) <= delta. Every mode above is
    #     satisfied by a p that merely spreads over S, because CE(q||U_S) is
    #     log|S| exactly, so they certify the support of q and nothing about its
    #     frequencies. This one is zero only at p == q. Use it when the target is
    #     the distribution itself, which is what a final readout collapsing onto
    #     the unique solution requires.
    config.promote_ceq_mode = os.environ.get("SUDOKU_PROMOTE_CEQ_MODE", "mean")
    config.promote_kl_delta = float(os.environ.get("SUDOKU_PROMOTE_KL", 0.05))
    #   mode "eps": promote when mean over puzzles of max(0, CE_i - log|S|_i +
    #     eps) <= SUDOKU_PROMOTE_HINGE_SLACK, i.e. CE <= log|S| - eps. Same
    #     family as "kl" (KL <= headroom - eps) but stated from the uniform
    #     ceiling, so eps is the nats of superposition the stage must give up.
    #     Keep eps below the stage headroom log|S| - H(q) or p == q also fails.
    config.promote_eps = float(os.environ.get("SUDOKU_PROMOTE_EPS", 0.0))
    # Puzzle-level nesting on the whole prefix 1->2->...->s, from stage 2 on.
    # Promote only if that fraction of puzzles have no empty cell whose
    # nucleus set grows on ANY step. 0 = off. 0.98 is the intended bar.
    # Stage 1 is skipped (one latent, no step). Each step is also logged.
    config.promote_nest_frac = float(os.environ.get("SUDOKU_PROMOTE_NEST", 0.0))
    # Bar for the OLDEST step, when a single bar for every step is too strict.
    # The all-steps AND (nest_clean) fails if any one earlier step slipped, so a
    # run can hold 1->2 at 0.94 and still read nest 0.40. With BASE set, each
    # step d gets its own bar, ramping BASE (oldest) -> PROMOTE_NEST (newest).
    # 0 = off, i.e. the old single-bar nest_clean gate.
    config.promote_nest_base = float(
        os.environ.get("SUDOKU_PROMOTE_NEST_BASE", 0.0))
    config.promote_nest_tau = float(os.environ.get("SUDOKU_NEST_TAU", 0.99))
    # Fraction of eval puzzles that must still hold the unique solution in the
    # nucleus at every depth. 1.0 = every puzzle. 0 = off (diagnostic only).
    config.promote_keep_frac = float(os.environ.get("SUDOKU_PROMOTE_KEEP", 0.0))
    #     The bar is relative by default, because the floor is a property of the
    #     data and moves with the stage: promote when
    #         hinge <= hinge_floor + SUDOKU_PROMOTE_HINGE_SLACK.
    #     Set SUDOKU_PROMOTE_HINGE > 0 to use an absolute bar in nats instead.
    config.promote_hinge_threshold = float(
        os.environ.get("SUDOKU_PROMOTE_HINGE", 0.0))
    config.promote_hinge_slack = float(
        os.environ.get("SUDOKU_PROMOTE_HINGE_SLACK", 0.05))
    config.promote_ceq_threshold = float(
        os.environ.get("SUDOKU_PROMOTE_CEQ", 0.0))
    config.promote_ceq_margin = float(
        os.environ.get("SUDOKU_PROMOTE_CEQ_MARGIN", 0.0))
    # Location floor, applied to loc_wave. 0 = off.
    config.promote_loc_threshold = float(
        os.environ.get("SUDOKU_PROMOTE_LOC_WAVE", 0.0))
    # 0 disables patience; instance arm never uses it regardless.
    config.promote_patience_steps = int(os.environ.get("SUDOKU_PATIENCE", 0))
    config.min_stage_steps = int(os.environ.get("SUDOKU_MIN_STAGE_STEPS", 2000))
    # Ceiling on one stage. With a KL gate the bar can be out of reach for a
    # while, and without this a single stage would consume the whole run.
    # 0 = no ceiling.
    config.max_stage_steps = int(
        os.environ.get("SUDOKU_MAX_STAGE_STEPS", 0))
    # Plateau promotion: advance once the frontier depth stops improving, rather
    # than on a fixed timer. A stage is "done" when it has not gained
    # plateau_delta over its best accuracy for plateau_steps steps. This is the
    # primary rule; the accuracy threshold is a fast path for mastery and
    # promote_patience_steps is a hard cap so a stuck stage cannot stall
    # training forever. Set plateau_steps=0 to disable.
    config.plateau_steps = int(os.environ.get("SUDOKU_PLATEAU_STEPS", 20000))
    config.plateau_delta = float(os.environ.get("SUDOKU_PLATEAU_DELTA", 0.005))
    # Train-time difficulty balancing. 0 (default) = draw puzzles uniformly from
    # the corpus, so the difficulty tag selects nothing. 1 = uniform over the 6
    # levels, which upsamples level 8 from 1.8% to 16.7% of batches. Eval is
    # always level-balanced so per-level accuracy stays measurable.
    config.level_balanced_sampling = int(
        os.environ.get("SUDOKU_LEVEL_BALANCED", 0))
    # Data curriculum over the puzzle POOL (as opposed to the latent-depth
    # curriculum over the supervision).
    #   "none"   = every puzzle available from step 0.
    #   "rounds" = stage t admits only puzzles whose solver round count falls in
    #              the first t of max_stage equal-count bins (rounds span 5..38,
    #              so 12 bins give a genuinely smooth 12-step ladder). This is
    #              the same propagation-depth axis the latent arm supervises,
    #              which makes the two arms directly comparable. Needs
    #              SUDOKU_TRAIN_META / SUDOKU_TEST_META.
    config.data_curriculum = os.environ.get("SUDOKU_DATA_CURRICULUM", "none")
    config.train_meta_path = os.environ.get("SUDOKU_TRAIN_META", "") or None
    config.test_meta_path = os.environ.get("SUDOKU_TEST_META", "") or None

    # Model related parameters
    config.block_size = 81
    config.seq_len = 3 * config.block_size + config.num_latent_slots
    config.vocab_size = 11

    # Model architecture
    config.num_heads = 8
    config.num_layers = 8
    config.emb_dim = 576
    config.qkv_dim = 576
    config.mlp_dim = 6 * config.emb_dim
    config.dropout_rate = float(os.environ.get("SUDOKU_DROPOUT", 0.2))
    config.attention_dropout_rate = float(
        os.environ.get("SUDOKU_ATTN_DROPOUT",
                       os.environ.get("SUDOKU_DROPOUT", 0.2)))

    # Training hyperparameters
    config.learning_rate = float(os.environ.get("SUDOKU_LR", 0.0002))  # Base learning rate.
    config.end_lr_factor = float(os.environ.get("SUDOKU_END_LR_FACTOR", 0.2))
    config.warmup_tokens = int(os.environ.get("SUDOKU_WARMUP", 10000))
    config.weight_decay = float(os.environ.get("SUDOKU_WD", 0.005))
    # Resume from a checkpoint (set SUDOKU_RESUME=1 and pass --ckpt_loc=<path>).
    config.resume_training = os.environ.get("SUDOKU_RESUME", "0") == "1"

    # Other hyperparameters
    config.seed = 7
    config.save_checkpoint = os.environ.get("SUDOKU_SAVE_CKPT", "1") == "1"
    config.save_every_steps = int(os.environ.get("SUDOKU_SAVE_EVERY", 10000))
    # How many checkpoints to retain. Large default so per-stage checkpoints are
    # never rolled off (disk is plentiful; ~0.5GB each).
    config.ckpt_keep = int(os.environ.get("SUDOKU_CKPT_KEEP", 100))
    config.use_wandb = False
    config.wandb_project_name = 'sudoku'

    # Evaluation related parameters
    config.eval_every_steps = int(os.environ.get("SUDOKU_EVAL_EVERY", 2000))
    config.eval_epochs = int(os.environ.get("SUDOKU_EVAL_EPOCHS", 5))

    # Path to dataset
    config.train_puzzle_path = os.environ.get(
        "SUDOKU_TRAIN_PATH", "datasets/train_sudoku_puzzles.npy")
    config.train_candidate_path = "datasets/train_sudoku_puzzles_candidate.npy"
    config.test_puzzle_path = os.environ.get(
        "SUDOKU_TEST_PATH", "datasets/test_sudoku_puzzles.npy")
    config.test_candidate_path = "datasets/test_sudoku_puzzles_candidate.npy"

    # Staged multi-candidate supervision (per-latent-slot BCE targets).
    # Empty string disables cand-mask loading (useful for K=0 baselines).
    config.train_cand_masks_path = os.environ.get(
        "SUDOKU_TRAIN_CAND", "datasets_multicandidate/train_cand_masks.npy") or None
    config.test_cand_masks_path = os.environ.get(
        "SUDOKU_TEST_CAND", "datasets_multicandidate/test_cand_masks.npy") or None
    # Superposition-instance targets. When set, the output prompt's value tokens
    # come from one sampled stage-k assignment instead of the unique solution:
    # the input prompt (clues) is fixed and the same puzzle recurs with different
    # legal completions, so the candidate set is represented across the batch
    # rather than supervised as a multi-hot set. Setting this should go with
    # SUDOKU_AUX_WEIGHT=0 (candidate head off) -- the masks are then read only
    # for the in-set metric. Empty string = classic single-solution targets.
    # If unset, reuse the candidate-mask folder (HF data/ has masks + instances).
    _data_dir = os.path.dirname(config.train_cand_masks_path or "")
    config.instance_dir = os.environ.get("SUDOKU_INSTANCE_DIR", "") or (
        _data_dir or None)
    # Large-sample q for the promotion gate (gen_bigq.py output). The stored
    # pool stops at coverage, ~5 grids per puzzle, which is far too small to
    # estimate the entropy the gate spends as headroom.
    config.gate_q_path = os.environ.get("SUDOKU_GATE_Q", "") or (
        os.path.join(_data_dir, "test_bigq300.npz") if _data_dir else None)
    # Instance-valued prefixes per eval example, averaged to form the marginal
    # p that CE(q||p) is scored against.
    config.eval_instance_reads = int(os.environ.get("SUDOKU_EVAL_READS", 4))
    # Weight of the auxiliary candidate-set BCE loss relative to the LM CE loss.
    # 1 = every train step averages CE over depths 1..stage, each against that
    # depth's own instance plane. Stops exclusive-replay from overwriting
    # earlier latents. 0 = last-slot CE only (the old default).
    config.depth_mix = int(os.environ.get("SUDOKU_DEPTH_MIX", 0))
    # ---- Stochastic latents -------------------------------------------------
    # 1 = every latent slot carries a learned diagonal Gaussian (its own mean
    # and log-sigma matrix) instead of a single vector, and the readout is
    # Monte-Carlo averaged over SUDOKU_MC_SAMPLES draws. Both weight matrices
    # are zero-init, so the mean starts as the identity on the deterministic
    # thought and sigma starts at exp(SUDOKU_NOISE_LOGSIGMA0): a warm start
    # reproduces the source checkpoint on step one.
    config.latent_noise = int(os.environ.get("SUDOKU_LATENT_NOISE", 0))
    config.latent_mc_samples = int(os.environ.get("SUDOKU_MC_SAMPLES", 4))
    config.latent_noise_logsigma0 = float(
        os.environ.get("SUDOKU_NOISE_LOGSIGMA0", -4.0))
    # "mean" = Monte-Carlo estimate of E_z[CE]. "min" = charge only the best of
    # the draws, so draws can specialize to different instances.
    config.latent_mc_objective = os.environ.get("SUDOKU_MC_OBJ", "mean")
    config.aux_cand_weight = float(os.environ.get("SUDOKU_AUX_WEIGHT", 1.0))
    # Positive-class weight inside the candidate BCE (counters the sparsity of
    # the multi-hot masks so the head doesn't collapse to predicting all-zeros).
    config.aux_pos_weight = float(os.environ.get("SUDOKU_CAND_POS_WEIGHT", 5.0))
    # Relative LM-CE weight on the (row, col) location tokens versus the value
    # token of each solution triple. The output is a serialized (r,c,v) stream, so
    # two thirds of the supervised tokens are locations scored against the solver's
    # one arbitrary ordering of an order-interchangeable wave, while the promotion
    # gate only ever reads the teacher-forced value CE(q||p). 1.0/1.0 = the
    # original uniform per-token mean.
    config.loc_loss_weight = float(os.environ.get("SUDOKU_LOC_WEIGHT", 1.0))
    config.value_loss_weight = float(os.environ.get("SUDOKU_VALUE_WEIGHT", 1.0))
    # Weight on -log(probability mass inside the candidate set), applied at each
    # empty cell's value token. CE(q||p) decomposes as leak + H(q) + within-set
    # KL, and the leak term is the one that grows with curriculum depth, so this
    # attacks the part of the promotion deficit nothing else penalizes. 0 = off.
    config.leak_weight = float(os.environ.get("SUDOKU_LEAK_WEIGHT", 0.0))

    # How many latent slots each example activates.
    #   "level" = k = clip(level-2, 1, K).  Difficulty-matched, but puzzle level
    #             explains only ~19% of the variance in solver round count, so
    #             most examples leave the majority of the K slots inert (a
    #             level-3 puzzle activates ONE slot for ~21 rounds of work).
    #   "depth"  = k = num_passes, uniform over the batch. Every slot the
    #             recurrence actually fills is active and supervised, and the
    #             curriculum advances reasoning depth rather than puzzle level.
    config.cand_slot_mode = os.environ.get("SUDOKU_CAND_SLOT_MODE", "depth")
    # Latent passes granted per curriculum stage: num_passes = min(pps*stage, K).
    # pps=2 with K=12 reaches all 12 slots by stage 6.
    config.passes_per_stage = int(os.environ.get("SUDOKU_PASSES_PER_STAGE", 1))
    # Superposition experiment: iterate every (puzzle, instance) pair of the
    # pinned stage in shuffled epochs, so each instance of each puzzle is seen
    # exactly `instance_epochs` times. 0 = old behaviour (draw a puzzle at
    # random, then one of its assignments at random). instance_puzzles caps the
    # puzzle count (0 = whole corpus) to shorten the experiment.
    config.instance_epochs = int(os.environ.get("SUDOKU_INSTANCE_EPOCHS", 0))
    config.instance_puzzles = int(os.environ.get("SUDOKU_INSTANCE_PUZZLES", 0))
    # Preferred source: N instances per puzzle synthesized on the fly, each cell
    # drawn uniformly from its candidate set, so the per-cell digit frequencies
    # are equal and the CE optimum at that cell is the uniform superposition.
    config.instance_uniform_draws = int(
        os.environ.get("SUDOKU_UNIFORM_DRAWS", 0))
    # Loss weight for candidate cells that did NOT change from the previous
    # stage. Consecutive stages are highly redundant (at 12 stages ~97% of the
    # target bits are copies of the previous slot), so plain BCE is dominated by
    # echoing the previous slot. <1.0 down-weights the copied cells and puts the
    # gradient on the digits actually eliminated at this stage. 1.0 = off.
    config.aux_delta_bg = float(os.environ.get("SUDOKU_CAND_DELTA_BG", 1.0))

    return config


def main(argv):
    if len(argv) > 1:
        raise app.UsageError('Too many command-line arguments.')

    # # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make
    # # it unavailable to JAX.
    tf.config.experimental.set_visible_devices([], 'GPU')

    cfgs = get_config()
    if cfgs.resume_training: 
        assert _CKPT_LOC.value is not None
  
    if cfgs.use_wandb:
        wandb.init(project=cfgs.wandb_project_name, name=_EXP_NAME.value, config=cfgs)

    logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())
    logging.info('JAX local devices: %r', jax.local_devices())

    # Add a note so that we can tell which task is which JAX host.
    # (Depending on the platform task 0 is not guaranteed to be host 0)
    platform.work_unit().set_task_status(f'process_index: {jax.process_index()}, '
                                       f'process_count: {jax.process_count()}')
    platform.work_unit().create_artifact(platform.ArtifactType.DIRECTORY,
                                       _WORKDIR.value, 'workdir')
  
    logging.info(cfgs)

    cfgs.workdir = _WORKDIR.value
    cfgs.ckpt_loc = _CKPT_LOC.value
    if int(getattr(cfgs, "backtrack", 0)):
        if str(getattr(cfgs, "backtrack_mode", "prob")) == "adaptive":
            train_backtrack.train_and_evaluate_backtrack_adaptive(cfgs, _WORKDIR.value)
        else:
            train_backtrack.train_and_evaluate_backtrack(cfgs, _WORKDIR.value)
    else:
        train_and_evaluate.train_and_evaluate(cfgs, _WORKDIR.value)

    if cfgs.use_wandb:
        wandb.finish()


if __name__ == '__main__':
    jax.config.config_with_absl()
    app.run(main)