Initial release: code for WCL2026-1544 (context-aware embedding masking via DRL)
This commit is contained in:
Ki-Ho Lee
@@ -0,0 +1,901 @@
|
||||
# =========================================================
|
||||
# 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)
|
||||
Reference in New Issue
Block a user