experiment: implement KP mixed-traffic innovation controls

3 files changed, 283 insertions, 28 deletions

diff --git a/experiments/analyze_kp_innovation_short.py b/experiments/analyze_kp_innovation_short.py
index 4c70575..5f21dea 100644
--- a/experiments/analyze_kp_innovation_short.py
+++ b/experiments/analyze_kp_innovation_short.py
@@ -82,12 +82,20 @@ def main():
         if len(tracking[rule]) != 20 or any(value is None for value in tracking[rule]):
             raise ValueError(f"MT-1 {rule} tracking trajectory is incomplete")
         for row, values in zip(records[rule]["epochs"], tracking[rule]):
+            mixed = row.get("mixed_apical")
+            if mixed is None:
+                raise ValueError(f"MT-1 {rule} mixed-apical trajectory is incomplete")
             trajectory_values.extend([
                 float(row["train_loss"]),
                 float(values["mean_feedback_forward_cosine"]),
                 float(values["mean_feedback_forward_relative_error"]),
                 float(values["min_feedback_forward_cosine"]),
                 float(values["max_feedback_forward_relative_error"]),
+                float(mixed["teaching_rms"]),
+                float(mixed["instruction_rms"]),
+                float(mixed["raw_apical_rms"]),
+                float(mixed["innovation_rms"]),
+                float(mixed["traffic_rms"]),
             ])
 
     all_finite = all(record["final"]["finite"] for record in records.values())
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index e8cf294..887674e 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -8,10 +8,12 @@ import torch
 import torch.nn.functional as F
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet,
-                       CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
+from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
+                       CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet,
+                       ConvSDILConfig,
                        channel_subspace_apical_calibration,
                        conv_hierarchical_step, conv_kolen_pollack_step,
+                       conv_kp_mixed_traffic_step,
                        conv_local_step,
                        hierarchical_mirror_observations,
                        hierarchical_parameter_subspace_calibration,
@@ -937,6 +939,128 @@ def kolen_pollack_checks():
     }
 
 
+def kp_mixed_traffic_checks():
+    """Mixed traffic isolates subtraction from norm and preserves KP locality."""
+    torch.manual_seed(211)
+    common = dict(
+        depth=8, base_width=2, seed=67, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0,
+        traffic_seed=4000)
+    net = CIFARKPMixedTrafficResNet(**common)
+    x = torch.randn(5, 3, 32, 32, dtype=torch.float64)
+    y = torch.tensor([0, 2, 4, 6, 8])
+    forward = net.forward(
+        x, return_cache=True, training=True, update_stats=False)
+    output_error = (torch.softmax(forward["logits"], dim=1)
+                    - F.one_hot(y, 10).to(torch.float64))
+    instruction = net.hierarchical_teaching(output_error, forward)
+
+    # Zero traffic and zero predictor collapse all three rules to clean KP.
+    zero_errors = []
+    for rule in ("raw", "matched", "innovation"):
+        components = net.mixed_apical_components(
+            instruction, forward["hiddens"], rule)
+        zero_errors.extend(float((left - right).abs().max()) for left, right in
+                           zip(components["used"], instruction))
+    assert max(zero_errors) == 0.0
+
+    calibration = net.calibrate_traffic_gain(
+        instruction, forward["hiddens"], target_ratio=4.0)
+    ratio_error = max(abs(value - 4.0) for value in
+                      calibration["realized_traffic_instruction_rms_ratio"])
+    assert ratio_error < 1e-12
+
+    # An exact per-unit predictor removes all predictable traffic.
+    for slope, gain, coefficient in zip(
+            net.P_traffic, net.traffic_gain, net.B_traffic):
+        slope.copy_(gain * coefficient)
+    exact = net.mixed_apical_components(
+        instruction, forward["hiddens"], "innovation")
+    exact_predictor_error = max(float((left - right).abs().max())
+                                for left, right in zip(
+        exact["innovation"], instruction))
+    assert exact_predictor_error < 1e-14
+
+    for slope, bias in zip(net.P_traffic, net.P_traffic_bias):
+        slope.zero_()
+        bias.zero_()
+    components = net.mixed_apical_components(
+        instruction, forward["hiddens"], "matched")
+    norm_errors = []
+    direction_errors = []
+    for raw, innovation, matched in zip(
+            components["raw"], components["innovation"],
+            components["matched"]):
+        raw_flat = raw.flatten(1)
+        innovation_flat = innovation.flatten(1)
+        matched_flat = matched.flatten(1)
+        norm_errors.append(float((
+            (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs()
+            / innovation_flat.norm(dim=1).clamp_min(1e-30)).max()))
+        direction_errors.append(float((F.cosine_similarity(
+            raw_flat, matched_flat, dim=1) - 1.0).abs().max()))
+    assert max(norm_errors) < 1e-12
+    assert max(direction_errors) < 1e-12
+
+    # All used signals retain equal independently recomputed local KP products.
+    correlation_errors = []
+    for rule in ("raw", "matched", "innovation"):
+        used = net.mixed_apical_components(
+            instruction, forward["hiddens"], rule)["used"]
+        forward_directions, _, _, output_weight, _ = (
+            net.local_ascent_directions(used, output_error, forward))
+        reciprocal, reciprocal_readout = net.reciprocal_feedback_directions(
+            used, output_error, forward)
+        correlation_errors.extend(float((left - right).abs().max())
+                                  for left, right in zip(
+            forward_directions[1:], reciprocal[1:]))
+        correlation_errors.append(float(
+            (reciprocal_readout + output_weight.t()).abs().max()))
+    assert max(correlation_errors) < 1e-14
+
+    # Predictor plasticity consumes only the supplied soma/traffic pair: once
+    # those are fixed, changing every forward/feedback weight has no effect.
+    left = CIFARKPMixedTrafficResNet(**common)
+    right = CIFARKPMixedTrafficResNet(**common)
+    for left_gain, right_gain, source in zip(
+            left.traffic_gain, right.traffic_gain, net.traffic_gain):
+        left_gain.copy_(source)
+        right_gain.copy_(source)
+    fixed_hiddens = [value.detach().clone() for value in forward["hiddens"]]
+    for value in right.W + right.Q + [right.W_out, right.R_out]:
+        value.add_(torch.randn_like(value))
+    left.predictor_step(fixed_hiddens, eta=0.1)
+    right.predictor_step(fixed_hiddens, eta=0.1)
+    predictor_independence_error = max(float((a - b).abs().max())
+                                       for a, b in zip(
+        left.P_traffic + left.P_traffic_bias,
+        right.P_traffic + right.P_traffic_bias))
+    assert predictor_independence_error == 0.0
+
+    net.traffic_rule = "innovation"
+    result = conv_kp_mixed_traffic_step(
+        net, x, y, ConvSDILConfig(
+            eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0,
+            weight_decay=0.0, learn_A=False, learn_P=True),
+        step=0, rule="innovation", predictor_every=16)
+    assert math.isfinite(result["loss"]) and result["did_predictor_update"]
+    assert all(not value.requires_grad for value in
+               net.W + net.Q + net.P_traffic + net.P_traffic_bias
+               + [net.W_out, net.R_out, net.b_out])
+    assert net.mixed_elementwise_ops_per_example("matched") > (
+        net.mixed_elementwise_ops_per_example("raw"))
+    return {
+        "kp_traffic_zero_limit_error": max(zero_errors),
+        "kp_traffic_ratio_error": ratio_error,
+        "kp_traffic_exact_predictor_error": exact_predictor_error,
+        "kp_traffic_matched_norm_error": max(norm_errors),
+        "kp_traffic_matched_direction_error": max(direction_errors),
+        "kp_traffic_reciprocal_correlation_error": max(correlation_errors),
+        "kp_traffic_predictor_parameter_independence_error": (
+            predictor_independence_error),
+    }
+
+
 def apical_learning_checks():
     torch.manual_seed(11)
     net = CIFARSDILResNet(depth=8, base_width=2, seed=6)
@@ -1046,6 +1170,7 @@ def main():
     report.update(hierarchical_parameter_calibration_checks())
     report.update(normalized_response_mirror_checks())
     report.update(kolen_pollack_checks())
+    report.update(kp_mixed_traffic_checks())
     report.update(apical_learning_checks())
     print(report)
     print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")
diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index 738c86d..d2c0653 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -12,11 +12,14 @@ import time
 import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet,
-                       CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
+from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
+                       CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet,
+                       ConvSDILConfig,
                        conv_alignment_report, conv_apical_calibration_step,
                        conv_hierarchical_alignment_report,
                        conv_hierarchical_step, conv_kolen_pollack_step,
+                       conv_kp_mixed_traffic_alignment_report,
+                       conv_kp_mixed_traffic_step,
                        conv_learned_hierarchical_step, conv_local_step,
                        evaluate_conv, hierarchical_feedback_tracking_report,
                        hierarchical_parameter_subspace_calibration)
@@ -108,17 +111,25 @@ def build(args):
         bn_momentum=args.bn_momentum, bn_eps=args.bn_eps)
     if args.mode == "bp":
         return CIFARLocalResNet(**common), None
-    if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"):
-        network_class = (
-            CIFARKPResNet if args.mode == "kp" else CIFARHierarchicalFAResNet)
-        net = network_class(
-            **common, feedback_seed=args.apical_seed,
-            feedback_scale=args.a_scale)
+    if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp", "kp_traffic"):
+        if args.mode == "kp_traffic":
+            net = CIFARKPMixedTrafficResNet(
+                **common, feedback_seed=args.apical_seed,
+                feedback_scale=args.a_scale, traffic_seed=args.traffic_seed)
+            net.traffic_rule = args.traffic_rule
+        else:
+            network_class = (
+                CIFARKPResNet if args.mode == "kp"
+                else CIFARHierarchicalFAResNet)
+            net = network_class(
+                **common, feedback_seed=args.apical_seed,
+                feedback_scale=args.a_scale)
         config = ConvSDILConfig(
             eta=args.lr, eta_output=args.output_lr,
             eta_A=(args.eta_A if args.mode == "lhfa" else 0.0),
+            eta_P=args.eta_P,
             momentum=args.momentum, weight_decay=args.weight_decay,
-            learn_A=args.mode == "lhfa", learn_P=False,
+            learn_A=args.mode == "lhfa", learn_P=args.mode == "kp_traffic",
             pert_sigma=args.pert_sigma, pert_every=args.pert_every,
             pert_directions=args.pert_directions,
             apical_calibration_mode="hierarchical_parameter_subspace")
@@ -145,7 +156,9 @@ def work_report(net, mode, counters):
     apical_macs = getattr(net, "apical_macs_per_example", 0)
     normal_forward = counters["ordinary_examples"] * forward_macs
     warmup_forward = ((counters["predictor_warmup_examples"]
-                       + counters["apical_warmup_examples"]) * forward_macs)
+                       + counters["apical_warmup_examples"]
+                       + counters["traffic_calibration_examples"]
+                       + counters["traffic_audit_examples"]) * forward_macs)
     calibration_forward = counters["perturbation_forward_examples"] * forward_macs
     if mode == "bp":
         bp_reverse = 2 * counters["ordinary_examples"] * forward_macs
@@ -156,7 +169,9 @@ def work_report(net, mode, counters):
         bp_reverse = 0
         local_correlation = counters["ordinary_examples"] * forward_macs
         apical_inference = ((counters["ordinary_examples"]
-                            + counters["apical_warmup_examples"]) * apical_macs)
+                            + counters["apical_warmup_examples"]
+                            + counters["traffic_calibration_examples"])
+                            * apical_macs)
         regression_multiplier = 2 if mode == "lhfa" else 1
         apical_regression = (regression_multiplier
                              * counters["calibration_event_examples"]
@@ -168,7 +183,17 @@ def work_report(net, mode, counters):
     mirror_prediction = mirror_forward if mode == "rrm" else 0
     mirror_correlation = mirror_forward
     kp_feedback_correlation = (
-        counters["ordinary_examples"] * apical_macs if mode == "kp" else 0)
+        counters["ordinary_examples"] * apical_macs
+        if mode in ("kp", "kp_traffic") else 0)
+    if mode == "kp_traffic":
+        elementwise_operations = (
+            counters["ordinary_examples"]
+            * net.mixed_elementwise_ops_per_example()
+            + (counters["predictor_warmup_examples"]
+               + counters["predictor_update_examples"])
+            * net.predictor_elementwise_ops_per_example)
+    else:
+        elementwise_operations = 0
     components = {
         "ordinary_forward_macs": normal_forward,
         "warmup_clean_forward_macs": warmup_forward,
@@ -187,12 +212,17 @@ def work_report(net, mode, counters):
         "apical_macs_per_example": apical_macs,
         "components": components,
         "total_macs_estimate": sum(components.values()),
+        "elementwise_operations_estimate": elementwise_operations,
         "total_clean_forward_examples": (
             counters["ordinary_examples"] + counters["predictor_warmup_examples"]
-            + counters["apical_warmup_examples"]),
+            + counters["apical_warmup_examples"]
+            + counters["traffic_calibration_examples"]
+            + counters["traffic_audit_examples"]),
         "total_forward_equivalent_examples": (
             counters["ordinary_examples"] + counters["predictor_warmup_examples"]
             + counters["apical_warmup_examples"]
+            + counters["traffic_calibration_examples"]
+            + counters["traffic_audit_examples"]
             + counters["perturbation_forward_examples"]),
         "logical_batch_loss_queries": counters["logical_batch_loss_queries"],
         "causal_scalar_observations": counters["causal_scalar_observations"],
@@ -203,18 +233,28 @@ def work_report(net, mode, counters):
             "multiply-accumulates in conv/linear maps; one local weight correlation "
             "equals one forward-weight MAC count; BP reverse is estimated as one "
             "weight-gradient plus one activation-gradient convolution per forward "
-            "convolution; elementwise nonlinearities and optimizer arithmetic excluded"),
+            "convolution; mixed-traffic/predictor elementwise arithmetic is reported "
+            "separately and is not folded into MACs"),
     }
 
 
 def run(args):
     if args.eval_split == "test" and args.eval_every:
         raise ValueError("test protocols must use --eval_every 0 (one final evaluation)")
-    if args.mode not in ("sdil", "lhfa") and (
+    if args.mode not in ("sdil", "lhfa", "kp_traffic") and (
             args.a_warmup_steps or args.learn_P):
         raise ValueError("apical/predictor warmup is restricted to SDIL")
     if args.mode == "lhfa" and args.learn_P:
         raise ValueError("predictor learning is not defined for learned HFA")
+    if args.mode == "kp_traffic":
+        if not args.learn_P or args.predictor_warmup_steps < 1:
+            raise ValueError("mixed-traffic KP requires predictor warmup")
+        if args.predictor_every < 1:
+            raise ValueError("mixed-traffic KP requires a predictor cadence")
+        if args.traffic_calibration_examples < 1 or args.traffic_ratio <= 0:
+            raise ValueError("invalid mixed-traffic calibration")
+    elif args.predictor_every:
+        raise ValueError("predictor cadence is restricted to mixed-traffic KP")
     if args.mode not in ("wm", "rrm") and args.mirror_warmup_steps:
         raise ValueError("mirror warmup is restricted to weight mirror mode")
     if args.mirror_every < 1 or args.mirror_batch_size < 1:
@@ -262,12 +302,17 @@ def run(args):
         "mirror_conv_examples": 0,
         "mirror_readout_examples": 0,
         "mirror_events": 0,
+        "traffic_calibration_examples": 0,
+        "traffic_audit_examples": 0,
+        "predictor_update_examples": 0,
     }
     log = {
         "schema_version": 1,
         "protocol_family": "oral_a_cifar_local_resnet_development",
         "calibration_metric_space": (
             None if config is None or args.mode == "hfa" else
+            "reciprocal_local_activity_products_with_mixed_apical_traffic"
+            if args.mode == "kp_traffic" else
             "hierarchical_feedback_parameters" if args.mode == "lhfa" else
             "reciprocal_local_activity_products" if args.mode == "kp" else {
                 "wm": "local_parent_child_response",
@@ -302,7 +347,8 @@ def run(args):
                 if args.mode == "hfa" else 0),
             "adaptive_feedback_parameters": (
                 getattr(net, "n_fixed_feedback_parameters", 0)
-                if args.mode in ("lhfa", "wm", "rrm", "kp") else 0),
+                if args.mode in ("lhfa", "wm", "rrm", "kp", "kp_traffic")
+                else 0),
         },
         "epochs": [],
     }
@@ -313,6 +359,33 @@ def run(args):
     apical_warmup_wall = 0.0
     mirror_warmup_wall = 0.0
     loader_state = train.g.get_state().clone()
+    traffic_calibration_batch = None
+    if args.mode == "kp_traffic":
+        count = min(args.traffic_calibration_examples, train.x.shape[0])
+        if count != args.traffic_calibration_examples:
+            raise ValueError("traffic calibration prefix is unavailable")
+        calibration_x = train.x[:count]
+        calibration_y = train.y[:count]
+        traffic_calibration_batch = (calibration_x, calibration_y)
+        calibration_forward = net.forward(
+            calibration_x, return_cache=True, training=True,
+            update_stats=False)
+        calibration_output_error = (
+            torch.softmax(calibration_forward["logits"], dim=1)
+            - torch.nn.functional.one_hot(
+                calibration_y, net.n_classes).to(
+                    calibration_forward["logits"].dtype))
+        calibration_instruction = net.hierarchical_teaching(
+            calibration_output_error, calibration_forward)
+        log["traffic_calibration"] = net.calibrate_traffic_gain(
+            calibration_instruction, calibration_forward["hiddens"],
+            args.traffic_ratio)
+        log["traffic_calibration"].update({
+            "examples": count,
+            "data_source": "first unaugmented development-training examples",
+            "uses_validation_endpoint": False,
+        })
+        counters["traffic_calibration_examples"] += count
     if args.mode in ("wm", "rrm") and args.mirror_warmup_steps:
         sync(args.device)
         mirror_start = time.time()
@@ -342,6 +415,7 @@ def run(args):
         sync(args.device)
         warmup_start = time.time()
         iterator = iter(train)
+        predictor_metrics = []
         for _ in range(args.predictor_warmup_steps):
             try:
                 x, _ = next(iterator)
@@ -349,10 +423,30 @@ def run(args):
                 iterator = iter(train)
                 x, _ = next(iterator)
             forward = net.forward(x, training=True, update_stats=False)
-            net.predictor_step(
-                forward["hiddens"], config.eta_P, config.nuisance_scale)
+            predictor_metric = (net.predictor_step(
+                forward["hiddens"], config.eta_P)
+                if args.mode == "kp_traffic" else net.predictor_step(
+                    forward["hiddens"], config.eta_P,
+                    config.nuisance_scale))
+            predictor_metrics.append(predictor_metric)
             counters["predictor_warmup_examples"] += x.shape[0]
         train.g.set_state(loader_state)
+        log["predictor_warmup"] = {
+            "steps": args.predictor_warmup_steps,
+            "first_mse": predictor_metrics[0],
+            "mean_mse": sum(predictor_metrics) / len(predictor_metrics),
+            "last_mse": predictor_metrics[-1],
+            "instruction_present": False,
+        }
+        if args.mode == "kp_traffic":
+            audit_x, _ = traffic_calibration_batch
+            audit_forward = net.forward(
+                audit_x, training=True, update_stats=False)
+            log["predictor_warmup"][
+                "post_warmup_traffic_residual_rms_ratio"] = (
+                    net.predictor_traffic_residual_rms_ratio(
+                        audit_forward["hiddens"]))
+            counters["traffic_audit_examples"] += audit_x.shape[0]
         sync(args.device)
         predictor_warmup_wall = time.time() - warmup_start
 
@@ -426,6 +520,7 @@ def run(args):
         examples = 0
         calibration_metrics = []
         mirror_metrics = []
+        signal_metrics = []
         for x, y in train:
             batch = x.shape[0]
             if args.mode == "bp":
@@ -476,6 +571,17 @@ def run(args):
                 result = conv_kolen_pollack_step(net, x, y, config)
                 loss = result["loss"]
                 did_perturb = False
+            elif args.mode == "kp_traffic":
+                result = conv_kp_mixed_traffic_step(
+                    net, x, y, config, step, args.traffic_rule,
+                    args.predictor_every)
+                loss = result["loss"]
+                did_perturb = False
+                signal_metrics.append({key: result[key] for key in (
+                    "teaching_rms", "instruction_rms", "raw_apical_rms",
+                    "innovation_rms", "traffic_rms")})
+                if result["did_predictor_update"]:
+                    counters["predictor_update_examples"] += batch
             else:
                 result = conv_local_step(
                     net, x, y, config, step, generator=perturb_generator)
@@ -521,7 +627,12 @@ def run(args):
                 key: sum(value[key] for value in mirror_metrics)
                 / len(mirror_metrics) for key in mirror_metrics[0]
             }
-        if args.mode == "kp":
+        if signal_metrics:
+            record["mixed_apical"] = {
+                key: sum(value[key] for value in signal_metrics)
+                / len(signal_metrics) for key in signal_metrics[0]
+            }
+        if args.mode in ("kp", "kp_traffic"):
             tracking = hierarchical_feedback_tracking_report(net)
             record["feedback_tracking"] = {
                 "mean_feedback_forward_cosine": tracking[
@@ -569,11 +680,15 @@ def run(args):
         probe = min(args.alignment_probe, train.x.shape[0])
         sync(args.device)
         diagnostic_start = time.time()
-        diagnostics = (conv_hierarchical_alignment_report(
-            net, train.x[:probe], train.y[:probe])
-            if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp") else
-            conv_alignment_report(
-                net, train.x[:probe], train.y[:probe], config))
+        if args.mode == "kp_traffic":
+            diagnostics = conv_kp_mixed_traffic_alignment_report(
+                net, train.x[:probe], train.y[:probe], args.traffic_rule)
+        elif args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"):
+            diagnostics = conv_hierarchical_alignment_report(
+                net, train.x[:probe], train.y[:probe])
+        else:
+            diagnostics = conv_alignment_report(
+                net, train.x[:probe], train.y[:probe], config)
         sync(args.device)
         diagnostics["wall_s"] = time.time() - diagnostic_start
         values = diagnostics["teaching_negative_gradient_cosine"]
@@ -620,8 +735,8 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--mode", choices=(
-            "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", "sdil",
-            "nodepert"),
+            "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp",
+            "kp_traffic", "sdil", "nodepert"),
         required=True)
     parser.add_argument("--out", required=True)
     parser.add_argument("--device", default="cpu")
@@ -673,6 +788,13 @@ def parse_args():
         choices=("unit_targets", "channel_subspace", "vectorizer_subspace"),
         default="unit_targets")
     parser.add_argument("--predictor_warmup_steps", type=int, default=0)
+    parser.add_argument("--predictor_every", type=int, default=0)
+    parser.add_argument("--traffic_rule",
+                        choices=("raw", "matched", "innovation"),
+                        default="innovation")
+    parser.add_argument("--traffic_seed", type=int, default=4000)
+    parser.add_argument("--traffic_ratio", type=float, default=4.0)
+    parser.add_argument("--traffic_calibration_examples", type=int, default=64)
     parser.add_argument("--a_warmup_steps", type=int, default=0)
     parser.add_argument("--mirror_warmup_steps", type=int, default=0)
     parser.add_argument("--mirror_every", type=int, default=16)