#!/usr/bin/env python3 # --------------------------------------------------------------- # Task-oriented evaluation for the WCL revision. # # Adds two things requested by the reviewers: # (R1.6) a task-oriented metric beyond CosSim: top-1 semantic # retrieval accuracy over the C=8,000 AG News pool # (chance = 1/C), i.e. each reconstructed user embedding # is classified against the full sentence pool by cosine. # (R2.6) an INDEPENDENT baseline that does not share the proposed # mask-learning machinery: FIXED mutually-orthogonal masks # (non-learned signatures, O(M)=0 by construction) with a # transceiver trained to them. This is the classical # "orthogonal-signature" reference point. # # Methods (all at U=4, matched 100-epoch budget, seed 0): # - Proposed DRL : load existing checkpoint (trx+actor), det. mean. # - Fixed-Orth : fixed orthogonal masks + trained transceiver. # - Static (Sem) : free masks + transceiver, MSE+CosSim loss. # - Static (CE) : free masks + transceiver, symbol-CE loss. # # Outputs: ../results_sweeps/task_oriented/taskmetric_sweep.csv # --------------------------------------------------------------- import os, math, csv, random, argparse import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from drl_mask_policy import (MaskedTransceiver, MaskActor, apply_channel, encode_state, orthogonality_penalty) SNRS = list(range(0, 31, 5)) def set_seed(s): torch.manual_seed(s); np.random.seed(s); random.seed(s) def load_emb(path, d_bert, pool_size, device): if path and os.path.exists(path): emb = torch.load(path, map_location="cpu") emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1) print(f"[INFO] Loaded {tuple(emb.shape)} embeddings from {path}") else: emb = F.normalize(torch.randn(pool_size, d_bert), p=2, dim=-1) print("[INFO] Using synthetic embeddings.") return emb.to(device) @torch.no_grad() def sweep_metrics(trx, emb, mask_fn, device, U, eval_trials, eval_seed=12345, channel="rayleigh"): """Return per-SNR (cossim, orthogonality, top1_accuracy). mask_fn(snr_db) -> (U, d_s) mask matrix for that context. A fixed eval_seed makes the channel draws identical across methods, giving a paired comparison. """ trx.eval() C = emb.size(0) rows = [] for snr_db in SNRS: set_seed(eval_seed) # identical channel/idx draws per method coss, orths, accs = [], [], [] for _ in range(eval_trials): idx = torch.randint(0, C, (U,), device=device) b = emb[idx] m = mask_fn(snr_db) b_hat, _ = trx(b, m, snr_db, channel) coss.append(F.cosine_similarity(b_hat, b, dim=-1).mean().item()) orths.append(orthogonality_penalty(m).item()) # top-1 retrieval over the full pool bn = F.normalize(b_hat, p=2, dim=-1) sim = bn @ emb.T # (U, C) pred = sim.argmax(dim=-1) # (U,) accs.append((pred == idx).float().mean().item()) rows.append((snr_db, float(np.mean(coss)), float(np.mean(orths)), float(np.mean(accs)))) print(f" [SNR {snr_db:>2d}dB] CosSim={rows[-1][1]:.4f} " f"Orth={rows[-1][2]:.5f} Top1Acc={rows[-1][3]:.4f}") return rows def train_transceiver(emb, args, device, mode, fixed_masks=None): """Train a transceiver (+ free masks unless fixed_masks given). mode in {'sem','ce','fixed'}. Returns (trx, mask_tensor_fn).""" set_seed(args.seed) U, d_bert, d_s = args.users, args.d_bert, args.d_bert * args.mux_factor C = emb.size(0) trx = MaskedTransceiver(U, d_bert, d_s).to(device) if fixed_masks is not None: masks_param = None m_fixed = fixed_masks.to(device) params = list(trx.parameters()) else: masks_param = nn.Embedding(U, d_s).to(device) nn.init.normal_(masks_param.weight, std=0.5) params = list(trx.parameters()) + list(masks_param.parameters()) opt = torch.optim.Adam(params, lr=args.lr) ar = torch.arange(U, device=device) tau = float(args.ce_tau) for ep in range(1, args.epochs + 1): for _ in range(args.steps_per_epoch): snr_db = random.choice(args.train_snr) idx = torch.randint(0, C, (U,), device=device) b = emb[idx] m = m_fixed if fixed_masks is not None else masks_param(ar) b_hat, _ = trx(b, m, snr_db, args.channel) if mode == "ce": b_hat_n = F.normalize(b_hat, p=2, dim=-1) loss = F.cross_entropy(tau * (b_hat_n @ emb.T), idx) else: # 'sem' and 'fixed' both use MSE+CosSim cos = F.cosine_similarity(b_hat, b, dim=-1).mean() loss = 0.5 * F.mse_loss(b_hat, b) + 0.5 * (1.0 - cos) opt.zero_grad(set_to_none=True) loss.backward() nn.utils.clip_grad_norm_(params, 1.0) opt.step() if ep % 20 == 0 or ep == 1: print(f" [{mode} ep {ep:03d}/{args.epochs}] loss={loss.item():.4f}") if fixed_masks is not None: return trx, (lambda snr: m_fixed) masks_const = masks_param(ar).detach() return trx, (lambda snr: masks_const) def build_fixed_orthogonal_masks(U, d_s, seed=0): """U mutually-orthogonal unit-RMS signature masks (O(M)=0).""" g = torch.Generator().manual_seed(seed) A = torch.randn(d_s, U, generator=g) Q, _ = torch.linalg.qr(A) # (d_s, U) orthonormal columns M = Q.T # (U, d_s), orthonormal rows M = M * math.sqrt(d_s) # scale to unit RMS per row return M def load_drl(ckpt_path, args, device): ck = torch.load(ckpt_path, map_location=device, weights_only=False) U, d_bert, d_s = args.users, args.d_bert, args.d_bert * args.mux_factor trx = MaskedTransceiver(U, d_bert, d_s).to(device) trx.load_state_dict(ck["trx"]) actor = MaskActor(U, d_s, hidden=args.hidden, rank=args.rank).to(device) actor.load_state_dict(ck["actor"]) actor.eval() @torch.no_grad() def mask_fn(snr_db): st = encode_state(snr_db, U, args.users_max, device=device) mean, _ = actor(st) return mean return trx, mask_fn def main(): p = argparse.ArgumentParser() p.add_argument("--users", type=int, default=4) p.add_argument("--users-max", type=int, default=8) p.add_argument("--mux-factor", type=int, default=4) p.add_argument("--d-bert", type=int, default=768) p.add_argument("--hidden", type=int, default=256) p.add_argument("--rank", type=int, default=64) p.add_argument("--embed-file", type=str, default="../bert_agnews_8000.pt") p.add_argument("--pool-size", type=int, default=8000) p.add_argument("--channel", default="rayleigh") p.add_argument("--train-snr", type=float, nargs="+", default=[0, 5, 10, 15, 20, 25]) p.add_argument("--lr", type=float, default=1e-3) p.add_argument("--ce-tau", type=float, default=16.0) p.add_argument("--epochs", type=int, default=100) p.add_argument("--steps-per-epoch", type=int, default=200) p.add_argument("--eval-trials", type=int, default=200) p.add_argument("--seed", type=int, default=0) p.add_argument("--drl-ckpt", type=str, default="../results_sweeps/drl_U4_100ep_s0/drl_mask_policy.pt") p.add_argument("--out-dir", type=str, default="../results_sweeps/task_oriented") args = p.parse_args() device = torch.device("mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")) print(f"[INFO] device={device}") U, d_s = args.users, args.d_bert * args.mux_factor emb = load_emb(args.embed_file, args.d_bert, args.pool_size, device) results = {} # method -> list of rows print("\n=== Proposed DRL (from checkpoint) ===") trx, mfn = load_drl(args.drl_ckpt, args, device) results["proposed_drl"] = sweep_metrics(trx, emb, mfn, device, U, args.eval_trials) print("\n=== Fixed orthogonal masks (independent baseline) ===") Mfix = build_fixed_orthogonal_masks(U, d_s, seed=0) trx, mfn = train_transceiver(emb, args, device, "fixed", fixed_masks=Mfix) results["fixed_orth"] = sweep_metrics(trx, emb, mfn, device, U, args.eval_trials) print("\n=== Static masking + semantic loss ===") trx, mfn = train_transceiver(emb, args, device, "sem") results["static_sem"] = sweep_metrics(trx, emb, mfn, device, U, args.eval_trials) print("\n=== Static masking + CE loss ===") trx, mfn = train_transceiver(emb, args, device, "ce") results["static_ce"] = sweep_metrics(trx, emb, mfn, device, U, args.eval_trials) os.makedirs(args.out_dir, exist_ok=True) out = os.path.join(args.out_dir, "taskmetric_sweep.csv") with open(out, "w", newline="") as f: w = csv.writer(f) w.writerow(["method", "snr_db", "cos_sim", "orthogonality", "top1_acc"]) for meth, rows in results.items(): for r in rows: w.writerow([meth, *r]) print(f"\n[DONE] wrote {out}") if __name__ == "__main__": main()