# ========================================================= # drl_mask_policy.py # PPO-based Mask Policy for Multi-User Semantic Multiplexing # # Companion to main_wcl.tex (IEEE WCL submission). # # Idea: # In the baseline (TVT paper), user masks {m_u} are plain # trainable parameters updated by gradient descent on the # reconstruction loss. Mask orthogonality emerges only as a # by-product, and the masks are not aware of the channel # state. # # Here we replace the fixed masks with a STATE-CONDITIONED # POLICY pi_phi(M | s), where the state s = (SNR, U) is the # wireless context. The policy is trained with PPO to maximize # # r(s, M) = sum_u CosSim(b_u, b_hat_u) # - beta * || M~ M~^T - I_U ||_F^2 / U^2, # # where M~ denotes the L2-normalized mask matrix. The first # term rewards semantic fidelity; the second explicitly rewards # mask orthogonality. # # The transceiver (projection, attention, reverse projection) is # identical to bert_semcom.py BertSemComMux and is trained jointly # with the policy via a shared optimizer, except that the mask # block is replaced by the policy output. # ========================================================= import math import os import csv import random import argparse import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------- # Utility: Rayleigh + AWGN channel (same as bert_semcom.py) # --------------------------------------------------------- def apply_channel(y, snr_db, channel="rayleigh"): snr_lin = 10.0 ** (snr_db / 10.0) noise_var = 1.0 / snr_lin if channel == "rayleigh": hr = torch.randn((), device=y.device) hi = torch.randn((), device=y.device) h_mag = torch.sqrt(hr ** 2 + hi ** 2) / math.sqrt(2.0) y = h_mag * y n = torch.randn_like(y) * math.sqrt(noise_var) return y + n # --------------------------------------------------------- # Transceiver without static masks (masks are injected) # --------------------------------------------------------- class MaskedTransceiver(nn.Module): """ Same projection + user-wise attention as BertSemComMux, but masks are SUPPLIED at forward time (by the policy). """ def __init__(self, U, d_bert, d_shared): super().__init__() self.U = U self.d_bert = d_bert self.d_shared = d_shared self.tx_proj = nn.Sequential( nn.Linear(d_bert, d_shared), nn.LayerNorm(d_shared), ) self.rx_proj = nn.Sequential( nn.LayerNorm(d_shared), nn.Linear(d_shared, d_bert), ) self.user_query = nn.Embedding(U, d_shared) nn.init.normal_(self.user_query.weight, std=0.5) def forward(self, b, masks, snr_db, channel="rayleigh"): """ b: (U, d_bert) user BERT embeddings masks: (U, d_shared) policy-generated masks """ e = self.tx_proj(b) # (U, d_s) x = e * masks # (U, d_s) y_tx = x.sum(dim=0) # (d_s,) y_tx = y_tx / (torch.sqrt((y_tx ** 2).mean() + 1e-8)) y_rx = apply_channel(y_tx, snr_db, channel) R = y_rx.unsqueeze(0) * masks # (U, d_s) Rn = F.normalize(R, p=2, dim=-1) idx = torch.arange(self.U, device=b.device) q = self.user_query(idx) # (U, d_s) qn = F.normalize(q, p=2, dim=-1) scores = (Rn @ qn.T) * 8.0 # (U_target, U_cand) attn = F.softmax(scores, dim=-1) # (U, U) z = attn @ R # (U, d_s) b_hat = self.rx_proj(z) # (U, d_bert) return b_hat, attn # --------------------------------------------------------- # Policy Actor: state -> mask mean + log-std # --------------------------------------------------------- class MaskActor(nn.Module): """ Contextual actor that maps state s=(SNR, U) to the parameters of a Gaussian policy over mask matrices M in R^{U x d_shared}. Implemented as a low-rank generator: M(s) = reshape( MLP(s) ) with a per-user amplitude. """ def __init__(self, U, d_shared, hidden=256, rank=64, log_std_init=-1.0): super().__init__() self.U, self.d_shared, self.rank = U, d_shared, rank state_dim = 2 # (SNR_norm, U_norm) self.trunk = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), ) # Output mean of mask matrix via low-rank factorization # Left: U x rank, Right: rank x d_shared self.left_head = nn.Linear(hidden, U * rank) self.right = nn.Parameter( torch.randn(rank, d_shared) / math.sqrt(rank)) # Learnable amplitude per user self.amp = nn.Parameter(torch.ones(U)) # Learnable log-std (state-independent). The initial value # controls exploration noise sigma = exp(log_std_init); smaller # values reduce the stochastic-sample bias that depresses # training-time CosSim and transceiver robustness to noise. self.log_std = nn.Parameter( float(log_std_init) * torch.ones(U * d_shared)) def forward(self, state): """ state: (state_dim,) tensor for a single context returns mean: (U, d_shared), std: (U, d_shared) """ h = self.trunk(state) L = self.left_head(h).view(self.U, self.rank) mean = (L @ self.right) * self.amp.view(-1, 1) std = torch.exp(self.log_std).view(self.U, self.d_shared) std = std.clamp(min=1e-3, max=1.0) return mean, std # --------------------------------------------------------- # Critic: state -> V(s) # --------------------------------------------------------- class MaskCritic(nn.Module): def __init__(self, state_dim=2, hidden=128): super().__init__() self.net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1), ) def forward(self, s): return self.net(s).squeeze(-1) # --------------------------------------------------------- # PPO trainer # --------------------------------------------------------- def encode_state(snr_db, U, U_max=8, snr_range=(0.0, 30.0), device="cpu"): s_norm = (float(snr_db) - snr_range[0]) / (snr_range[1] - snr_range[0]) u_norm = float(U) / float(U_max) return torch.tensor([s_norm, u_norm], device=device, dtype=torch.float) def orthogonality_penalty(masks): """|| M~ M~^T - I ||_F^2 / U^2, with row-normalized M~.""" Mn = F.normalize(masks, p=2, dim=-1) G = Mn @ Mn.T U = G.size(0) I = torch.eye(U, device=G.device) return ((G - I) ** 2).sum() / (U * U) def reward(b_hat, b, masks, beta=1.0): cos = F.cosine_similarity(b_hat, b, dim=-1).mean() orth = orthogonality_penalty(masks) return cos - beta * orth, cos.item(), orth.item() class RolloutBuffer: def __init__(self): self.states, self.actions, self.logps = [], [], [] self.rewards, self.values = [], [] def add(self, s, a, lp, r, v): self.states.append(s) self.actions.append(a) self.logps.append(lp) self.rewards.append(r) self.values.append(v) def clear(self): self.states, self.actions, self.logps = [], [], [] self.rewards, self.values = [], [] def ppo_update(actor, critic, buffer, actor_opt, critic_opt, clip=0.2, epochs=4, entropy_coef=1e-3): S = torch.stack(buffer.states) A = torch.stack(buffer.actions).view(len(buffer.states), -1) old_lp = torch.stack(buffer.logps).detach() R = torch.tensor(buffer.rewards, device=S.device, dtype=torch.float) V = torch.stack(buffer.values).detach() adv = R - V if adv.numel() > 1: adv = (adv - adv.mean()) / (adv.std() + 1e-6) for _ in range(epochs): # Re-evaluate logp and value new_lps, new_vs = [], [] ent_total = 0.0 for i in range(len(buffer.states)): mean, std = actor(buffer.states[i]) dist = torch.distributions.Normal(mean.view(-1), std.view(-1)) new_lps.append(dist.log_prob(A[i]).sum()) new_vs.append(critic(buffer.states[i])) ent_total = ent_total + dist.entropy().sum() new_lps = torch.stack(new_lps) new_vs = torch.stack(new_vs) ratio = torch.exp(new_lps - old_lp) s1 = ratio * adv s2 = torch.clamp(ratio, 1 - clip, 1 + clip) * adv actor_loss = -torch.min(s1, s2).mean() - entropy_coef * \ ent_total / len(buffer.states) critic_loss = F.mse_loss(new_vs, R) actor_opt.zero_grad(set_to_none=True) actor_loss.backward(retain_graph=True) nn.utils.clip_grad_norm_(actor.parameters(), 1.0) actor_opt.step() critic_opt.zero_grad(set_to_none=True) critic_loss.backward() nn.utils.clip_grad_norm_(critic.parameters(), 1.0) critic_opt.step() # --------------------------------------------------------- # Joint baseline: masks as free nn.Embedding parameters # --------------------------------------------------------- def train_joint(args): """Gradient-trained baseline: masks are plain trainable parameters, jointly optimized with the transceiver on the semantic loss (MSE + 1-CosSim), as in the TVT baseline.""" device = torch.device( "cuda" if args.cuda and torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")) torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) d_bert = args.d_bert d_shared = d_bert * args.mux_factor if args.embed_file and os.path.exists(args.embed_file): emb = torch.load(args.embed_file, map_location="cpu") emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1) print(f"[INFO] Loaded {emb.shape} embeddings.") else: print("[INFO] Using synthetic embeddings.") emb = F.normalize(torch.randn(args.pool_size, d_bert), p=2, dim=-1) emb = emb.to(device) trx = MaskedTransceiver(args.users, d_bert, d_shared).to(device) masks = nn.Embedding(args.users, d_shared).to(device) nn.init.normal_(masks.weight, std=0.5) params = list(trx.parameters()) + list(masks.parameters()) opt = torch.optim.Adam(params, lr=args.lr) train_snrs = list(args.train_snr) log_rows = [] for ep in range(1, args.epochs + 1): ep_cos, ep_orth, ep_loss = [], [], [] for _ in range(args.steps_per_epoch): snr_db = random.choice(train_snrs) idx = torch.randint(0, emb.size(0), (args.users,)) b = emb[idx] m = masks(torch.arange(args.users, device=device)) b_hat, _ = trx(b, m, snr_db, args.channel) cos = F.cosine_similarity(b_hat, b, dim=-1).mean() mse = F.mse_loss(b_hat, b) loss = 0.5 * mse + 0.5 * (1.0 - cos) opt.zero_grad(set_to_none=True) loss.backward() nn.utils.clip_grad_norm_(params, 1.0) opt.step() with torch.no_grad(): orth = orthogonality_penalty(m).item() ep_cos.append(cos.item()) ep_orth.append(orth) ep_loss.append(loss.item()) mc, mo, ml = (float(np.mean(ep_cos)), float(np.mean(ep_orth)), float(np.mean(ep_loss))) print(f"[Joint ep {ep:03d}/{args.epochs}] loss={ml:.4f} " f"CosSim={mc:.4f} Orth={mo:.4f}") log_rows.append([ep, ml, mc, mo]) os.makedirs(args.save_dir, exist_ok=True) with open(os.path.join(args.save_dir, "joint_train_log.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["epoch", "loss", "cos_sim", "orthogonality"]) w.writerows(log_rows) # SNR sweep sweep_rows = [] trx.eval() with torch.no_grad(): for snr_db in range(0, 31, 5): coss, orths = [], [] for _ in range(args.eval_trials): idx = torch.randint(0, emb.size(0), (args.users,)) b = emb[idx] m = masks(torch.arange(args.users, device=device)) b_hat, _ = trx(b, m, snr_db, args.channel) coss.append(F.cosine_similarity( b_hat, b, dim=-1).mean().item()) orths.append(orthogonality_penalty(m).item()) sweep_rows.append([snr_db, float(np.mean(coss)), float(np.mean(orths))]) print(f"[SNR {snr_db:>4.1f}dB] " f"CosSim={sweep_rows[-1][1]:.4f} " f"Orth={sweep_rows[-1][2]:.4f}") with open(os.path.join(args.save_dir, "joint_snr_sweep.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["snr_db", "cos_sim", "orthogonality"]) w.writerows(sweep_rows) print(f"Saved joint baseline to {args.save_dir}") # --------------------------------------------------------- # CE baseline: TVT-style cross-entropy over the AG News pool # --------------------------------------------------------- def train_joint_ce(args): """Cross-entropy baseline matching the TVT/JSAC formulation L = - sum_u sum_t sum_i x_{u,t}^(i) log p_hat_{u,t}^(i), instantiated for the WCL setup with C = pool_size classes (the 8,000 AG News sentences) and T = 1 token per user per step. Logits are computed as the temperature-scaled cosine similarity between the reconstructed user embedding and every pool entry, p_hat_u = softmax(tau * b_hat_u E_pool^T). Masks remain free trainable parameters (as in train_joint), so the only difference from the static-masking baseline is the loss surface (CE on symbol identity vs MSE+1-CosSim). The transceiver, optimizer, channel, and mask init are identical to train_joint, isolating the loss as the variable. """ device = torch.device( "cuda" if args.cuda and torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")) torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) d_bert = args.d_bert d_shared = d_bert * args.mux_factor if args.embed_file and os.path.exists(args.embed_file): emb = torch.load(args.embed_file, map_location="cpu") emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1) print(f"[INFO] Loaded {emb.shape} embeddings.") else: print("[INFO] Using synthetic embeddings.") emb = F.normalize(torch.randn(args.pool_size, d_bert), p=2, dim=-1) emb = emb.to(device) C = emb.size(0) trx = MaskedTransceiver(args.users, d_bert, d_shared).to(device) masks = nn.Embedding(args.users, d_shared).to(device) nn.init.normal_(masks.weight, std=0.5) params = list(trx.parameters()) + list(masks.parameters()) opt = torch.optim.Adam(params, lr=args.lr) tau = float(args.ce_tau) train_snrs = list(args.train_snr) log_rows = [] for ep in range(1, args.epochs + 1): ep_cos, ep_orth, ep_loss = [], [], [] for _ in range(args.steps_per_epoch): snr_db = random.choice(train_snrs) idx = torch.randint(0, C, (args.users,), device=device) b = emb[idx] m = masks(torch.arange(args.users, device=device)) b_hat, _ = trx(b, m, snr_db, args.channel) b_hat_n = F.normalize(b_hat, p=2, dim=-1) logits = tau * (b_hat_n @ emb.T) # (U, C) loss = F.cross_entropy(logits, idx) opt.zero_grad(set_to_none=True) loss.backward() nn.utils.clip_grad_norm_(params, 1.0) opt.step() with torch.no_grad(): cos = F.cosine_similarity( b_hat, b, dim=-1).mean().item() orth = orthogonality_penalty(m).item() ep_cos.append(cos) ep_orth.append(orth) ep_loss.append(loss.item()) mc, mo, ml = (float(np.mean(ep_cos)), float(np.mean(ep_orth)), float(np.mean(ep_loss))) print(f"[JointCE ep {ep:03d}/{args.epochs}] loss={ml:.4f} " f"CosSim={mc:.4f} Orth={mo:.4f}") log_rows.append([ep, ml, mc, mo]) os.makedirs(args.save_dir, exist_ok=True) with open(os.path.join(args.save_dir, "joint_ce_train_log.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["epoch", "loss", "cos_sim", "orthogonality"]) w.writerows(log_rows) sweep_rows = [] trx.eval() with torch.no_grad(): for snr_db in range(0, 31, 5): coss, orths = [], [] for _ in range(args.eval_trials): idx = torch.randint(0, C, (args.users,), device=device) b = emb[idx] m = masks(torch.arange(args.users, device=device)) b_hat, _ = trx(b, m, snr_db, args.channel) coss.append(F.cosine_similarity( b_hat, b, dim=-1).mean().item()) orths.append(orthogonality_penalty(m).item()) sweep_rows.append([snr_db, float(np.mean(coss)), float(np.mean(orths))]) print(f"[SNR {snr_db:>4.1f}dB] " f"CosSim={sweep_rows[-1][1]:.4f} " f"Orth={sweep_rows[-1][2]:.4f}") with open(os.path.join(args.save_dir, "joint_ce_snr_sweep.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["snr_db", "cos_sim", "orthogonality"]) w.writerows(sweep_rows) print(f"Saved CE baseline to {args.save_dir}") # --------------------------------------------------------- # MAML baseline: first-order MAML across SNR tasks # --------------------------------------------------------- def train_maml(args): """First-order MAML (FOMAML) where each SNR value is a task. At every meta step we take one inner gradient step on a support batch, then evaluate query loss with the fast params. Meta gradient is the average of query losses across tasks. This matches the MAML baseline of ~\cite{lee2025sharedembedding,finn2017maml} but is now run on the same real BERT embeddings as DRL and Joint.""" try: from torch.func import functional_call except Exception: from torch.nn.utils.stateless import functional_call device = torch.device( "cuda" if args.cuda and torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")) torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) d_bert = args.d_bert d_shared = d_bert * args.mux_factor if args.embed_file and os.path.exists(args.embed_file): emb = torch.load(args.embed_file, map_location="cpu") emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1) print(f"[INFO] Loaded {emb.shape} embeddings.") else: emb = F.normalize(torch.randn(args.pool_size, d_bert), p=2, dim=-1) emb = emb.to(device) trx = MaskedTransceiver(args.users, d_bert, d_shared).to(device) masks = nn.Embedding(args.users, d_shared).to(device) nn.init.normal_(masks.weight, std=0.5) # Combine all trainable modules into a single dict of # parameters so that the inner-loop adaptation and outer-loop # meta-update can be expressed uniformly. def _named_modules(): return {"trx": trx, "masks": masks} outer_params = list(trx.parameters()) + list(masks.parameters()) outer_opt = torch.optim.Adam(outer_params, lr=args.lr) train_snrs = list(args.train_snr) inner_lr = args.inner_lr meta_batch = min(args.meta_batch, len(train_snrs)) log_rows = [] def _forward_with_params(fast, b, snr_db): """Forward pass of the transceiver with the mask matrix derived from fast['masks']['weight'].""" m = functional_call( masks, fast["masks"], (torch.arange(args.users, device=device),)) # trx.forward uses masks as an argument, not as a submodule # parameter, so we only need functional_call for trx tx_proj, # rx_proj, and user_query submodules. Pass fast['trx'] as # the state dict. b_hat, _ = functional_call( trx, fast["trx"], (b, m, snr_db, args.channel)) return b_hat for ep in range(1, args.epochs + 1): ep_cos, ep_loss = [], [] for _ in range(args.steps_per_epoch): # Sample a meta-batch of SNR tasks. tasks = random.sample(train_snrs, meta_batch) \ if meta_batch <= len(train_snrs) \ else [random.choice(train_snrs) for _ in range(meta_batch)] meta_loss = torch.zeros((), device=device) for snr in tasks: # Snapshot current params. fast = { "trx": {n: p for n, p in trx.named_parameters()}, "masks": {n: p for n, p in masks.named_parameters()}, } # --- Inner loop: 1 gradient step on support batch idx_s = torch.randint(0, emb.size(0), (args.users,)) b_s = emb[idx_s] b_hat_s = _forward_with_params(fast, b_s, snr) cos_s = F.cosine_similarity( b_hat_s, b_s, dim=-1).mean() mse_s = F.mse_loss(b_hat_s, b_s) loss_s = 0.5 * mse_s + 0.5 * (1.0 - cos_s) flat = [] meta_index = [] for key in ("trx", "masks"): for n, p in fast[key].items(): flat.append(p) meta_index.append((key, n)) grads = torch.autograd.grad( loss_s, flat, create_graph=False) # Update fast params (detach for first-order). for (key, n), g in zip(meta_index, grads): fast[key][n] = fast[key][n] - inner_lr * \ g.detach() # --- Query loss with fast params idx_q = torch.randint(0, emb.size(0), (args.users,)) b_q = emb[idx_q] b_hat_q = _forward_with_params(fast, b_q, snr) cos_q = F.cosine_similarity( b_hat_q, b_q, dim=-1).mean() mse_q = F.mse_loss(b_hat_q, b_q) loss_q = 0.5 * mse_q + 0.5 * (1.0 - cos_q) meta_loss = meta_loss + loss_q ep_cos.append(cos_q.item()) meta_loss = meta_loss / float(len(tasks)) outer_opt.zero_grad(set_to_none=True) meta_loss.backward() nn.utils.clip_grad_norm_(outer_params, 1.0) outer_opt.step() ep_loss.append(meta_loss.item()) mc = float(np.mean(ep_cos)) ml = float(np.mean(ep_loss)) print(f"[MAML ep {ep:03d}/{args.epochs}] " f"meta_loss={ml:.4f} CosSim(q)={mc:.4f}") log_rows.append([ep, ml, mc, 0.0]) # orthogonality not # tracked for MAML os.makedirs(args.save_dir, exist_ok=True) with open(os.path.join(args.save_dir, "maml_train_log.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["epoch", "loss", "cos_sim", "orthogonality"]) w.writerows(log_rows) # --- SNR sweep: for each eval SNR, adapt with 1 inner step # then evaluate on fresh query samples (deterministic eval). sweep_rows = [] for snr_db in range(0, 31, 5): coss = [] for _ in range(args.eval_trials): fast = { "trx": {n: p for n, p in trx.named_parameters()}, "masks": {n: p for n, p in masks.named_parameters()}, } # Inner adapt idx_s = torch.randint(0, emb.size(0), (args.users,)) b_s = emb[idx_s] b_hat_s = _forward_with_params(fast, b_s, snr_db) cos_s = F.cosine_similarity( b_hat_s, b_s, dim=-1).mean() mse_s = F.mse_loss(b_hat_s, b_s) loss_s = 0.5 * mse_s + 0.5 * (1.0 - cos_s) flat, meta_index = [], [] for key in ("trx", "masks"): for n, p in fast[key].items(): flat.append(p) meta_index.append((key, n)) grads = torch.autograd.grad(loss_s, flat, create_graph=False) for (key, n), g in zip(meta_index, grads): fast[key][n] = fast[key][n] - \ args.eval_inner_lr * g.detach() # Query with torch.no_grad(): idx_q = torch.randint(0, emb.size(0), (args.users,)) b_q = emb[idx_q] b_hat_q = _forward_with_params(fast, b_q, snr_db) coss.append(F.cosine_similarity( b_hat_q, b_q, dim=-1).mean().item()) # orthogonality of adapted masks (skip: we report deterministic # outer masks for fairness with Joint) with torch.no_grad(): m = masks(torch.arange(args.users, device=device)) orth = orthogonality_penalty(m).item() sweep_rows.append([snr_db, float(np.mean(coss)), orth]) print(f"[MAML SNR {snr_db:>4.1f}dB] " f"CosSim={sweep_rows[-1][1]:.4f} O={orth:.4f}") with open(os.path.join(args.save_dir, "maml_snr_sweep.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["snr_db", "cos_sim", "orthogonality"]) w.writerows(sweep_rows) print(f"Saved MAML baseline to {args.save_dir}") # --------------------------------------------------------- # Main training loop (DRL-guided mask policy) # --------------------------------------------------------- def train_drl(args): device = torch.device( "cuda" if args.cuda and torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")) torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) d_bert = args.d_bert d_shared = d_bert * args.mux_factor # Synthetic (or preloaded) centered, L2-normalized embeddings. # The DRL contribution is orthogonal to the BERT extraction step, # so we can use cached embeddings stored as a tensor. If the user # has real AG News embeddings from bert_semcom.py, they can be # loaded from --embed-file. if args.embed_file and os.path.exists(args.embed_file): emb = torch.load(args.embed_file, map_location="cpu") emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1) print(f"[INFO] Loaded {emb.shape} embeddings from " f"{args.embed_file}") else: print("[INFO] Using synthetic embeddings " "(pass --embed-file for real BERT).") emb = F.normalize(torch.randn(args.pool_size, d_bert), p=2, dim=-1) emb = emb.to(device) # Transceiver + policy trx = MaskedTransceiver(args.users, d_bert, d_shared).to(device) actor = MaskActor(args.users, d_shared, hidden=args.hidden, rank=args.rank, log_std_init=args.log_std_init).to(device) critic = MaskCritic(state_dim=2, hidden=128).to(device) trx_opt = torch.optim.Adam(trx.parameters(), lr=args.lr) actor_opt = torch.optim.Adam(actor.parameters(), lr=args.lr_actor) critic_opt = torch.optim.Adam(critic.parameters(), lr=args.lr_critic) train_snrs = list(args.train_snr) buf = RolloutBuffer() log_rows = [] for ep in range(1, args.epochs + 1): ep_cos, ep_orth, ep_rew = [], [], [] for step in range(args.steps_per_epoch): # --- Sample context --- snr_db = random.choice(train_snrs) state = encode_state(snr_db, args.users, args.users_max, device=device) # --- Actor: sample mask matrix --- mean, std = actor(state) dist = torch.distributions.Normal(mean.view(-1), std.view(-1)) a = dist.sample() logp = dist.log_prob(a).sum() masks = a.view(args.users, d_shared) # --- Transceiver forward --- idx = torch.randint(0, emb.size(0), (args.users,)) b = emb[idx] b_hat, _ = trx(b, masks, snr_db, args.channel) r_t, cos_v, orth_v = reward(b_hat, b, masks, beta=args.beta) # --- Critic value --- v = critic(state) buf.add(state.detach(), a.detach(), logp.detach(), r_t.item(), v) ep_cos.append(cos_v) ep_orth.append(orth_v) ep_rew.append(r_t.item()) # --- Transceiver update on differentiable path --- # For the transceiver (projection + attention + rx_proj), # we can use the deterministic mean mask as a low-variance # auxiliary loss. This keeps trx trained well while the # stochastic policy continues to explore mask placements. with torch.enable_grad(): masks_det = mean.detach() + 0.0 # cut PPO graph masks_det.requires_grad_(False) b_hat_aux, _ = trx(b, masks_det, snr_db, args.channel) aux = 1.0 - F.cosine_similarity( b_hat_aux, b, dim=-1).mean() trx_opt.zero_grad(set_to_none=True) aux.backward() trx_opt.step() # --- PPO update periodically --- if len(buf.states) >= args.ppo_batch: ppo_update(actor, critic, buf, actor_opt, critic_opt, clip=args.clip, epochs=args.ppo_epochs, entropy_coef=args.entropy_coef) buf.clear() mean_cos = float(np.mean(ep_cos)) mean_orth = float(np.mean(ep_orth)) mean_rew = float(np.mean(ep_rew)) # --- Deterministic-mu diagnostic eval (no gradient, no sampling) # Small per-epoch pass to log what the *deployed* policy achieves, # i.e. using the mean mu_phi(s) instead of a sampled M. This is # what the paper reports at evaluation time; overlaying both # curves in Fig 2(a) reveals that the sampling noise -- not the # policy itself -- is what depresses the training-time CosSim. actor.eval(); trx.eval() det_cos = [] with torch.no_grad(): for _ in range(args.det_eval_trials): snr_db_d = random.choice(train_snrs) st_d = encode_state(snr_db_d, args.users, args.users_max, device=device) mean_d, _ = actor(st_d) idx_d = torch.randint(0, emb.size(0), (args.users,)) b_d = emb[idx_d] b_hat_d, _ = trx(b_d, mean_d, snr_db_d, args.channel) det_cos.append(F.cosine_similarity( b_hat_d, b_d, dim=-1).mean().item()) actor.train(); trx.train() mean_cos_det = float(np.mean(det_cos)) if det_cos else mean_cos print(f"[DRL ep {ep:03d}/{args.epochs}] " f"R={mean_rew:.4f} CosSim(sample)={mean_cos:.4f} " f"CosSim(det)={mean_cos_det:.4f} Orth={mean_orth:.4f}") log_rows.append([ep, mean_rew, mean_cos, mean_orth, mean_cos_det]) os.makedirs(args.save_dir, exist_ok=True) with open(os.path.join(args.save_dir, "drl_train_log.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["epoch", "reward", "cos_sim", "orthogonality", "cos_sim_det"]) w.writerows(log_rows) # SNR sweep at convergence sweep_rows = [] actor.eval(); critic.eval(); trx.eval() with torch.no_grad(): for snr_db in range(0, 31, 5): coss, orths = [], [] for _ in range(args.eval_trials): state = encode_state(snr_db, args.users, args.users_max, device=device) mean, _ = actor(state) idx = torch.randint(0, emb.size(0), (args.users,)) b = emb[idx] b_hat, _ = trx(b, mean, snr_db, args.channel) coss.append(F.cosine_similarity( b_hat, b, dim=-1).mean().item()) orths.append(orthogonality_penalty(mean).item()) sweep_rows.append( [snr_db, float(np.mean(coss)), float(np.mean(orths))]) print(f"[SNR {snr_db:>4.1f}dB] CosSim={sweep_rows[-1][1]:.4f} " f"Orth={sweep_rows[-1][2]:.4f}") with open(os.path.join(args.save_dir, "drl_snr_sweep.csv"), "w", newline="") as f: w = csv.writer(f) w.writerow(["snr_db", "cos_sim", "orthogonality"]) w.writerows(sweep_rows) # Save trained modules torch.save({ "actor": actor.state_dict(), "critic": critic.state_dict(), "trx": trx.state_dict(), "args": vars(args), }, os.path.join(args.save_dir, "drl_mask_policy.pt")) print(f"\nSaved to {args.save_dir}") # --------------------------------------------------------- # CLI # --------------------------------------------------------- if __name__ == "__main__": p = argparse.ArgumentParser() # model 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("--log-std-init", type=float, default=-1.0, help="Initial log-std for the Gaussian policy. sigma = exp(.). Default -1.0 (~0.37); try -2.0 (~0.14).") # data p.add_argument("--embed-file", type=str, default=None) p.add_argument("--pool-size", type=int, default=4000) # channel p.add_argument("--channel", choices=["awgn", "rayleigh"], default="rayleigh") p.add_argument("--train-snr", type=float, nargs="+", default=[0, 5, 10, 15, 20, 25]) # optim p.add_argument("--lr", type=float, default=1e-3) p.add_argument("--lr-actor", type=float, default=3e-4) p.add_argument("--lr-critic", type=float, default=1e-3) # PPO p.add_argument("--ppo-batch", type=int, default=64) p.add_argument("--ppo-epochs", type=int, default=4) p.add_argument("--clip", type=float, default=0.2) p.add_argument("--entropy-coef", type=float, default=1e-3) p.add_argument("--beta", type=float, default=0.5, help="Orthogonality penalty weight") # run 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("--det-eval-trials", type=int, default=32, help="Per-epoch deterministic-mu diagnostic trials") p.add_argument("--save-dir", type=str, default="results_drl") p.add_argument("--cuda", action="store_true") p.add_argument("--seed", type=int, default=0) p.add_argument("--mode", choices=["drl", "joint", "joint_ce", "maml"], default="drl") p.add_argument("--inner-lr", type=float, default=5e-4, help="MAML inner-loop learning rate") p.add_argument("--eval-inner-lr", type=float, default=5e-4, help="MAML inner-loop LR at evaluation") p.add_argument("--meta-batch", type=int, default=4, help="MAML: number of SNR tasks per meta step") p.add_argument("--ce-tau", type=float, default=16.0, help="Logit temperature for joint_ce mode " "(scales b_hat E_pool^T before softmax).") args = p.parse_args() if args.mode == "joint": train_joint(args) elif args.mode == "joint_ce": train_joint_ce(args) elif args.mode == "maml": train_maml(args) else: train_drl(args)